neuron-nki-docs

name	neuron-nki-docs
description	Research NKI documentation for API lookups, tutorials, error codes, architecture, and optimization guides. Use when user asks "what does <API> do", "how to do <task>", "what is error <code>", "NKI <API> signature", "find NKI tutorial for <topic>", "look up <symbol>", or needs any NKI documentation reference.
argument-hint	[query or API name]
context	fork

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:

Query Pattern	Start With	Then Read
nl.* or nisa.* or API name	references/indices/symbol-lookup.md	Linked API doc
"how to [task]"	references/indices/task-routing.md	Linked tutorial
"NCC_*" or error code	references/debugging/error-codes-index.md	Specific error doc
gen2/gen3/gen4/trn1/trn2/trn3	references/architecture/	Specific arch doc
optimize/profile/performance	references/optimization/	Relevant guide
browse/list all docs	references/indices/hierarchical-toc.md	Navigate tree

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

Directory	Contents
references/architecture/	Hardware architecture guides for Trainium/Inferentia generations
references/debugging/	Error codes index and individual error documentation
references/downloads/	Valid Python kernel examples (deprecated patterns excluded)
references/indices/	Navigation aids and lookup tables
references/optimization/	Performance tuning, profiling, migration guides
references/programming/	Core NKI concepts, tutorials, and API reference
references/programming/api/	Detailed API documentation by category
references/programming/tutorials/	Step-by-step kernel implementation tutorials
references/reference/	FAQ and release notes

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:

Constraint	Limit	Notes
Partition dimension (P)	≤ 128	First axis of on-chip tensors
PSUM free dimension (F)	≤ 512 (gen2/3) / ≤ 4096 (gen4)	Second axis in PSUM buffer
SBUF free dimension (F)	≤ 32767	Second axis in SBUF buffer
MatMul K dimension	≤ 2048	Contraction dimension per tile

Hardware Generations

Generation	Devices	Key Features
gen2 (v2)	Trn1, Inf2	Baseline NKI support
gen3 (v3)	Trn2	FP8 support, Double FP8 mode
gen4 (v4)	Trn3	MXFP8/MXFP4, Quad-MX mode

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

Skill	Purpose
/neuron-nki-writing	Write NKI kernels from specifications
/neuron-nki-debugging	Debug compiler errors on device
/neuron-nki-profiling	Profile kernel performance
/neuron-nki-profile-querying	Query and analyze kernel profile data