This document is relevant for: Trn2, Trn3

NKI API reference manual#

The NKI (Neuron Kernel Interface) API reference documents every module you use to write, compile, and run custom kernels on AWS Trainium and Inferentia. It covers the high-level kernel language, the low-level hardware ISA, multi-core collectives, and the decorators that compile and simulate kernels. Use the sections below to jump to the module you need.

Compile and run kernels#

nki

Top-level entry points for compiling and running NKI kernels, including the jit decorator and the simulate CPU simulator for debugging.

Write kernels#

nki.language

High-level constructs for writing kernels: tensor creation, indexing, type casting, math operations, and loop constructs the compiler maps to hardware.

Hardware ISA#

nki.isa

Low-level ISA instructions for compute, data movement, and synchronization, mapping to Tensor, Vector, Scalar, and DMA engine operations.

NKI ISA common fields

Shared fields used across nki.isa APIs, including the supported NKI data types accepted by the dtype parameter.

Multi-core communication#

nki.collectives

Collective communication operations such as all-reduce and all-gather for multi-rank kernels that run across NeuronCores.