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.

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:

For tensors exceeding limits, use explicit tiling with TiledRange for remainder-safe iteration (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

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
• K-dimension loop: always nl.affine_range(), never nl.sequential_range (serializes execution)
• Multiple nisa.nc_matmul() writes to same PSUM buffer triggers hardware accumulation
• 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

References are tiered to minimize overhead on simple tasks. Load only what you need based on the Complexity Assessment above.

Always load (core references):

• references/nki-language-constraint.md - MANDATORY: Required and forbidden API patterns for NKI 0.4.0, reference kernel template
• references/common-patterns.md - Full code examples: matmul PSUM accumulation, fused ScalarE, associative scan, production patterns
• references/api-translation.md - PyTorch/NumPy to NKI operation mapping
• references/kernel-template.md - Standard kernel template with self-contained utilities
• references/indexing-patterns.md - Complete indexing guide: memory-type rules (HBM/SBUF/PSUM), operation constraints (matmul/transpose/reduce), dynamic indexing with DGE modes

Load when tiling or DMA patterns are needed (medium+ complexity):

• references/memory-patterns.md - DMA and tiling patterns with code examples
• references/nkilib/core/tiled-range.md - TiledRange: dimension tiling with remainder handling
• references/nkilib/core/kernel-helpers.md - Math helpers, SPMD, dtype utilities

Load when layout manipulation is needed:

• references/transpose-and-layout.md - Transpose and layout transformation guide: nc_transpose, TensorView, array patterns, strided DMA, decision trees for layout operations
• references/nkilib/core/tensor-view.md - TensorView: zero-copy tensor manipulation