[mark14wu]

Summary

The legacy NKI interpreter (frontend="nki") runs only kernels written in a restricted style. Kernels in the style of AWS's NKI samples fail before their first tile operation. Of seven tutorial kernels from the samples (softmax, rmsnorm, two layernorms, transpose2d, average pool 2-D, SPMD tensor addition), only the SPMD addition ran under the tracer on main. The failures come from a few idioms:

Idiom	Failure on main
nl.store(o[i], value=nl.copy(nl.load(x[i]))): a load or store nested in another call	NKITransformer rewrote nl.load(x[...]) to nl.load(x, keys) only when the call was a whole statement, so the nested call lost its keys: TypeError: missing ... 'keys'
nl.load(t), nl.store(t, value=v): whole-tensor access	no keys at all, same error
nl.load(x[i * 128 + ix, iy], mask=(i * 128 + ix < n)): a [128, 1] row mask on a [128, F] tile	NKI broadcasts masks; masked_load asserts that the mask has the tile's shape
nl.load(x[...]) without a mask	the tracer read mask.data of None
out = nl.ndarray(..., buffer=nl.shared_hbm) ... return out	NKIInterpretedFunction.run dropped the kernel's return value, so the launch returned None
nl.max(t, axis=[1], keepdims=True), nl.sum(t, axis=[1])	numpy rejects a list of axes; and NKI keeps the last free axis: reducing the free axis of a [128, F] tile gives [128, 1], so t / nl.sum(t, axis=[1]) broadcasts in NKI and failed here
nl.max, nl.min, nl.mean, nl.add, nl.subtract, nl.divide, nl.maximum, nl.minimum, nl.static_cast, nl.par_dim, nl.sequential_range, nl.static_range	not patched, so the kernel called NKI's real implementation: RuntimeError: calling NKI API outside of NKI kernels is not supported

This PR:

• NKITransformer.visit_Call rewrites nl.load / nl.store / nl.load_transpose2d with a subscripted first argument wherever the call appears.
• Load and store keys default to the whole tensor (full_keys), in the interpreter and in the frontend's Load / Store adapters, so clients see the same keys.
• tile_shape / tile_mask (in core/masked_load_store.py) give the shape of ndarray[keys] even when some lanes are out of bounds, and broadcast an NKI mask to it. The interpreter broadcasts masks and stored values with them; the tracer records the broadcast mask, or an all-true mask for an access without one.
• The runner returns the last program's return value, converted to numpy arrays. Every program returns the same shared-HBM arrays. It also returns them when a client stops the grid early (second commit); before, the early return dropped them.
• Reductions go through Builder._reduce: a list of axes becomes a tuple, mask becomes numpy's where (with the dtype's lowest or highest value as the identity of max / min), dtype casts the result, and a result that would lose its last free axis gets [P, 1], as NKI's simulator gives. Element-wise operations take dtype; nl.load takes dtype too.
• The operations in the last row of the table are implemented and patched. par_dim returns the size, and the three ranges iterate like range.

The shapes and results were checked against NKI's own CPU simulator (neuronxcc.nki.simulate_kernel). For example, nl.sum over axis [1] of an [8, 4] tile gives (8, 1); over axes [1, 2] of an [8, 4, 3] tile, (8, 1) without keepdims and (8, 1, 1) with it. With this PR, six of the seven tutorial kernels run under the tracer. The seventh, layernorm v2, uses nisa.bn_stats, which the interpreter does not implement yet.

Tests

• New tests/nki/test_nki_kernel_idioms.py: four kernels in the samples' style run under TileLens and under nki.simulate_kernel; their outputs must match. The kernels cover a masked softmax over 200 rows, whole-tensor and nested loads and stores, an SPMD addition on a 2×2 grid, and a kernel with loops, casts, and binary operations. All four fail on main. A fifth test stops the grid after the first program with a client whose post_run_callback returns False and checks that the launch still returns that program's output; it fails without the second commit.
• tests/nki/test_nki_tracer.py::test_tracer_records_reduce_sum now expects the reduction's output shape (4, 1) instead of (4,), as NKI's simulator gives. Storing that result into the test's 1-D output slice still works: a stored value with the tile's number of elements is reshaped to the tile.
• pytest tests/ -m "" (NKI tests included): 621 passed, 2 skipped on main; 626 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.

🤖 Generated with Claude Code

[github-actions]

Performance Benchmark

Benchmark	main (min)	PR (min)	Change	Samples
gemm	0.054s	0.055s	+1.2%	20 / 20
gemm_oob	0.060s	0.061s	+0.8%	20 / 20
indirect_load	0.011s	0.011s	-1.8%	20 / 20
nested_loop	0.111s	0.111s	+0.7%	20 / 20
block_pointer_loop_advance	0.059s	0.060s	+1.6%	20 / 20
liger_jsd	0.076s	0.076s	-0.2%	20 / 20
flaggems_layernorm	0.190s	0.199s	+4.6%	20 / 20
swiglu	0.092s	0.092s	+0.3%	20 / 20
cross_entropy	0.517s	0.517s	+0.0%	20 / 20
fused_linear_jsd	0.117s	0.117s	-0.6%	20 / 20
Total	1.288s	1.299s	+0.8%	N/A

Iterations: 1 warmup + 20 measured
Samples are shown as main / PR; long pytest benchmarks may use fewer samples.