NKI Compiler Documentation#

The NKI (Neuron Kernel Interface) Compiler is a key component of the AWS Neuron SDK that enables direct programming of NeuronDevices on AWS Trn/Inf instances. It compiles NKI kernel functions marked with nki.jit into optimized code that can run efficiently on NeuronCores. This section provides an overview of the NKI Compiler’s architecture, its integration with the Neuron SDK, and best practices for developing high-performance kernels.

About the NKI Compiler

NKI has its own compiler! Learn more about it here.

How the NKI Compiler Works with the Neuron Graph Compiler

Learn how the NKI compiler integrates with the Neuron SDK graph compiler to optimize and execute kernel functions on NeuronDevices.

Migrate NKI Beta 1 Kernels to the New NKI Compiler

Learn how to migrate AWS Neuron NKI Beta 1 kernels to the new NKI compiler in Beta 2.