Research $ARGUMENTS thoroughly:

1. Find relevant files through provided indices using Glob and Grep
2. Read and analyze document and code
3. Summarize findings with specific file references

NKI Documentation Lookup

This skill provides comprehensive access to NKI (Neuron Kernel Interface) documentation for AWS Trainium/Inferentia kernel development. Use this skill to answer questions about NKI APIs, tutorials, hardware architecture, error codes, and optimization techniques.

Query Type Routing

Route queries to the appropriate documentation based on the query pattern:

Search Strategy

1. Start with indices - Use the index files in references/indices/ to find the exact documentation needed
2. Follow links - Read the linked documentation file for full details
3. Cross-reference - For complex topics, combine information from multiple sources

Index Files

• indices/symbol-lookup.md - Alphabetical API function index with direct links
• indices/task-routing.md - Task-oriented guide ("I want to...")
• indices/hierarchical-toc.md - Complete documentation tree structure

Directory Guide

Quick API Reference

Most Common APIs

Data Movement:

• nki.isa.dma_copy - Load/store data between HBM and SBUF
• nki.isa.dma_transpose - Transpose during DMA transfer

Tensor Operations:

• nki.isa.nc_matmul - Matrix multiplication on Tensor Engine
• nki.isa.tensor_tensor - Element-wise operations
• nki.isa.tensor_scalar - Broadcast scalar operations
• nki.isa.tensor_reduce - Reduction along axes

Memory Allocation:

• nl.ndarray - Create tensor in SBUF
• nl.zeros - Create zero-initialized tensor

Loop Constructs:

• range - Standard loop iterator (recommended)
• nl.affine_range / nl.sequential_range / nl.static_range - Legacy aliases for range (all have identical effect in NKI 0.3.0+)

SPMD:

• nl.program_id - Get current program index
• nl.num_programs - Get total program count

Module Aliases

The documentation uses these standard aliases:

import nki
import nki.language as nl
import nki.isa as nisa

Activation Function Operators

Important: Symbols like nl.exp, nl.sigmoid, nl.tanh, etc. are op specifiers, not standalone callable functions. They must be passed as the op parameter to nisa.activation():

# CORRECT: Use as op parameter
result = nisa.activation(op=nl.exp, data=tensor, scale=scale_tensor, bias=bias_tensor)

# INCORRECT: These are NOT callable functions
# result = nl.exp(tensor)  # This does NOT work

Hardware Constraints

Critical limits to remember when answering questions:

The PSUM free-dimension limit is generation- and dtype-gated: 512 on gen2/gen3 (one PSUM bank), and on gen4 up to 4096 for a float32 dst or 8192 for a bfloat16 dst (the entire PSUM). It is not a queryable tile_size field (psum_bank_fmax is a fixed 512). Authoritative source: the nc_matmul / nc_matmul_mx API blocks in programming/api/.

The per-instruction MatMul K (contraction) limit is the partition dimension: K ≤ 128 (the systolic array is 128×128). double_row (FP8, gen3/gen4) and Quad-MX (gen4) pack the contraction across the partition dim and a small free dim, giving an effective K of 256 / 512 respectively. Authoritative source: the nc_matmul tile-size block in programming/api/api-nki-isa-tensor.md. The SBUF free-dimension capacity is a per-partition byte limit (nl.tile_size.sbuf_fmax / sbuf_fmax_bytes, generation- and dtype-dependent), not a fixed element count.

Hardware Generations

Source NKI Kernel

The Python example files in references/downloads/ use the latest NKI API patterns and can be used as reference implementations:

• average_pool2d_nki_kernels.py — 2D average pooling kernel
• matrix_multiplication_nki_kernels.py — Tiled matrix multiplication kernel
• transpose2d_nki_kernels.py — 2D transpose kernel
• mamba_nki_kernels.py — Fused Mamba state-space model kernel
• test_nki_isa_local_gather.py — Local gather ISA test

Note: The fused_mamba tutorial (references/programming/tutorials/fused_mamba.md) still uses Beta 1 syntax (nl.load/nl.store). The optimization concepts (loop reordering, data reuse, tiling for temporal locality) are valid, but do NOT copy its code patterns — use nisa.dma_copy from /neuron-nki-writing for any code generation.

Common Query Examples

API Lookup

Query: "What is nisa.nc_matmul?" → Read indices/symbol-lookup.md → Find link → Read programming/api/api-nki-isa-tensor.md#nki-isa-nc_matmul

Error Code

Query: "What does NCC_EVRF001 mean?" → Read debugging/error-codes-index.md → Read debugging/error-codes/EVRF001.md

Tutorial

Query: "How do I implement matrix multiplication?" → Read indices/task-routing.md → Read programming/tutorials/matrix_multiplication.md

Architecture

Query: "What's different in Trainium3?" → Read architecture/trainium3_arch.md

Optimization

Query: "How do I profile my kernel?" → Read optimization/use-neuron-profile.md

Response Guidelines

When answering NKI documentation queries:

1. Cite sources - Reference the specific documentation file used
2. Include signatures - For API queries, include function signatures and key parameters
3. Note hardware requirements - Mention if APIs are generation-specific (v2/v3/v4)
4. Link related docs - Point to related tutorials or guides for context
5. Warn about deprecated patterns - If documentation shows old patterns, note this

Error Code Format

NKI compiler errors follow the pattern NCC_<category><number>:

• EOOM* - Out of memory errors
• EVRF* - Verification/validation errors
• EUOC* - Unsupported operation errors
• EBVF* - Backend verification failures
• ESPP* - Shape/type errors

The full index is in debugging/error-codes-index.md with individual files in debugging/error-codes/.

Related Skills