.. meta:: :description: Reference for the pre-built NKI Library kernels included with the AWS Neuron SDK. :date-modified: 05/21/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``. 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:`KV-Parallel Segmented Prefill ` - KV-parallel segmented prefill attention kernel. 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. Transformer Kernels ~~~~~~~~~~~~~~~~~~~~ .. list-table:: :widths: 40 60 * - :doc:`Transformer TKG ` - Multi-layer transformer forward pass megakernel for token generation. 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. 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. 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. 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. 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. .. toctree:: :maxdepth: 1 :hidden: Argsort Unstable Attention Block TKG Attention CTE Attention Segmented CTE Attention TKG Blockwise MM Backward Build All-to-All-V Metadata Conv1D Conv3D Cross Entropy Cumsum Depthwise Conv1D Dynamic Elementwise Add FGCC Find Nonzero Indices Fine-Grained All-Gather Foreach Elementwise Foreach Norm FP8 Quantize Fused Adam KV-Parallel Segmented Prefill Linear Scan Matmul MXFP8 MLP MLP Backward MXFP8 MLP Forward MXFP8 MoE CTE MoE TKG MX MoE Block TKG Wrapper Output Projection CTE Output Projection TKG Pad Permute Routed Tokens QKV Quantize MXFP8 Ring Attention Backward Ring Attention Forward RMSNorm-Quant RNG RoPE Router Top-K SBUF-to-SBUF All-Gather Selective Scan SSD Top-K Reduce Transformer TKG