[mark14wu]

Stacked on #539 (base branch fix/nki-official-kernels); retarget to main after #539 merges.

Summary

#539 lets the legacy NKI interpreter (frontend="nki") run kernels written in the style of AWS's NKI samples, but the samples' larger kernels use NKI ISA instructions (nisa.*) that the interpreter does not implement. Without them, the kernel calls NKI's real implementation and fails with RuntimeError: calling NKI API outside of NKI kernels is not supported. This PR implements them, plus the nl operations those kernels use, with NKI's semantics:

• nisa: nc_matmul (stationary.T @ moving, fp32), nc_transpose, activation / activation_reduce (op(data * scale + bias), optional reduction into reduce_res), tensor_tensor, tensor_scalar / tensor_scalar_reduce, scalar_tensor_tensor, tensor_reduce, tensor_tensor_scan, tensor_copy, memset, affine_select, reciprocal, iota, dma_copy, dma_transpose, get_nc_version (gen2, a Trn1 NeuronCore).
• nl: ds, num_programs, program_ndim, full, ones, shared_constant, shared_identity_matrix, transpose, floor, mod.
• Operators given as numpy ufuncs or as nl functions (nl.add, nl.max, nl.sum, ...) map to numpy ufuncs by name.
• nisa.dma_copy and nisa.dma_transpose read and write through load / store, so analysis clients see their HBM accesses as loads and stores. NKITransformer rewrites nisa.dma_copy(dst=o[...], src=x[...]) and nisa.dma_transpose(x[...]) to pass the subscripts as dst_keys / src_keys.

The samples' kernels also needed these fixes:

• 1-D tiles. A 1-D tile has only the partition dimension; NKI broadcasts it along the free axes. numpy aligned it with the last axis, so nisa.activation(op=nl.exp, data=t, scale=A[c0:c1, j]) failed with operands could not be broadcast together with shapes (128,512) (128,). Binary operations, scale / bias, and scalar operands now align a 1-D tile with axis 0.
• par_dim. Kernels import it directly (from neuronxcc.nki.language import par_dim), so nl.ndarray((par_dim(128), n), ...) received NKI's ParDim objects; shapes now take their value.
• nl.mgrid. Its result has .p and .x (and .y, .z) for its axes, as in NKI.
• Assignment. Assigning a [P, 1] result to a [P] slot reshapes it, as store already does.
• Slices past the tensor. NKI keeps the tile's size, but numpy truncated the slice and changed the tile's shape. A key of ints and slices with an out-of-range slice is now indexed with arrays, so the tile keeps its shape and the out-of-bounds lanes reach masked_load. An unmasked out-of-bounds load now raises IndexError there, as NKI's simulator does, instead of failing on mask.shape.
• Shared HBM aliasing. Builder.ndarray named a shared-HBM tensor after its caller's frame so that every program of the grid gets the same array. That caller was the op-patching wrapper, so every shared-HBM tensor of a kernel shared one name, and one array. It now uses the kernel's allocation site.
• Tracer. It accepts a boolean mask (a condition on Python ints) and slices whose bounds are computed (nl.ds(start, size)).

Tests

• New tests/nki/test_nki_isa.py: six kernels run under TileLens and under nki.simulate_kernel, and their outputs must match. They cover matmul with DMA, the vector-engine instructions, a grid with nl.ds, DMA with subscripts, miscellaneous operations, and partition vectors with mgrid and the scan. Two more tests check that a DMA is traced as a load and a store, and that an unmasked slice past the tensor raises IndexError, as it does in NKI. All eight fail on [FIX] Run NKI kernels written like the official samples in the NKI interpreter #539's head.
• On AWS's NKI samples, the interpreter's outputs now match nki.simulate_kernel for the eleven attention tutorial kernels (attn_fwd_v1 ... v11), the three fused Mamba kernels, and the Stable Diffusion attention tutorial, besides [FIX] Run NKI kernels written like the official samples in the NKI interpreter #539's kernels.
• pytest tests/ -m "" (NKI tests included): 626 passed, 2 skipped on [FIX] Run NKI kernels written like the official samples in the NKI interpreter #539; 634 passed, 2 skipped with this PR. Environment: Python 3.12, Triton 3.8.0, neuronx-cc 2.22.12471, nki 0.2.0.
• ruff check and ruff format --check (v0.1.9) pass on the changed files.

Not modeled: kernels that allocate SBUF by address (ncc.sbuf.mod_alloc) and rely on two tiles sharing the same memory. Such tiles are separate arrays here.

🤖 Generated with Claude Code

https://claude.ai/code/session_0169QRxKtsivkLDCS5f7rrsm