This document is relevant for: Trn2, Trn3

NKI Library Supported Kernel Reference

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

Source code for these kernel APIs can be found at: aws-neuron/nki-library

Core Kernels

Normalization and Quantization Kernels

RMSNorm-Quant	Performs optional RMS normalization followed by quantization to fp8.
RMSNorm MX Prefill	Fused RMSNorm + MX quantization (+ optional router top-K) for prefill, with optional residual add.

QKV Projection Kernels

QKV	Performs Query-Key-Value projection with optional normalization and RoPE fusion.

Attention Kernels

Attention CTE	Implements attention optimized for Context Encoding (prefill) use cases.
Attention Segmented CTE	Segmented attention with block-based KV cache and prefix caching for decode.
Attention TKG	Implements attention optimized for Token Generation (decode) use cases with small active sequence lengths.
Attention KV-Parallel Segmented CTE	KV-parallel segmented prefill attention with paged KV cache, sliding-window, and round-robin KV distribution.

Rotary Position Embedding (RoPE) Kernels

RoPE	Applies Rotary Position Embedding to input embeddings with flexible layout support.

Multi-Layer Perceptron (MLP) Kernels

MLP	Implements Multi-Layer Perceptron with optional normalization fusion and quantization support.

Output Projection Kernels

Output Projection CTE	Computes output projection optimized for Context Encoding use cases.
Output Projection TKG	Computes output projection optimized for Token Generation use cases.

Mixture of Experts (MoE) Kernels

Router Top-K	Computes router logits, applies activation functions, and performs top-K selection for MoE models.
MoE CTE	Implements Mixture of Experts MLP operations optimized for Context Encoding use cases.
MoE TKG	Implements Mixture of Experts MLP operations optimized for Token Generation use cases.

Quantization Kernels

FP8 Quantize

Static and row-wise dynamic FP8 quantization with pre-combined dequantization scale support.

Cumulative Sum Kernels

Cumsum

Computes cumulative sum along the last dimension with optimized tiling.

Core Subkernels

Find Nonzero Indices

Finds indices of nonzero elements along the T dimension using GpSimd nonzero_with_count ISA.

Experimental Kernels

Note

Experimental kernels are under active development and their APIs may change in future releases.

Attention Kernels

Attention Block TKG	Fused attention block for Token Generation that keeps all intermediate tensors in SBUF to minimize HBM traffic.
Ring Attention Forward	Ring attention forward pass for context parallelism across multiple workers.
Ring Attention Backward	Ring attention backward pass SPMD kernel for context parallelism.
Ring Attention Unpermute	Reorders striped ring-attention output back to contiguous sequence order.
QKV CTE MLA	DeepSeek Multi-head Latent Attention (MLA) QKV projection with MX quantization for context encoding.
SWA Fused CTE	Fused GPT-OSS sliding-window-attention block with QKV projection, RoPE, attention with sink, and output projection; updates the K/V caches in place.
Attention MXFP8 TKG	MXFP8 flash-decode attention over a quantized block KV cache with packed-Q eviction (Trainium 3).
Sparse Attention Indexer (MX + BF16 Score)	DeepSeek Sparse Attention Indexer with MX projections and a BF16 score path.
MLA QKV CTE	DeepSeek MLA QKV projection (MX) for context encoding, emitting absorbed latents.
MLA Sparse Attention CTE	Standalone sparse latent + RoPE attention (kernel A of the split DeepSeek-V3.2 sparse-MLA forward).
MLA V-Up O-Proj CTE	Standalone MX V-up + MX output projection (kernel B of the split DeepSeek-V3.2 sparse-MLA forward).

Deformable Attention Kernels

MS Deformable Attention	Multi-scale deformable attention using bilinear interpolation and indirect DMA transpose.
MS Deformable Attention Backward	Backward pass for multi-scale deformable attention.

Transformer Kernels

Transformer TKG

Multi-layer transformer forward pass megakernel for token generation.

Convolution Kernels

Conv1D	1D convolution using tensor engine with replication strategy.
Conv3D	3D convolution using tensor engine with K-replication strategy and W-contiguous tiling.
Depthwise Conv1D	Implements depthwise 1D convolution using implicit GEMM algorithm.
Conv3D Temporal Unroll	3D convolution with temporal unrolling and column tiling for small C_out.
Conv3D Transpose	3D transposed convolution (ConvTranspose3d) via an embedded 3D convolution core with remapped parameters.

