aws-neuron/neuron-agentic-development/neuron-nki-writing

Writing NKI Kernels

This skill guides writing and modifying NKI (Neuron Kernel Interface) kernels — from new kernel creation (PyTorch/NumPy/natural language translation) to editing existing kernels (adding shape/dtype support, refactoring tiling, implementing new features). Focus on correctness using documented APIs.

Critical: NKI Language Constraints

BEFORE writing any NKI code, read references/nki-language-constraint.md for the complete list of required and forbidden API patterns. Each rule is labelled with the release that made it a breaking change (0.x.0+), so the file covers every current release rather than one migration step. Violating ANY rule is a compilation failure.

Hardware limits: don't hardcode the PSUM free dim

The PSUM free-dimension limit is generation- and dtype-gated and has changed across releases. Current values:

• gen2 / gen3: 512 (one PSUM bank)
• gen4: 4096 for a float32 dst, 8192 for a bfloat16 dst (entire PSUM)

This is not a queryable tile_size field — nl.tile_size.psum_bank_fmax is a fixed 512 (fp32 elements per bank) and is neither dtype- nor generation-aware, so do not use it as the free-dim limit. Authoritative source: the nc_matmul, nc_matmul_mx, and dma_transpose API blocks via /neuron-nki-docs. Treat any PSUM free-dim number inlined in this skill's reference files as illustrative; confirm against those API blocks.

tile_size does authoritatively report other limits — use it for those: nl.tile_size.pmax (128), nl.tile_size.psum_num_banks (bank cycling), nl.tile_size.gemm_moving_fmax (matmul moving-operand SBUF free dim), and nl.tile_size.sbuf_fmax / sbuf_fmax_bytes (SBUF capacity).

Quick Start

Minimal working kernel structure:

import nki
import nki.isa as nisa
import nki.language as nl
@nki.jit
def my_kernel(input_hbm: nl.ndarray) -> nl.ndarray:
    """One-line description of kernel operation."""
    # 1. Allocate SBUF tile
    tile = nl.ndarray(input_hbm.shape, dtype=input_hbm.dtype, buffer=nl.sbuf)
    # 2. Load from HBM to SBUF
    nisa.dma_copy(dst=tile, src=input_hbm[0:input_hbm.shape[0], 0:input_hbm.shape[1]])
    # 3. Compute (example: element-wise exp)
    result = nl.ndarray(tile.shape, dtype=tile.dtype, buffer=nl.sbuf)
    nisa.activation(dst=result, data=tile, op=nl.exp)
    # 4. Allocate and store to HBM
    output = nl.ndarray(input_hbm.shape, dtype=input_hbm.dtype, buffer=nl.shared_hbm)
    nisa.dma_copy(dst=output, src=result)
    return output

Complexity Assessment

Before reading references, classify the task to avoid unnecessary overhead:

Simple (element-wise op, single reduction, activation, layernorm, add/multiply):

• Use the Quick Start template and Step 4 API table directly
• Skip utility library references entirely
• Start writing code immediately — consult references only when stuck
• Target: working kernel within 5 minutes

Medium (matmul, softmax, multi-step fusion, transpose with tiling):

• Read references/common-patterns.md, references/api-translation.md and references/memory-patterns.md
• Skip utility library references unless tiling is complex
• Target: working kernel within 15 minutes

Complex (multi-head attention, transformer blocks, state-space models, MoE):

• Full reference loading appropriate
• Read utility selection guide and relevant patterns
• Target: working kernel within 30 minutes

Start writing code as soon as possible. Reference reading should supplement coding, not precede it. Write the kernel structure first, then consult references for specific API details as needed.

Translation Workflow

Step 1: Identify Operations

Map PyTorch/NumPy operations to NKI equivalents using references/api-translation.md.

Step 2: Design Tiling Strategy

NKI operates on tiles with hardware constraints:

SBUF free dim and MatMul K (contraction) are not fixed constants — look

them up rather than hardcoding. SBUF free-dim capacity is a per-partition byte

limit (nl.tile_size.sbuf_fmax / sbuf_fmax_bytes, generation-/dtype-dependent).

Per-instruction MatMul K is the partition dim, K ≤ 128 (effective 256/512

with double_row / Quad-MX packing). Authoritative source: the nc_matmul

tile-size block in the /neuron-nki-docs skill (api-nki-isa-tensor.md).

For tensors exceeding these limits, tile explicitly. Two mechanisms, both remainder-safe:

• `nl.tiles` / `nl.blocks` (NKI 0.7.0+, language-native) — view the tensor as a grid

addressed by tile coordinate, e.g. src_tiles = nl.tiles(src, tile_shape=(128, 256)) then src_tiles[i, j].data. Prefer this when the alternative is hand-computed offsets and strides, and use nl.alloc_tiles / nl.alloc_blocks to allocate on-chip grids. Boundary tiles clamp to the real extent automatically. See references/indexing-patterns.md ("Tiled and Blocked Layouts").

• `TiledRange` (nkilib helper) — iterate a tiled dimension, giving .start_offset /

.size per step for plain slicing. Still valid and widely used in nkilib kernels; see Utility Selection Guide below.

Step 3: Implement Memory Access

