This document is relevant for: Trn2, Trn3

Neuron Kernel Interface (NKI)#

NKI is a bare-metal language and compiler for programming AWS Trainium and Inferentia NeuronDevices directly, giving you instruction-level control to develop, optimize, and run custom operators on NeuronCores.

Before you write a kernel, check the NKI Library — it may already have an optimized kernel for your operation. New to NKI? Start with Get started. For how NKI fits into the Neuron compilation stack, see How NKI works below.

Start here#

Get started

Set up your environment, then implement and run your first NKI kernel.

NKI Library

Use prebuilt, optimized kernels for common operations. Check here before writing your own.

Learn and build#

Tutorials

Step-by-step walkthroughs for building NKI kernels, from basics to advanced patterns.

How-to guides

Task-focused guides: insert kernels into models, use the CPU simulator, profile, and control scheduling.

Deep dives

Performance optimization, access patterns, DMA, dynamic loops, and the NKI compiler.

NKI FAQ

Common questions about programming with NKI.

Reference#

NKI API reference manual

Complete API reference for nki, nki.language, nki.isa, and nki.collectives.

How NKI works#

NKI provides developers with direct access to the NeuronCore ISA (Instruction Set Architecture), accessible from a Python-based programming environment, which has syntax and tile-level semantics that are similar to Triton and NumPy. This enables developers to get started quickly and optimize performance in a familiar environment, while at the same time get full control of the underlying hardware. At the hardware level, NeuronCore’s tensorized memory access capability enables efficient reading and writing of multi-dimensional arrays on a per instruction basis, which makes NKI’s tile-based programming highly suitable for the NeuronCore instruction set.

For comparison, before NKI was introduced, the only way to program NeuronDevices was through defining high-level ML models in frameworks such as PyTorch and JAX. Neuron Compiler takes such high-level model definitions as input, performs multiple rounds of optimization, and eventually generates a NEFF (Neuron Executable File Format) that is executable on NeuronDevices. At a high level, Neuron Compiler runs the following optimization stages in order:

  1. Hardware-agnostic graph-level optimizations. These transformations are done in the compiler front-end, using XLA, including optimizations like constant propagation, re-materialization and operator fusion.

  2. Loop-level optimization. Compiler turns the optimized graph from Step 1 into a series of loop nests and performs layout, tiling and loop fusion optimizations.

  3. Hardware intrinsics mapping. Compiler maps the architecture-agnostic loop nests from Step 2 into architecture-specific instructions.

  4. Hardware-specific optimizations. These optimizations are mainly done at the instruction level in compiler back-end, with a key goal of reducing memory pressure and improving instruction-level parallelism. For example, memory allocation and instruction scheduling are done in this stage.

NKI kernels bypass the first 3 steps, and are compiled into IRs (intermediate representations) that the compiler’s back-end (Step 4 above) can directly consume. Advanced features in NKI, such as direct allocation, also allow programmers to bypass certain compiler passes in Step 4. As a result, NKI developers can now have great control over NeuronDevices down to the instruction level.

Note

Neuron highly recommends developers study the underlying hardware architecture before optimizing performance of their NKI kernels. To start, read the NKI architecture guides for Trainium and the NKI performance guide.

This document is relevant for: Trn2, Trn3