Collective Communication Kernels

Fine-Grained All-Gather	Ring-based all-gather for TRN2 with double-buffered collective permute.
FGCC (All-Gather + Matmul)	Fused all-gather and matrix multiplication for TRN2.
SBUF-to-SBUF All-Gather	SBUF-to-SBUF all-gather with variants for small and large tensors.
Collective Communication Kernels	HBM-based all-reduce, all-gather, reduce-scatter, and all-to-all kernels.
Permute A2AV	MoE training dispatch: permute tokens by destination EP rank and exchange via all-to-all-v.
Unpermute A2AV	MoE training combine: all-to-all-v exchange and unpermute to original token order.
QKV Batch Shard	Q projection transition from tensor-parallel to tensor-and-data-parallel layout for batch-sharded attention.

Foreach Kernels

Foreach Elementwise	Suite of fused elementwise operations (add, sub, mul, div, addcdiv, addcmul, lerp, sqrt) with SPMD tiling.
Foreach Norm	L1, L2, and Linf norm computation kernels with SPMD parallelization.

Matmul and MLP MXFP8 Kernels

Matmul MXFP8	Generic matrix multiplication with MXFP8 quantization, supporting BF16 and pre-quantized inputs.
MLP Forward MXFP8	MXFP8 SwiGLU MLP forward pass with optional activation checkpointing.
MLP Backward MXFP8	MXFP8 SwiGLU MLP backward pass with 4-phase gradient computation.
Matmul MXFP8 Backward	Backward pass for a linear layer with MXFP8 quantization, computing both input (dX) and weight (dW) gradients.

MoE Kernels

MX MoE Block TKG Wrapper	Wrapper that bitcasts unsigned integer weights to MX x4 dtype for MoE block.
Blockwise MM Shard-on-Block	Block-sharded blockwise matrix multiplication for context-encoding MoE layers.
RMSNorm Router Top-K A2AV	Fused RMSNorm + router top-K for MoE token generation feeding the all-to-all-v dispatch path.
RMSNorm Router Top-K TKG	Fused RMSNorm (+ optional MX quantize) + router top-K for MoE token generation.
Blockwise MM Forward MXFP8	MXFP8 forward pass for blockwise (dropless) MoE; emits activation checkpoints for the MXFP8 MoE backward.
Build All2All Dispatch Metadata	Builds per-rank send counts (with token deduplication) and displacements from expert_index for MoE all-to-all dispatch.

Indexing Kernels

Gather	Gather rows from a 2D input by a 1D index using indirect DMA load.
Scatter-Add	Scatter-add rows from src into a 2D input by a 1D index.

Top-K Kernels

GPSIMD Top-K

Top-k over the last dimension using the GpSIMD nisa.topk instruction.

Optimizer Kernels

Fused Adam/AdamW

Fused Adam (L2 regularization) and AdamW (decoupled weight decay) optimizer kernels.

Padding Kernels

Pad	Multi-mode tensor padding (constant, replicate, reflect, circular) following PyTorch semantics.

Quantization Kernels

Quantize MXFP8

Block-wise BF16-to-MXFP8 quantization kernel with packed scale support.

RNG Kernels

RNG	Random number generation kernels using GPSIMD engine with state management.

Scan Kernels (State Space Models)

Linear Scan	First-order linear recurrence computation along the last dimension.
Selective Scan	Selective scan (SSM) as in Mamba models.
SSD	State Space Duality scan for Mamba-2 models.

MoE Subkernels

Top-K Reduce	MoE Top-K reduction across sparse all-to-all collective output.
Argsort Unstable	Unstable argsort on 1D input buffer for MoE routing.
Build All-to-All-V Metadata	Builds metadata buffer for all_to_all_v collective from MoE routing decisions.
Permute Routed Tokens	Sorts tokens by expert and packs hidden states for MoE dispatch.

Dynamic Shape Kernels

Dynamic Elementwise Add

Elementwise addition with runtime-variable M-dimension tiling.

Loss Kernels

Cross Entropy

Memory-efficient cross entropy loss forward and backward passes using online log-sum-exp algorithm.

MoE Backward Kernels

Blockwise MM Backward	Computes backward pass for blockwise matrix multiplication in Mixture of Experts layers.
Blockwise MM Backward MXFP8	MXFP8-quantized backward pass for blockwise Mixture of Experts.

This document is relevant for: Trn2, Trn3