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 covering Beta 1 → Beta 2, Beta 2 → NKI 0.3.0, and NKI 0.3.0 → NKI 0.4.0 migration rules. 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