[mark14wu]

Stacked on #540. With this, the legacy NKI interpreter runs the AWS NKI samples' tutorials, reference kernels, and several Neuron-team kernels as the samples write them (@nki.jit kernels and their helper kernels), with outputs matching NKI's own CPU simulator (nki.simulate_kernel).

Operations (checked against nki.simulate_kernel)

• nl.where, nl.equal, nl.loop_reduce (dst[...] = op(dst[...], x)), nl.dropout / nl.random_seed (zeroes without rescaling, like NKI; the random numbers do not match the hardware's, as in NKI's simulator), nl.all_reduce over a grid of one program (several programs raise NotImplementedError, since programs run one at a time).
• nisa.bn_stats / bn_aggr (even/odd statistics, masked lanes excluded), tensor_partition_reduce, select_reduce (the Vector Engine accumulator carries across calls: reset_reduce, reduce), tensor_copy_predicated, and the HBM DMA instructions _tiled_offloaded_memcpy / _tiled_offloaded_fma (traced as loads and stores).
• nisa.nc_matmul with perf_mode="double_row_gen3" (operands [K, 2, M] and [K, 2, N], both rows contracted).
• nl.abs as a tensor_scalar operator; get_nc_version() follows NEURON_PLATFORM_TARGET_OVERRIDE=trn2 (gen3), as NKI's simulator does.

Kernels as the samples write them

• tilelens.trace(frontend="nki") accepts an @nki.jit / @nki.baremetal kernel. The NKI kernels it calls by name (@nki.jit(mode="trace") helpers, also imported from other modules) are rewritten too, with one rewritten namespace per module so kernels that call each other find the rewrites.
• a = t[idx]; nl.load(a, mask=m): subscripting an HBM tensor gives an access pattern that the load/store transfers with its mask (kernel arguments are HBM tensors). Before, the subscript read the masked-off out-of-range lanes eagerly and raised.
• t.reshape(shape=...), t.view(dtype), a [P] vector indexed as v[:, 0], a predicate subscript t[mask] (not numpy's boolean selection), tensors passed by keyword (they were not converted to arrays), [P] / [P, 1] operands of activation, and ncc.multi_buffer directives.
• ISA instructions with a destination or reduce_res write through to the tensor an indexed slice came from.

Fixes

• A client whose pre_run_callback declines a program (the profiler's block sampling) now skips that program instead of ending the grid, in the legacy and beta2 NKI runners (the Triton frontend already skips). Before, the profiler ran nothing unless the sampled program was the first.
• masked_load's fast path called np.iinfo on float tensors, so every masked float load took the slow path, which copies the whole tensor; masked_store built an index grid of the whole (padded) tensor for every masked store (690 MB for a (10, 2242, 1280) output). Both now have a fast path for in-bounds accesses, and the store's slow path indexes zero-stride views.
• The tracer no longer reads numpy masks and keys through .data (a memoryview).
• Builder.store's value annotation is NDArray | None (mypy).

Beta2 (import nki): the matrix-multiplication tutorial

• tilelens.trace(frontend="nki_beta2") accepts an @nki.jit kernel of the nki package too.
• The beta2 runner returns the kernel's outputs (as numpy arrays).
• nl.shared_hbm is an HBM buffer; nisa.memset; nl.sequential_range / nl.static_range iterate like affine_range; nisa.tensor_copy accepts the dtype that NKI's tracer accepts and the tutorial passes.
• With these, the tutorial's five kernels run under TileLens and match the tutorial's own check (NKI 0.2.0 has no CPU simulator). Test: an @nki.jit kernel written like the tutorial (checked by NKI's tracer first) returns its output.

Tensors placed by address, and the paged-attention kernel

• SBUF tensors placed with ncc.sbuf.mod_alloc / ncc.sbuf.alloc are views of a modeled physical SBUF at the addresses NKI's modulo allocation gives, so tensors and logical tiles at one address share values and a tile declared in a loop keeps its values across iterations (AddressedTiles when logical tiles outnumber physical ones). PSUM tensors start from zero (they were uninitialized numpy arrays, zero only when fresh pages backed them). The allocated SD attention, the allocated fused RMSNorm-QKV, and the attention tutorial's v10 now match NKI's simulator.
• oob_mode.skip for loads, stores, and DMA (out-of-bounds lanes are not moved; a tile assigned a skipping load keeps its values there; clients see the access as performed); scalar(t) (also imported by name) for dynamic loops; plain helpers that call NKI kernels through their module's globals reach the rewrites (kernels in user modules are swapped for their rewrites during a launch); an unrewritten nl.store(t[k], v); operands of fewer dimensions align from the partition dimension; a 0-d index tile slices like an int (numpy copied, losing writes); NDArray.itemsize. With these, flash_paged_attention_varlen runs and matches its test's torch reference (NKI's simulator cannot run it).
• The module scans these need (user modules whose kernels to swap; modules importing scalar by name) are cached until more modules are loaded: scanning every module cost about 9 ms per launch with torch loaded.

Tests

tests/nki/test_nki_more_ops.py (20 tests): selection, statistics, loop reduction, and HBM DMA kernels compared with nki.simulate_kernel; DMA traced as loads and stores; dropout; all_reduce across programs reported unsupported; a masked gather through an access pattern; a @nki.jit kernel calling a @nki.jit(mode="trace") helper; the tile idioms above; the trn2 target; a client skipping a program, a double-row matmul on the trn2 target, address placement (a double-buffered kernel vs the simulator; shared addresses; PSUM reuse), skip-mode loads and a plain helper vs the simulator, a dynamic trip count, left-aligned operands, and a 0-d index keeping a view. Full suite with -m "": 655 passed, 2 skipped.

🤖 Generated with Claude Code

https://claude.ai/code/session_0169QRxKtsivkLDCS5f7rrsm