NKI Library

The NKI Library provides pre-built reference kernels you can use directly in your model development with the AWS Neuron SDK and NKI. These kernel APIs provide the default classes, functions, and parameters you can use to integrate the NKL kernels into your models. More details can be found in the NKI Library Documentation

Kernel Reference

Kernel API	Description
Attention CTE Kernel	The kernel implements attention with support for multiple variants and optimizations.
Attention KV Parallel Segmented CTE Kernel	The kernel implements KV-parallel segmented prefill attention with online softmax merging for context parallelism.
Attention TKG Kernel	The kernel implements attention specifically optimized for token generation use cases.
MLP Kernel	The kernel implements a Multi-Layer Perceptron with optional normalization fusion and various optimizations.
MoE CTE Kernel	The kernel implements Mixture of Experts optimized for Context Encoding use cases.
MoE TKG Kernel	The kernel implements Mixture of Experts optimized for Token Generation use cases.
Output Projection CTE Kernel	The kernel computes the output projection operation optimized for Context Encoding use cases.
Output Projection TKG Kernel	The kernel computes the output projection operation optimized for Token Generation use cases.
QKV Kernel	The kernel performs Query-Key-Value projection with optional normalization fusion.
RMSNorm-Quant Kernel	The kernel performs optional RMS normalization followed by quantization to fp8.
RMSNorm MX Prefill Kernel	The kernel fuses RMSNorm with MX quantization and optional router top-K in token-major [T, H] layout for prefill.
RoPE Kernel	The kernel applies Rotary Position Embedding to input embeddings with optional LNC sharding.
Router Top-K Kernel	The kernel computes router logits and top-K selection for Mixture of Experts models.
Cumsum Kernel	The kernel computes cumulative sum along the last dimension.

Experimental Kernels

Kernel API	Description
Attention Block TKG Kernel	The kernel implements fused attention block for TKG with RMSNorm, QKV, RoPE, and output projection.
Cross Entropy Kernel	The kernel implements memory-efficient cross entropy loss forward and backward passes for large vocabularies.
Depthwise Conv1D Kernel	The kernel implements depthwise 1D convolution using implicit GEMM.
Blockwise MM Backward Kernel	The kernel implements blockwise matrix multiplication backward pass for dropless Mixture of Experts.
Conv1D Kernel	The kernel implements 1D convolution using a filter replication strategy.
Conv3D Kernel	The kernel implements 3D convolution using a filter replication strategy.
Conv3D Transpose Kernel	The kernel implements 3D transposed convolution by performing a 3D convolution with input dilation.
Conv3D Temporal Unroll Kernel	The kernel implements 3D convolution with temporal unrolling and column tiling for small C_out configurations.
Dynamic Shape Kernels	The kernels dynamic input shapes with dynamic loop tiling on dynamic dimension.
Fine-Grained AllGather Kernel	The kernel implements fine-grained ring-based all-gather.
FGCC Kernel	The kernel implements fused all-gather and matrix multiplication (Fine-Grained Gather Collective Compute).
Ring Attention Backward Kernel	The kernel implements the backward pass for ring attention using collective permute operations.
RNG Kernels	The kernels provide GPSIMD engine RNG state management and random number generation.
Transformer TKG Kernel	The kernel implements a transformer forward pass megakernel optimized for token generation (TKG).
Ring Attention Forward Kernel	The kernel implements ring attention forward using attention_cte with HBM I/O and online softmax reduction for context parallelism.
Fused Adam Kernel	The kernel implements a fused Adam/AdamW/AMSGrad optimizer step with SPMD tiling and Scalar Engine fusion.
MXFP8 Matmul Kernel	The kernel implements MXFP8 matrix multiplication with configurable tiling and quantization.
MXFP8 MLP Kernels	The kernels implement MXFP8 MLP forward and backward passes with recompute support.
MXFP8 MoE Backward Kernel	The kernel implements blockwise matrix multiplication backward pass for dropless Mixture of Experts using MXFP8 quantized matmuls.
MXFP8 Quantize Kernel	The kernel implements block-wise MXFP8 quantization with scale packing.
Foreach Norm Kernels	The kernels implement L1, L2, and Linf norm computation with SPMD tiling and fused activation-reduce.
Foreach Elementwise Kernels	The kernels implement elementwise add, sub, mul, div, addcdiv, addcmul and sqrt operations with SPMD tiling for scalar and tensor operands.
Linear Scan Kernel	The kernel computes first-order linear recurrence along the last dimension using tensor_tensor_scan.
Selective Scan Kernel	The kernel implements fused Mamba-style discretization, recurrence, and output projection.
SSD Kernel	The kernel implements Mamba-2 chunk-wise parallel SSD computation with TensorE matmuls and VectorE scans, with optional LNC sharding across heads.
Gather Kernel	The kernel gathers rows from a 2D input tensor based on a 1D index tensor using indirect DMA load.
Scatter-Add Kernel	The kernel scatter-adds values from a source tensor into a destination tensor using a gather-accumulate-scatter pattern.
NeuroTile	Tile-iterator library that abstracts HBM/SBUF/PSUM tiling, sharding, and access-pattern construction for kernel authors; tutorials and example kernels live under src/nkilib_src/nkilib/experimental/neurotile/examples/.
MSDeformableAttention Kernel	The kernel implements multi-scale deformable attention with an indirect DMA transpose strategy.
MSDeformableAttentionBwd Kernel	The kernel implements multi-scale deformable attention backward with an indirect DMA transpose and combined bilinear corner scatter-add strategy.
GpSIMD Top-K Kernel	The kernel computes top-k over the last dimension using the GpSIMD nisa.topk instruction (bfloat16, gen3+).

Integration with the Neuron Compiler

The Neuron compiler includes a bundled version of this package within neuronx-cc, accessible under the nkilib Python namespace (for example, import nkilib). This bundled version is referred to as "bundled nkilib" throughout this guide. Bundled nkilib has been validated to work with that particular compiler version and can be used out of the box.

If you want to contribute a kernel change or use the latest kernels, you can integrate with this package directly.

Note: Unlike bundled nkilib, kernels from this package are not guaranteed to be compatible with the latest release of the Neuron compiler. To start from a known good commit compatible with your compiler version, find the branch corresponding to your compiler version in this repository.

Installation

1. Install neuronx-cc as usual (most likely already done). For more information, see the Neuron Quick Start Guide.
2. Install this package into the same virtual environment as the rest of your project:

   pip install nki-library

3. Import and use kernels as usual. This package automatically replaces bundled nkilib kernels with the content of this package. No code changes are required.

Uninstalling

To uninstall, run the following command:

pip uninstall nki-library

After uninstalling, the compiler falls back to the bundled nkilib.

Controlling which package gets loaded

To temporarily revert to the bundled version of nkilib, set the NKILIB_FORCE_BUNDLED_LIBRARY environment variable to a truthy value:

export NKILIB_FORCE_BUNDLED_LIBRARY=true

On the next execution of neuronx-cc, it will use the bundled version of nkilib. To go back to the kernels from this package, unset NKILIB_FORCE_BUNDLED_LIBRARY

unset NKILIB_FORCE_BUNDLED_LIBRARY