.. meta:: :description: Reference for the pre-built NKI Library kernels included with the AWS Neuron SDK. :date-modified: 10/19/2026 .. _nkl_api_ref_home: 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**: https://github.com/aws-neuron/nki-library Core Kernels ------------- Normalization and Quantization Kernels ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ .. list-table:: :widths: 40 60 * - :doc:`RMSNorm-Quant ` - Performs optional RMS normalization followed by quantization to ``fp8``. * - :doc:`RMSNorm MX Prefill ` - Fused RMSNorm + MX quantization (+ optional router top-K) for prefill, with optional residual add. QKV Projection Kernels ~~~~~~~~~~~~~~~~~~~~~~~~ .. list-table:: :widths: 40 60 * - :doc:`QKV ` - Performs Query-Key-Value projection with optional normalization and RoPE fusion. Attention Kernels ~~~~~~~~~~~~~~~~~ .. list-table:: :widths: 40 60 * - :doc:`Attention CTE ` - Implements attention optimized for Context Encoding (prefill) use cases. * - :doc:`Attention Segmented CTE ` - Segmented attention with block-based KV cache and prefix caching for decode. * - :doc:`Attention TKG ` - Implements attention optimized for Token Generation (decode) use cases with small active sequence lengths. * - :doc:`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 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ .. list-table:: :widths: 40 60 * - :doc:`RoPE ` - Applies Rotary Position Embedding to input embeddings with flexible layout support. Multi-Layer Perceptron (MLP) Kernels ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ .. list-table:: :widths: 40 60 * - :doc:`MLP ` - Implements Multi-Layer Perceptron with optional normalization fusion and quantization support. Output Projection Kernels ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ .. list-table:: :widths: 40 60 * - :doc:`Output Projection CTE ` - Computes output projection optimized for Context Encoding use cases. * - :doc:`Output Projection TKG ` - Computes output projection optimized for Token Generation use cases. Mixture of Experts (MoE) Kernels ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ .. list-table:: :widths: 40 60 * - :doc:`Router Top-K ` - Computes router logits, applies activation functions, and performs top-K selection for MoE models. * - :doc:`MoE CTE ` - Implements Mixture of Experts MLP operations optimized for Context Encoding use cases. * - :doc:`MoE TKG ` - Implements Mixture of Experts MLP operations optimized for Token Generation use cases. Quantization Kernels ~~~~~~~~~~~~~~~~~~~~~ .. list-table:: :widths: 40 60 * - :doc:`FP8 Quantize ` - Static and row-wise dynamic FP8 quantization with pre-combined dequantization scale support. Cumulative Sum Kernels ~~~~~~~~~~~~~~~~~~~~~~~ .. list-table:: :widths: 40 60 * - :doc:`Cumsum ` - Computes cumulative sum along the last dimension with optimized tiling. Core Subkernels ~~~~~~~~~~~~~~~~ .. list-table:: :widths: 40 60 * - :doc:`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 ~~~~~~~~~~~~~~~~~ .. list-table:: :widths: 40 60 * - :doc:`Attention Block TKG ` - Fused attention block for Token Generation that keeps all intermediate tensors in SBUF to minimize HBM traffic. * - :doc:`Ring Attention Forward ` - Ring attention forward pass for context parallelism across multiple workers. * - :doc:`Ring Attention Backward ` - Ring attention backward pass SPMD kernel for context parallelism. * - :doc:`Ring Attention Unpermute ` - Reorders striped ring-attention output back to contiguous sequence order. * - :doc:`QKV CTE MLA ` - DeepSeek Multi-head Latent Attention (MLA) QKV projection with MX quantization for context encoding. * - :doc:`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. * - :doc:`Attention MXFP8 TKG ` - MXFP8 flash-decode attention over a quantized block KV cache with packed-Q eviction (Trainium 3). * - :doc:`Sparse Attention Indexer (MX + BF16 Score) ` - DeepSeek Sparse Attention Indexer with MX projections and a BF16 score path. * - :doc:`MLA QKV CTE ` - DeepSeek MLA QKV projection (MX) for context encoding, emitting absorbed latents. * - :doc:`MLA Sparse Attention CTE ` - Standalone sparse latent + RoPE attention (kernel A of the split DeepSeek-V3.2 sparse-MLA forward). * - :doc:`MLA V-Up O-Proj CTE ` - Standalone MX V-up + MX output projection (kernel B of the split DeepSeek-V3.2 sparse-MLA forward). * - :doc:`Attention Const-Max ` - Attention with a caller-supplied global softmax maximum, supporting sequence-major Q/K layouts. * - :doc:`Flash Attention 2 ` - FlashAttention-2 forward for unmasked bidirectional attention over long key sequences (Trainium 3). * - :doc:`MLA Decode QKV ` - DeepSeek-V3.2 MLA query and key/value projections for decode, with RoPE, query absorption, and paged KV-cache write. * - :doc:`MLA Decode Attention ` - DeepSeek-V3.2 MLA decode attention over indexer-selected keys, with V up-projection and output projection. * - :doc:`MLA Attention Block Gather Decode ` - End-to-end DeepSeek-V3.2 MLA decode attention block (``mla_decode_qkv`` + ``mla_decode_attention``). * - :doc:`DeepSeek-V3.2 Sparse Indexer ` - DeepSeek Sparse Attention lightning indexer for decode (Trainium 3). Linear Attention Kernels ~~~~~~~~~~~~~~~~~~~~~~~~ .. list-table:: :widths: 40 60 * - :doc:`KDA TKG ` - Kimi Delta Attention token generation: advances the KDA recurrence by one token per request. * - :doc:`KDA Block TKG ` - Full Kimi Delta Attention block for token generation with all intermediates kept in SBUF. * - :doc:`GDN CTE ` - Chunked gated delta-rule (GDN) prefill with a selectable intra-chunk inverse algorithm. * - :doc:`GDN TKG ` - Fused single-token gated delta-rule decode step. * - :doc:`GDN Block TKG ` - Fused Qwen3.5 GDN decode layer (input projections, conv1d, GQA, gated delta-rule). * - :doc:`GDN Conv1D ` - Depthwise causal conv1d + SiLU + ``conv_state`` update for GDN prefill and decode. Deformable Attention Kernels ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ .. list-table:: :widths: 40 60 * - :doc:`MS Deformable Attention ` - Multi-scale deformable attention using bilinear interpolation and indirect DMA transpose. * - :doc:`MS Deformable Attention Backward ` - Backward pass for multi-scale deformable attention. Transformer Kernels ~~~~~~~~~~~~~~~~~~~~ .. list-table:: :widths: 40 60 * - :doc:`Transformer TKG ` - Multi-layer transformer forward pass megakernel for token generation. * - :doc:`GPT-OSS Tail Chunked Megakernel ` - GPT-OSS decode sampling tail (RMSNorm, ``lm_head``, top-K, sampler) with a single all-to-all collective. Convolution Kernels ~~~~~~~~~~~~~~~~~~~~ .. list-table:: :widths: 40 60 * - :doc:`Conv1D ` - 1D convolution using tensor engine with replication strategy. * - :doc:`Conv3D ` - 3D convolution using tensor engine with K-replication strategy and W-contiguous tiling. * - :doc:`Depthwise Conv1D ` - Implements depthwise 1D convolution using implicit GEMM algorithm. * - :doc:`Conv3D Temporal Unroll ` - 3D convolution with temporal unrolling and column tiling for small ``C_out``. * - :doc:`Conv3D Transpose ` - 3D transposed convolution (ConvTranspose3d) via an embedded 3D convolution core with remapped parameters. * - :doc:`Conv3D Backward ` - Input, weight, and bias gradients of a 3D convolution with a geometry-driven planner. Normalization Kernels ~~~~~~~~~~~~~~~~~~~~~~ .. list-table:: :widths: 40 60 * - :doc:`Batch Norm Backward ` - Batch normalization backward over the channel dimension, with optional fused ReLU. * - :doc:`RMSNorm Residual Prefill ` - Pre-norm residual add followed by RMSNorm for prefill, with the add performed during the DMA load. Pooling Kernels ~~~~~~~~~~~~~~~~ .. list-table:: :widths: 40 60 * - :doc:`Max Pooling 3D ` - 3D max pooling over (D, H, W) in NCDHW or NDHWC layout. * - :doc:`Max Pooling 3D Backward ` - Backward pass of 3D max pooling. * - :doc:`Average Pooling 3D ` - 3D average pooling over (D, H, W) in NCDHW or NDHWC layout. * - :doc:`Average Pooling 3D Backward ` - Backward pass of 3D average pooling. Collective Communication Kernels ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ .. list-table:: :widths: 40 60 * - :doc:`Fine-Grained All-Gather ` - Ring-based all-gather for TRN2 with double-buffered collective permute. * - :doc:`FGCC (All-Gather + Matmul) ` - Fused all-gather and matrix multiplication for TRN2. * - :doc:`SBUF-to-SBUF All-Gather ` - SBUF-to-SBUF all-gather with variants for small and large tensors. * - :doc:`Collective Communication Kernels ` - HBM-based all-reduce, all-gather, reduce-scatter, and all-to-all kernels. * - :doc:`Permute A2AV ` - MoE training dispatch: permute tokens by destination EP rank and exchange via all-to-all-v. * - :doc:`Unpermute A2AV ` - MoE training combine: all-to-all-v exchange and unpermute to original token order. * - :doc:`Build Packed Routing ` - Builds the destination-major ``packed_source_indices`` consumed by ``permute_a2av`` and ``unpermute_a2av``. * - :doc:`QKV Batch Shard ` - Q projection transition from tensor-parallel to tensor-and-data-parallel layout for batch-sharded attention. Foreach Kernels ~~~~~~~~~~~~~~~~ .. list-table:: :widths: 40 60 * - :doc:`Foreach Elementwise ` - Suite of fused elementwise operations (add, sub, mul, div, addcdiv, addcmul, lerp, sqrt) with SPMD tiling. * - :doc:`Foreach Norm ` - L1, L2, and Linf norm computation kernels with SPMD parallelization. Matmul and MLP MXFP8 Kernels ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ .. list-table:: :widths: 40 60 * - :doc:`Matmul MXFP8 ` - Generic matrix multiplication with MXFP8 quantization, supporting BF16 and pre-quantized inputs. * - :doc:`MLP Forward MXFP8 ` - MXFP8 SwiGLU MLP forward pass with optional activation checkpointing. * - :doc:`MLP Backward MXFP8 ` - MXFP8 SwiGLU MLP backward pass with 4-phase gradient computation. * - :doc:`Matmul MXFP8 Backward ` - Backward pass for a linear layer with MXFP8 quantization, computing both input (dX) and weight (dW) gradients. * - :doc:`MLP DeepSeek MX ` - DeepSeek-V3.2 shared-experts and first dense MLP layers for prefill with MXFP8 weights and packed MX input. MoE Kernels ~~~~~~~~~~~~~ .. list-table:: :widths: 40 60 * - :doc:`MX MoE Block TKG Wrapper ` - Wrapper that bitcasts unsigned integer weights to MX x4 dtype for MoE block. * - :doc:`Blockwise MM Shard-on-Block ` - Block-sharded blockwise matrix multiplication for context-encoding MoE layers. * - :doc:`RMSNorm Router Top-K A2AV ` - Fused RMSNorm + router top-K for MoE token generation feeding the all-to-all-v dispatch path. * - :doc:`RMSNorm Router Top-K TKG ` - Fused RMSNorm (+ optional MX quantize) + router top-K for MoE token generation. * - :doc:`Blockwise MM Forward MXFP8 ` - MXFP8 forward pass for blockwise (dropless) MoE; emits activation checkpoints for the MXFP8 MoE backward. * - :doc:`Build All2All Dispatch Metadata ` - Builds per-rank send counts (with token deduplication) and displacements from ``expert_index`` for MoE all-to-all dispatch. * - :doc:`DeepSeek-V3 Router ` - DeepSeek-V3.2 group-limited ``noaux_tc`` routing for one decode step. * - :doc:`Blockwise MM Block MX ` - DeepSeek-V3.2 blockwise MXFP4/MXFP8 MoE for context encoding on a single logical core (LNC1). Indexing Kernels ~~~~~~~~~~~~~~~~~ .. list-table:: :widths: 40 60 * - :doc:`Gather ` - Gather rows from a 2D input by a 1D index using indirect DMA load. * - :doc:`Scatter-Add ` - Scatter-add rows from src into a 2D input by a 1D index. Top-K Kernels ~~~~~~~~~~~~~ .. list-table:: :widths: 40 60 * - :doc:`GPSIMD Top-K ` - Top-k over the last dimension using the GpSIMD ``nisa.topk`` instruction. * - :doc:`Chunked Top-K ` - Single-core top-K over a vocabulary wider than the ``nisa.topk`` limit, with no collective. Optimizer Kernels ~~~~~~~~~~~~~~~~~~ .. list-table:: :widths: 40 60 * - :doc:`Fused Adam/AdamW ` - Fused Adam (L2 regularization) and AdamW (decoupled weight decay) optimizer kernels. Padding Kernels ~~~~~~~~~~~~~~~~ .. list-table:: :widths: 40 60 * - :doc:`Pad ` - Multi-mode tensor padding (constant, replicate, reflect, circular) following PyTorch semantics. Quantization Kernels ~~~~~~~~~~~~~~~~~~~~~ .. list-table:: :widths: 40 60 * - :doc:`Quantize MXFP8 ` - Block-wise BF16-to-MXFP8 quantization kernel with packed scale support. RNG Kernels ~~~~~~~~~~~~ .. list-table:: :widths: 40 60 * - :doc:`RNG ` - Random number generation kernels using GPSIMD engine with state management. Scan Kernels (State Space Models) ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ .. list-table:: :widths: 40 60 * - :doc:`Linear Scan ` - First-order linear recurrence computation along the last dimension. * - :doc:`Selective Scan ` - Selective scan (SSM) as in Mamba models. * - :doc:`SSD ` - State Space Duality scan for Mamba-2 models. MoE Subkernels ~~~~~~~~~~~~~~~ .. list-table:: :widths: 40 60 * - :doc:`Top-K Reduce ` - MoE Top-K reduction across sparse all-to-all collective output. * - :doc:`Argsort Unstable ` - Unstable argsort on 1D input buffer for MoE routing. * - :doc:`Build All-to-All-V Metadata ` - Builds metadata buffer for all_to_all_v collective from MoE routing decisions. * - :doc:`Permute Routed Tokens ` - Sorts tokens by expert and packs hidden states for MoE dispatch. Dynamic Shape Kernels ~~~~~~~~~~~~~~~~~~~~~~ .. list-table:: :widths: 40 60 * - :doc:`Dynamic Elementwise Add ` - Elementwise addition with runtime-variable M-dimension tiling. Loss Kernels ~~~~~~~~~~~~~ .. list-table:: :widths: 40 60 * - :doc:`Cross Entropy ` - Memory-efficient cross entropy loss forward and backward passes using online log-sum-exp algorithm. MoE Backward Kernels ~~~~~~~~~~~~~~~~~~~~~~ .. list-table:: :widths: 40 60 * - :doc:`Blockwise MM Backward ` - Computes backward pass for blockwise matrix multiplication in Mixture of Experts layers. * - :doc:`Blockwise MM Backward MXFP8 ` - MXFP8-quantized backward pass for blockwise Mixture of Experts. .. toctree:: :maxdepth: 1 :hidden: Argsort Unstable Attention Block TKG Attention Const-Max Attention CTE Attention KV-Parallel Segmented CTE Attention MXFP8 TKG Attention Segmented CTE Attention TKG Average Pooling 3D Average Pooling 3D Backward Batch Norm Backward Blockwise MM Backward Blockwise MM Backward MXFP8 Blockwise MM Block MX Blockwise MM Forward MXFP8 Blockwise MM Shard-on-Block Build All-to-All-V Metadata Build All2All Dispatch Metadata Build Packed Routing Chunked Top-K Collective Communication Kernels Conv1D Conv3D Conv3D Backward Conv3D Temporal Unroll Conv3D Transpose Cross Entropy Cumsum DeepSeek-V3 Router DeepSeek-V3.2 Sparse Indexer Depthwise Conv1D Dynamic Elementwise Add FGCC Find Nonzero Indices Fine-Grained All-Gather Flash Attention 2 Foreach Elementwise Foreach Norm FP8 Quantize Fused Adam Gather GDN Block TKG GDN Conv1D GDN CTE GDN TKG GPSIMD Top-K GPT-OSS Tail Chunked Megakernel KDA Block TKG KDA TKG Linear Scan Matmul MXFP8 Matmul MXFP8 Backward Max Pooling 3D Max Pooling 3D Backward MLA Attention Block Gather Decode MLA Decode Attention MLA Decode QKV MLA QKV CTE MLA Sparse Attention CTE MLA V-Up O-Proj CTE MLP MLP Backward MXFP8 MLP DeepSeek MX MLP Forward MXFP8 MoE CTE MoE TKG MS Deformable Attention MS Deformable Attention Backward MX MoE Block TKG Wrapper Output Projection CTE Output Projection TKG Pad Permute A2AV Permute Routed Tokens QKV QKV Batch Shard QKV CTE MLA Quantize MXFP8 Ring Attention Backward Ring Attention Forward Ring Attention Unpermute RMSNorm MX Prefill RMSNorm Residual Prefill RMSNorm Router Top-K A2AV RMSNorm Router Top-K TKG RMSNorm-Quant RNG RoPE Router Top-K SBUF-to-SBUF All-Gather Scatter-Add Selective Scan Sparse Attention Indexer MX BF16Score SSD SWA Fused CTE Top-K Reduce Transformer TKG Unpermute A2AV