Consult the Utility Selection Guide to choose the right utilities for the kernel's access patterns, then follow references/memory-patterns.md and references/transpose-and-layout.md:

• Contiguous DMA: For aligned, sequential access (most efficient)
• Strided DMA: Use TensorView.slice(step=N) for gather/scatter patterns (see transpose-and-layout.md)
• Tiling loops: Use TiledRange for dimensions requiring tiling with remainder handling
• Transpose operations: Use nisa.nc_transpose() for P↔F transpose, TensorView for layout manipulation
• Partition broadcast: Use stream_shuffle_broadcast when a bias/scale in partition 0 must reach all PEs
• Buffer management: Use SbufManager when the kernel has 4+ SBUF allocations or shared sub-functions

Step 4: Add Compute Operations

Use ISA functions with explicit dst parameter:

# Element-wise
nisa.activation(dst=result, data=input, op=nl.exp)
nisa.tensor_tensor(dst=result, data1=a, data2=b, op=nl.add)
nisa.tensor_scalar(dst=result, data=input, op0=nl.multiply, operand0=2.0)
# Reductions
nisa.tensor_reduce(dst=result, data=input, op=nl.add, axis=1)
# Matrix multiply
nisa.nc_matmul(dst=psum_result, stationary=a, moving=b)

Step 5: Validate Complex Translations

Important: Always compute reference results on CPU, not on the XLA device. Every on-device XLA graph generates a separate NEFF file, making it hard to identify the NKI kernel's NEFF during profiling. Use tensor.cpu() before computing PyTorch/NumPy references.

For simple kernels (single operation, few tiles), comparing the final output against a PyTorch/NumPy reference is usually sufficient. For complex kernels with multiple computation stages, validate incrementally:

1. Identify logical stages in the source operation. For example, a fused attention kernel has: QK matmul → scale → mask → softmax → AV matmul.
2. Translate and validate one stage at a time. Write the first stage, store its output to HBM, and compare against the corresponding PyTorch intermediate. Only proceed to the next stage once the current one matches.
3. Compose validated stages. Once each stage is verified independently, connect them (keeping intermediates in SBUF instead of round-tripping through HBM) and validate the final output.

This catches translation errors early — a mismatch in the final output of a 5-stage kernel is much harder to diagnose than a mismatch after stage 2.

Use multiple complementary checks (atol/rtol, max absolute difference, tensor norm of the difference, cosine similarity) rather than relying on a single metric.

Hardware Constraints Quick Reference

Loop Types

`nl.affine_range` / `nl.sequential_range` / `nl.static_range` are deprecated aliases for `range()` — use `range()`. All three are literally return range(start, stop, step) on the default frontend and in the CPU simulator, and each carries an explicit "This API is deprecated and will be removed in future releases. Prefer using `range()` directly instead" warning. They do not differ in unrolling or scheduling:

• There is no "no unroll" or "partial unroll" variant. Every one of them fully unrolls.
• nl.sequential_range does not serialize execution or express a loop-carried dependency.

Choosing it over nl.affine_range changes nothing. To constrain scheduling, use the nl.no_reorder() region; to carry state across iterations, keep it in SBUF/PSUM and mutate in place — correctness comes from the data dependency, not from the loop function.

Runtime/data-dependent loops (NKI 0.7.0+): use nl.fori_loop (counted, runtime bound) or nl.while_loop (condition-driven), NOT for i in nl.dynamic_range(...) or bare while reg: — those cannot work on the default standard frontend, where a runtime register has no value at trace time. See references/nki-language-constraint.md for the before/after patterns and rules, and /neuron-nki-docs (api-nki-language-dims.md) for the fori_loop / while_loop API signatures.

Common Patterns

For detailed code examples, anti-patterns, and production patterns (cumsum, rmsnorm_quant), see references/common-patterns.md.

Element-wise Operations

• Reshape to 2D, tile P dimension (≤128), use nisa.activation() / nisa.tensor_tensor()

Matrix Multiply — Key Rules

• Allocate PSUM: nl.ndarray(..., buffer=nl.psum) — uninitialized is correct; do NOT `nisa.memset` the PSUM tile before the K loop
• K-dimension loop: plain range() — the deprecated nl.*_range aliases are the same function and none of them serializes
• Control accumulation explicitly with accumulate=(k_idx > 0): the first K tile overwrites (initializes) PSUM, later tiles accumulate — this replaces any separate memset/zero-init of PSUM
• Never write PSUM to HBM between accumulation steps
• Operands: stationary [K≤128, M≤128], moving [K≤128, N≤512/4096], result [M, N] in PSUM
• Copy PSUM→SBUF via nisa.tensor_copy() before further ops
• See examples/simple_matmul.py for complete examples

Fused ScalarE Operations

• nisa.activation(op=nl.exp, data=x, scale=s) → exp(x * s) in one instruction
• Available: scale, bias, or both before any activation function

Sequential Operations & Associative Scan

• Use nisa.tensor_tensor_scan instead of explicit sequential loops
• Pattern: out[i] = op0(data[i], out[i-1]) op1 data1[i]
• Multi-tile: pass final state as initial= to next tile's scan
• See examples/associative_scan.py for complete pattern

Skill References