From a550f3600bc191c347e002bc11009dd50d8e98fa Mon Sep 17 00:00:00 2001
From: Hao Wu <hwu27@gmu.edu>
Date: Sun, 8 Nov 2026 18:09:47 -0500
Subject: [PATCH 1/6] [FEAT] Interpret the remaining NKI operations of the
 official samples and run their kernels as written

Operations, checked against NKI's own CPU simulator (nki.simulate_kernel):
- nl.where, nl.equal, nl.loop_reduce (dst[...] = op(dst[...], x)), nl.dropout and
  nl.random_seed (zeroes without rescaling, like NKI), and nl.all_reduce over a grid of one
  program (several programs raise NotImplementedError: programs run one at a time).
- nisa.bn_stats / bn_aggr (even/odd statistics, masked lanes excluded),
  tensor_partition_reduce, select_reduce (with the Vector Engine accumulator carried across
  calls), tensor_copy_predicated, and the HBM DMA instructions _tiled_offloaded_memcpy /
  _tiled_offloaded_fma (traced as loads and stores).
- 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 (or @nki.baremetal) kernel, and the
  NKI kernels it calls by name (@nki.jit helpers, also imported from other modules) are
  rewritten too, one rewritten namespace per module.
- `a = t[idx]; nl.load(a, mask=m)`: subscripting an HBM tensor gives an access pattern that
  the load or store transfers with its mask (kernel arguments are HBM tensors).
- `t.reshape(shape=...)`, `t.view(dtype)`, a [P] vector indexed as `v[:, 0]`, a predicate
  subscript `t[mask]`, tensors passed by keyword, [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 (legacy and beta2 NKI runners).
- masked_load's fast path no longer fails on float tensors (np.iinfo), and masked_store no
  longer builds an index grid of the whole tensor; both have a fast path for in-bounds
  accesses.
- The tracer no longer reads numpy masks and keys through `.data` (a memoryview).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0169QRxKtsivkLDCS5f7rrsm
---
 tests/nki/test_nki_more_ops.py          | 334 ++++++++++++++++
 tilelens/clients/tracer/tracer.py       |  16 +-
 tilelens/core/frontend/nki.py           |  22 +-
 tilelens/core/frontend/nki_transform.py |  30 +-
 tilelens/core/masked_load_store.py      |  81 +++-
 tilelens/core/simulation/nki.py         | 498 +++++++++++++++++++++---
 tilelens/core/simulation/nki_beta2.py   |   2 +-
 tilelens/core/trace.py                  |   6 +
 8 files changed, 897 insertions(+), 92 deletions(-)
 create mode 100644 tests/nki/test_nki_more_ops.py

diff --git a/tests/nki/test_nki_more_ops.py b/tests/nki/test_nki_more_ops.py
new file mode 100644
index 00000000..35b77b43
--- /dev/null
+++ b/tests/nki/test_nki_more_ops.py
@@ -0,0 +1,334 @@
+"""Selection, statistics, loop and partition reductions, predicated copies, HBM DMA, and
+dropout in the NKI interpreter, checked against NKI's own CPU simulator
+(``nki.simulate_kernel``) on small kernels."""
+
+import numpy as np
+import pytest
+
+try:
+    from neuronxcc import nki
+    import neuronxcc.nki.isa as nisa
+    import neuronxcc.nki.language as nl
+except ModuleNotFoundError:
+    pytest.skip(
+        "NeuronX dependencies are missing. Install tilelens[nki] to run these tests.",
+        allow_module_level=True,
+    )
+
+import tilelens
+from tilelens.clients import Tracer
+
+pytestmark = pytest.mark.nki  # only run at "pytest -m nki"
+
+
+def selection_kernel(x, y, pred):
+    # nl.where, nl.equal, nisa.select_reduce with the accumulator carried from one call to
+    # the next, and nisa.tensor_copy_predicated.
+    sel = nl.ndarray(x.shape, dtype=x.dtype, buffer=nl.shared_hbm)
+    eq = nl.ndarray(x.shape, dtype=nl.float32, buffer=nl.shared_hbm)
+    picked = nl.ndarray(x.shape, dtype=x.dtype, buffer=nl.shared_hbm)
+    maxes = nl.ndarray((x.shape[0], 2), dtype=x.dtype, buffer=nl.shared_hbm)
+    copied = nl.ndarray(x.shape, dtype=x.dtype, buffer=nl.shared_hbm)
+    a = nl.load(x)
+    b = nl.load(y)
+    p = nl.load(pred)
+    nl.store(sel, value=nl.where(a > b, a, -5.0))
+    nl.store(eq, value=nl.equal(a, b))
+    dst = nl.ndarray(x.shape, dtype=x.dtype)
+    res = nl.ndarray((x.shape[0], 2), dtype=x.dtype)
+    i_p = nl.arange(x.shape[0])[:, None]
+    nisa.select_reduce(
+        dst=dst,
+        predicate=p,
+        on_true=a,
+        on_false=-9.0,
+        reduce_cmd=nisa.reduce_cmd.reset_reduce,
+        reduce_res=res[:, 0],
+        reduce_op=np.max,
+    )
+    nisa.select_reduce(
+        dst=dst,
+        predicate=p,
+        on_true=b,
+        on_false=-9.0,
+        reverse_pred=True,
+        reduce_cmd=nisa.reduce_cmd.reduce,
+        reduce_res=res[i_p, 1],  # an indexed slice: a copy in numpy, written through
+        reduce_op=np.max,
+    )
+    nl.store(picked, value=dst)
+    nl.store(maxes, value=res)
+    c = nl.full(x.shape, 7.0, dtype=x.dtype)
+    nisa.tensor_copy_predicated(src=a, dst=c, predicate=p)
+    nl.store(copied, value=c)
+    return sel, eq, picked, maxes, copied
+
+
+def stats_kernel(x):
+    # nisa.bn_stats over two chunks (the second masked), nisa.bn_aggr, and
+    # nisa.tensor_partition_reduce.
+    stats = nl.ndarray((x.shape[0], 12), dtype=nl.float32, buffer=nl.shared_hbm)
+    mean_var = nl.ndarray((x.shape[0], 2), dtype=nl.float32, buffer=nl.shared_hbm)
+    col_sum = nl.ndarray((1, x.shape[1]), dtype=x.dtype, buffer=nl.shared_hbm)
+    col_max = nl.ndarray((1, x.shape[1]), dtype=x.dtype, buffer=nl.shared_hbm)
+    a = nl.load(x)
+    s = nl.ndarray((x.shape[0], 12), dtype=nl.float32)
+    i_p = nl.arange(x.shape[0])[:, None]
+    i_s = nl.arange(6)[None, :]
+    i_f = nl.arange(64)[None, :]
+    s[i_p, i_s] = nisa.bn_stats(a[i_p, i_f], dtype=nl.float32)
+    s[i_p, 6 + i_s] = nisa.bn_stats(a[i_p, 64 + i_f], mask=(i_f < 37), dtype=nl.float32)
+    nl.store(stats, value=s)
+    nl.store(mean_var, value=nisa.bn_aggr(s))
+    nl.store(col_sum, value=nisa.tensor_partition_reduce(np.add, a))
+    nl.store(col_max, value=nisa.tensor_partition_reduce(np.max, a))
+    return stats, mean_var, col_sum, col_max
+
+
+def loop_reduce_kernel(x):
+    # nl.loop_reduce into a subscripted destination: dst = op(dst, x), also into one slice
+    # of a larger tile per outer iteration.
+    total = nl.ndarray((128, x.shape[2]), dtype=nl.float32, buffer=nl.shared_hbm)
+    per_j = nl.ndarray((2, 128, x.shape[2]), dtype=nl.float32, buffer=nl.shared_hbm)
+    acc = nl.full((128, x.shape[2]), 1.0, dtype=nl.float32, buffer=nl.sbuf)
+    for i in nl.affine_range(x.shape[0]):
+        acc[:, :] = nl.loop_reduce(nl.load(x[i]), op=np.add, loop_indices=[i])
+    nl.store(total, value=acc)
+    for j in nl.affine_range(2):
+        tile = nl.zeros((128, 2 * x.shape[2]), dtype=nl.float32, buffer=nl.sbuf)
+        for i in nl.affine_range(x.shape[0]):
+            tile[:, j * x.shape[2] : (j + 1) * x.shape[2]] = nl.loop_reduce(
+                nl.load(x[i]), op=np.add, loop_indices=[i], dtype=nl.float32
+            )
+        nl.store(per_j[j], value=tile[:, j * x.shape[2] : (j + 1) * x.shape[2]])
+    return total, per_j
+
+
+def hbm_dma_kernel(x):
+    # nisa._tiled_offloaded_memcpy and nisa._tiled_offloaded_fma between HBM slices.
+    copy = nl.ndarray(x.shape, dtype=x.dtype, buffer=nl.shared_hbm)
+    fma = nl.ndarray((x.shape[0], x.shape[1] // 2), dtype=x.dtype, buffer=nl.shared_hbm)
+    h = x.shape[1] // 2
+    nisa._tiled_offloaded_memcpy(dst=copy[:, 0:h], src=x[:, 0:h])
+    nisa._tiled_offloaded_memcpy(dst=copy[:, h : 2 * h], src=x[:, h : 2 * h])
+    nisa._tiled_offloaded_fma(x[:, 0:h], x[:, h : 2 * h], dst=fma, scales=[1.0, 2.0])
+    return copy, fma
+
+
+rng = np.random.default_rng(0)
+_x = rng.standard_normal((128, 64)).astype(np.float32)
+_y = np.where(rng.random((128, 64)) < 0.5, _x, rng.standard_normal((128, 64))).astype(
+    np.float32
+)
+
+
+@pytest.mark.parametrize(
+    "kernel, args",
+    [
+        (selection_kernel, (_x, _y, (rng.random((128, 64)) < 0.5).astype(np.uint8))),
+        (stats_kernel, (rng.standard_normal((128, 128)).astype(np.float32),)),
+        (loop_reduce_kernel, (rng.standard_normal((4, 128, 64)).astype(np.float32),)),
+        (hbm_dma_kernel, (rng.standard_normal((128, 64)).astype(np.float32),)),
+    ],
+    ids=["selection", "stats", "loop_reduce", "hbm_dma"],
+)
+def test_kernels_match_nki_simulator(kernel, args):
+    tilelens.clear()
+    got = tilelens.trace(client=Tracer(), frontend="nki")(kernel)[(1,)](*args)
+    want = nki.simulate_kernel(nki.jit(kernel)[(1,)], *args)
+    got = got if isinstance(got, tuple) else (got,)
+    want = want if isinstance(want, tuple) else (want,)
+    assert len(got) == len(want)
+    for g, w in zip(got, want):
+        np.testing.assert_allclose(
+            g.astype(np.float32), w.astype(np.float32), rtol=1e-4, atol=1e-4
+        )
+
+
+def test_hbm_dma_is_traced_as_loads_and_stores():
+    tilelens.clear()
+    from tilelens.core.data import Load, Store
+    from tilelens.core.trace import launches
+
+    x = rng.standard_normal((128, 64)).astype(np.float32)
+    tilelens.trace(client=Tracer(), frontend="nki")(hbm_dma_kernel)[(1,)](x)
+    records = launches[-1].records
+    # Two memcpys (a load and a store each) and an fma of two sources (two loads, a store).
+    assert len([r for r in records if isinstance(r, Load)]) == 4
+    assert len([r for r in records if isinstance(r, Store)]) == 3
+
+
+def dropout_kernel(x, rate):
+    out = nl.ndarray(x.shape, dtype=x.dtype, buffer=nl.shared_hbm)
+    nl.random_seed(seed=nl.full((1, 1), 3, dtype=nl.int32))
+    nl.store(out, value=nl.dropout(nl.load(x), rate=rate))
+    return out
+
+
+def test_dropout_zeroes_without_rescaling():
+    # NKI's random numbers differ between its simulator and the hardware; only the
+    # distribution is checked.
+    x = rng.random((128, 512)).astype(np.float32) + 1.0
+    tilelens.clear()
+    kept = tilelens.trace(client=Tracer(), frontend="nki")(dropout_kernel)[(1,)](x, 0.0)
+    np.testing.assert_array_equal(kept, x)
+    tilelens.clear()
+    half = tilelens.trace(client=Tracer(), frontend="nki")(dropout_kernel)[(1,)](x, 0.5)
+    assert 0.45 < (half == 0).mean() < 0.55
+    np.testing.assert_array_equal(half[half != 0], x[half != 0])
+    tilelens.clear()
+    again = tilelens.trace(client=Tracer(), frontend="nki")(dropout_kernel)[(1,)](
+        x, 0.5
+    )
+    np.testing.assert_array_equal(again, half)  # same seed, same zeros
+
+
+def all_reduce_kernel(x):
+    out = nl.ndarray(x.shape, dtype=x.dtype, buffer=nl.shared_hbm)
+    shard_id = nl.program_id(0)
+    nl.store(out, value=nl.all_reduce(nl.load(x), op=np.add, program_axes=[shard_id]))
+    return out
+
+
+def test_all_reduce_across_programs_is_reported_unsupported():
+    tilelens.clear()
+    x = rng.standard_normal((128, 64)).astype(np.float32)
+    with pytest.raises(NotImplementedError, match="all_reduce"):
+        tilelens.trace(client=Tracer(), frontend="nki")(all_reduce_kernel)[(2,)](x)
+
+
+def gather_kernel(x):
+    # `a = t[idx, j]` then `nl.load(a, mask=m)`: an HBM access pattern with computed indices,
+    # out of range on lanes the mask turns off.
+    out = nl.ndarray((256, x.shape[1]), dtype=x.dtype, buffer=nl.shared_hbm)
+    for t in nl.affine_range(2):
+        i_p = nl.arange(128)[:, None]
+        i_f = nl.arange(x.shape[1])[None, :]
+        pos = t * 128 + i_p
+        src = nisa.iota(pos, dtype=nl.int32)
+        valid = pos < x.shape[0]
+        rows = x[src, i_f]
+        nl.store(out[pos, i_f], value=nl.load(rows, mask=valid), mask=valid)
+    return out
+
+
+def test_masked_gather_through_an_access_pattern():
+    x = rng.standard_normal((200, 64)).astype(np.float32)
+    tilelens.clear()
+    got = tilelens.trace(client=Tracer(), frontend="nki")(gather_kernel)[(1,)](x)
+    want = nki.simulate_kernel(nki.jit(gather_kernel)[(1,)], x)
+    np.testing.assert_array_equal(got[:200], want[:200])  # rows 200.. are never written
+
+
+@nki.jit(mode="trace")
+def _scaled_rows(t_hbm, rows, scale):
+    # a helper kernel, called by name: TileLens rewrites it as it rewrites the kernel
+    return nl.multiply(nl.load(t_hbm[rows, :]), scale)
+
+
+@nki.jit
+def kernel_with_helper(x):
+    out = nl.ndarray(x.shape, dtype=x.dtype, buffer=nl.shared_hbm)
+    for t in nl.affine_range(x.shape[0] // 128):
+        rows = nl.ds(t * 128, 128)
+        nl.store(out[rows, :], value=_scaled_rows(x, rows, 3.0))
+    return out
+
+
+def test_jit_kernel_and_its_helper_kernels_run_as_written():
+    x = rng.standard_normal((256, 32)).astype(np.float32)
+    tilelens.clear()
+    got = tilelens.trace(client=Tracer(), frontend="nki")(kernel_with_helper)[(1,)](x)
+    want = nki.simulate_kernel(kernel_with_helper[(1,)], x)
+    np.testing.assert_allclose(got, want, rtol=1e-6)
+
+
+import neuronxcc.nki.compiler as ncc  # noqa: E402
+
+
+def tile_idioms_kernel(x, *, scale):
+    # a tensor passed by keyword; `t.reshape(shape=...)` of an HBM tensor; a [P] vector
+    # indexed as `v[:, 0]`; an activation of a [P] vector with a [P, 1] bias; a predicate
+    # subscript `t[mask]`; a range directive; and a bit view of a tile.
+    out = nl.ndarray(x.shape, dtype=x.dtype, buffer=nl.shared_hbm)
+    col = nl.ndarray((x.shape[0], 1), dtype=x.dtype, buffer=nl.shared_hbm)
+    bits = nl.ndarray((x.shape[0], 1), dtype=nl.int32, buffer=nl.shared_hbm)
+    flat = x.reshape(shape=(x.shape[0], x.shape[1] * x.shape[2]))
+    t = nl.load(flat)
+    s = nl.load(scale)  # [P, 1]
+    rows = nl.ndarray((x.shape[0], 2), dtype=nl.float32)
+    v = rows[:, 0]  # a [P] vector
+    v[:, 0] = nisa.tensor_reduce(np.add, t, axis=[1])
+    e = nisa.activation(op=nl.copy, data=v, bias=s, scale=2.0)
+    i_f = nl.arange(t.shape[1])[None, :]
+    keep = (
+        i_f < t.shape[1]
+    )  # every lane on: the predicate only must not change the tile
+    for k in nl.sequential_range(2, directives=ncc.multi_buffer(2)):
+        nl.store(out.reshape(shape=flat.shape), value=nl.multiply(t[keep], 2.0))
+    nl.store(col, value=e)
+    nl.store(bits, value=nl.copy(s).view(nl.int32))
+    return out, col, bits
+
+
+def test_tile_idioms_match_nki_simulator():
+    x = rng.standard_normal((128, 4, 8)).astype(np.float32)
+    scale = rng.standard_normal((128, 1)).astype(np.float32)
+    tilelens.clear()
+    got = tilelens.trace(client=Tracer(), frontend="nki")(tile_idioms_kernel)[(1,)](
+        x, scale=scale
+    )
+    want = nki.simulate_kernel(nki.jit(tile_idioms_kernel)[(1,)], x, scale=scale)
+    for g, w in zip(got, want):
+        np.testing.assert_allclose(g, w, rtol=1e-5, atol=1e-5)
+
+
+def abs_kernel(x):
+    out = nl.ndarray(x.shape, dtype=x.dtype, buffer=nl.shared_hbm)
+    amax = nl.ndarray((x.shape[0], 1), dtype=x.dtype, buffer=nl.shared_hbm)
+    acc = nl.ndarray((x.shape[0], 1), dtype=x.dtype)
+    t = nisa.tensor_scalar_reduce(
+        data=nl.load(x), op0=nl.abs, operand0=0.0, reduce_op=nl.max, reduce_res=acc
+    )
+    nl.store(out, value=t)
+    nl.store(amax, value=acc)
+    return out, amax
+
+
+def test_trn2_target_abs_operator_and_nc_version(monkeypatch):
+    from tilelens.core.simulation.nki import nki_builder
+
+    assert nki_builder.get_nc_version() == nisa.nc_version.gen2
+    monkeypatch.setenv("NEURON_PLATFORM_TARGET_OVERRIDE", "trn2")
+    assert nki_builder.get_nc_version() == nisa.nc_version.gen3
+    x = rng.standard_normal((128, 64)).astype(np.float32)
+    tilelens.clear()
+    got = tilelens.trace(client=Tracer(), frontend="nki")(abs_kernel)[(1,)](x)
+    want = nki.simulate_kernel(nki.jit(abs_kernel)[(1,)], x)
+    for g, w in zip(got, want):
+        np.testing.assert_allclose(g, w)
+
+
+def rows_kernel(x):
+    out = nl.ndarray(x.shape, dtype=x.dtype, buffer=nl.shared_hbm)
+    rows = nl.ds(nl.program_id(0) * 128, 128)
+    nl.store(out[rows, :], value=nl.load(x[rows, :]))
+    return out
+
+
+def test_a_program_a_client_skips_does_not_stop_the_grid():
+    # The profiler samples programs: a skipped program must not end the grid.
+    class SkipFirstProgram(Tracer):
+        def grid_idx_callback(self, grid_idx):
+            super().grid_idx_callback(grid_idx)
+            self.current = grid_idx
+
+        def pre_run_callback(self, fn):
+            return self.current != (0, 0, 0)
+
+    x = rng.standard_normal((384, 16)).astype(np.float32)
+    tilelens.clear()
+    got = tilelens.trace(client=SkipFirstProgram(), frontend="nki")(rows_kernel)[(3,)](
+        x
+    )
+    np.testing.assert_array_equal(got[128:], x[128:])
diff --git a/tilelens/clients/tracer/tracer.py b/tilelens/clients/tracer/tracer.py
index 509334b4..8380f3d7 100644
--- a/tilelens/clients/tracer/tracer.py
+++ b/tilelens/clients/tracer/tracer.py
@@ -17,13 +17,21 @@
 import numpy as np
 
 
+def _nki_value(x):
+    """The numpy value of an NKI interpreter array; numpy arrays and scalars as they are
+    (their own `.data` is a memoryview)."""
+    if isinstance(x, (np.ndarray, np.generic)):
+        return x
+    return getattr(x, "data", x)
+
+
 def _nki_offsets_and_mask(tensor, keys, mask):
     """Element offsets and mask of an NKI access ``tensor[keys]``. NKI masks broadcast
     to the tile shape, and an access without a mask touches every lane."""
     offsets = tensor.get_offsets().data
     mask_data = tile_mask(
-        offsets, keys, getattr(mask, "data", mask)
-    )  # NDArray, bool, or None
+        offsets, keys, _nki_value(mask)
+    )  # NDArray, numpy array, bool, or None
     if mask_data is None:
         mask_data = np.ones(tile_shape(offsets, keys), dtype=bool)
     return masked_load(offsets, keys, mask=mask_data), mask_data
@@ -108,10 +116,10 @@ def _convert_keys_to_numpy(keys):
                 part = (
                     lambda v: v
                     if v is None or isinstance(v, int)
-                    else int(np.asarray(getattr(v, "data", v)))
+                    else int(np.asarray(_nki_value(v)))
                 )  # noqa: E731
                 return slice(part(keys.start), part(keys.stop), part(keys.step))
-            return keys.data if hasattr(keys, "data") else keys
+            return _nki_value(keys)
 
         @self.lock_fn
         def pre_load_callback(ptr, mask, keys):
diff --git a/tilelens/core/frontend/nki.py b/tilelens/core/frontend/nki.py
index 42e51707..50dd70eb 100644
--- a/tilelens/core/frontend/nki.py
+++ b/tilelens/core/frontend/nki.py
@@ -19,7 +19,7 @@
 HAS_NKI = False
 nki_builder = None
 try:
-    from tilelens.core.simulation.nki import NDArray, full_keys, nki_builder  # type: ignore
+    from tilelens.core.simulation.nki import NDArray, access, nki_builder  # type: ignore
 
     HAS_NKI = True
 except ModuleNotFoundError:
@@ -33,6 +33,13 @@ def _nki_dot_adapter(x: Any, y: Any, *_args: Any, **_kwargs: Any) -> AdapterResu
     return AdapterResult(x, y)
 
 
+def _nki_access_adapter(tensor: Any, keys: Any, mask: Any) -> AdapterResult:
+    # `nl.load(t)` and `nl.store(t, value=v)` access the whole tensor, and
+    # `a = t[k]; nl.load(a)` accesses t[k].
+    tensor, keys = access(tensor, keys)
+    return AdapterResult(tensor, mask, keys)
+
+
 def _nki_reduce_sum_adapter(
     input_tensor: Any, *args: Any, mask: Any = None, **kwargs: Any
 ) -> AdapterResult:
@@ -62,22 +69,15 @@ def _nki_reduce_sum_adapter(
     NKI_ADAPTERS = {
         ProgramId: lambda axis, *_args, **_kwargs: AdapterResult(axis),
         Allocate: lambda *_args, **_kwargs: AdapterResult(),
-        # `nl.load(t)` and `nl.store(t, value=v)` access the whole tensor.
-        Load: lambda src, keys=None, *, mask=None, **_kwargs: AdapterResult(
-            src,
-            mask,
-            full_keys(src) if keys is None else keys,
+        Load: lambda src, keys=None, *, mask=None, **_kwargs: _nki_access_adapter(
+            src, keys, mask
         ),
         Store: lambda dst,
         keys=None,
         value=None,
         *,
         mask=None,
-        **_kwargs: AdapterResult(
-            dst,
-            mask,
-            full_keys(dst) if keys is None else keys,
-        ),
+        **_kwargs: _nki_access_adapter(dst, keys, mask),
         Dot: _nki_dot_adapter,
         ReduceSum: _nki_reduce_sum_adapter,
     }
diff --git a/tilelens/core/frontend/nki_transform.py b/tilelens/core/frontend/nki_transform.py
index 8b2ed703..0c5565be 100644
--- a/tilelens/core/frontend/nki_transform.py
+++ b/tilelens/core/frontend/nki_transform.py
@@ -10,8 +10,10 @@ class NKITransformer(ast.NodeTransformer):
     - `nl.<func>(x[...], ...)` -> `nl.<func>(x, slice_obj, ...)`
        where <func> is one of {load, store, load_transpose2d}, wherever the call
        appears (e.g. nested in another call, as in `nl.store(o[i], nl.load(x[i]))`)
-    - `nisa.dma_copy(dst=o[...], src=x[...])` and `nisa.dma_transpose(x[...])` -> keys passed
-       as `dst_keys` / `src_keys`
+    - `nisa.dma_copy(dst=o[...], src=x[...])`, `nisa.dma_transpose(x[...])`, and the
+       HBM-to-HBM DMA calls `nisa._tiled_offloaded_memcpy(dst=o[...], src=x[...])` and
+       `nisa._tiled_offloaded_fma(x[...], y[...], dst=o[...], ...)` -> keys passed as
+       `dst_keys` / `src_keys` (a list, one per source, for `_tiled_offloaded_fma`)
     - strips out kernel decorators
     """
 
@@ -26,18 +28,40 @@ def visit_Call(self, node: ast.Call) -> ast.AST:
             new_call = self._rewrite_dma_copy_call(node)
         return node if new_call is None else new_call
 
+    _DMA_CALLS = (
+        "dma_copy",
+        "dma_transpose",
+        "_tiled_offloaded_memcpy",
+        "_tiled_offloaded_fma",
+    )
+
     def _rewrite_dma_copy_call(self, call_node: ast.Call) -> ast.Call | None:
         """`nisa.dma_copy(dst=o[...], src=x[...])` -> `nisa.dma_copy(dst=o, dst_keys=..., src=x, src_keys=...)`"""
         if not (
             isinstance(call_node.func, ast.Attribute)
             and isinstance(call_node.func.value, ast.Name)
             and call_node.func.value.id == "nisa"
-            and call_node.func.attr in ("dma_copy", "dma_transpose")
+            and call_node.func.attr in self._DMA_CALLS
         ):
             return None
         args = list(call_node.args)
         keywords = []
         changed = False
+        if call_node.func.attr == "_tiled_offloaded_fma" and any(
+            isinstance(a, ast.Subscript) for a in args
+        ):
+            # positional sources -> `src_keys=[keys of each source, or None]`
+            keys = [
+                cast(ast.expr, self._create_slice_value_node(a.slice))
+                if isinstance(a, ast.Subscript)
+                else ast.Constant(value=None)
+                for a in args
+            ]
+            args = [a.value if isinstance(a, ast.Subscript) else a for a in args]
+            keywords.append(
+                ast.keyword(arg="src_keys", value=ast.List(elts=keys, ctx=ast.Load()))
+            )
+            changed = True
         if (
             call_node.func.attr == "dma_transpose"
             and args
diff --git a/tilelens/core/masked_load_store.py b/tilelens/core/masked_load_store.py
index c73bc3dd..5d089721 100644
--- a/tilelens/core/masked_load_store.py
+++ b/tilelens/core/masked_load_store.py
@@ -67,6 +67,31 @@ def _get_valid_indices(
     return valid_mask
 
 
+def _has_negative_index(keys: tuple) -> bool:
+    """Whether an integer or integer-array key is negative: numpy would wrap it around,
+    where it is an out-of-bounds index."""
+    for key in keys:
+        if isinstance(key, (int, np.integer)) and key < 0:
+            return True
+        if (
+            isinstance(key, np.ndarray)
+            and key.dtype.kind == "i"
+            and key.size
+            and key.min() < 0
+        ):
+            return True
+    return False
+
+
+def _fill_value(dtype):
+    """Value of lanes a masked load does not read."""
+    if np.issubdtype(dtype, np.integer):
+        return np.iinfo(dtype).max
+    if np.issubdtype(dtype, np.floating):
+        return np.finfo(dtype).max
+    return 0
+
+
 def tile_shape(ndarray: np.ndarray, keys) -> tuple:
     """Shape of ``ndarray[keys]``, also when some integer-array keys are out of bounds."""
     if not isinstance(keys, tuple):
@@ -80,7 +105,10 @@ def tile_shape(ndarray: np.ndarray, keys) -> tuple:
         ):
             key = np.clip(key, 0, max(ndarray.shape[dim] - 1, 0))
         clipped.append(key)
-    return ndarray[tuple(clipped)].shape
+    # index a one-byte, zero-stride view: only the shape is needed, not the values
+    return np.broadcast_to(np.zeros((), dtype=bool), ndarray.shape)[
+        tuple(clipped)
+    ].shape
 
 
 def tile_mask(ndarray: np.ndarray, keys, mask) -> np.ndarray | None:
@@ -107,19 +135,20 @@ def masked_load(
     Returns:
         Indexed array with masked error handling
     """
-    try:  # fast path case - if keys aren't OOB, just go with that
-        out = ndarray[keys].copy()
-        if mask is None:
-            return out
-        out[~mask] = np.iinfo(ndarray.dtype).max
-        return out
-    except Exception:
-        pass
-
     # Convert keys to tuple if it's not already
     if not isinstance(keys, tuple):
         keys = (keys,)
 
+    if not _has_negative_index(keys):
+        try:  # fast path case - if keys aren't OOB, just go with that
+            out = ndarray[keys].copy()
+            if mask is None:
+                return out
+            out[~mask] = _fill_value(ndarray.dtype)
+            return out
+        except Exception:
+            pass
+
     # normalize (remove Nones in slices)
     keys, singleton_dims = normalize_slice(ndarray, keys)
     target_shape = _calculate_target_shape(keys, ndarray.shape)
@@ -152,13 +181,7 @@ def masked_load(
         )
 
     valid_mask = mask & in_bounds_mask
-    # Use appropriate max value based on dtype
-    if np.issubdtype(ndarray.dtype, np.integer):
-        result[~valid_mask] = np.iinfo(ndarray.dtype).max
-    elif np.issubdtype(ndarray.dtype, np.floating):
-        result[~valid_mask] = np.finfo(ndarray.dtype).max
-    else:
-        result[~valid_mask] = 0  # false
+    result[~valid_mask] = _fill_value(ndarray.dtype)
 
     return np.expand_dims(result, singleton_dims)
 
@@ -194,16 +217,32 @@ def masked_store(
     if not isinstance(keys, tuple):
         keys = (keys,)
 
+    if all(
+        k is None or isinstance(k, (int, np.integer, slice)) for k in keys
+    ) and not _has_negative_index(keys):
+        # fast path: basic indexing within bounds gives a view to store into
+        try:
+            target = ndarray[keys]
+        except IndexError:
+            target = None
+        if target is not None and target.shape == mask.shape:
+            target[mask] = value[mask]
+            return
+
     flat_mask = mask.ravel()  # can only index with bool tensors if they're 1d
 
     # normalize (remove Nones in slices)
     keys, singleton_dims = normalize_slice(ndarray, keys)
     target_shape = _calculate_target_shape(keys, ndarray.shape)
-    mgrid = np.mgrid[tuple(slice(0, dim, 1) for dim in target_shape)]
     idxs = []
-    for mgrid_dim in mgrid:
-        # get the section that would be sliced if the array was big enough
-        sliced_section = mgrid_dim[keys]
+    for dim, size in enumerate(target_shape):
+        # get the section that would be sliced if the array was big enough. The index grid
+        # of each dimension is a broadcast (zero-stride) view, so only the section is
+        # materialized, not a grid of the whole (padded) array.
+        axis_index = np.arange(size).reshape(
+            [-1 if d == dim else 1 for d in range(len(target_shape))]
+        )
+        sliced_section = np.broadcast_to(axis_index, target_shape)[keys]
 
         # get just the section where mask=True
         dim_idxs = sliced_section.ravel()[flat_mask]
diff --git a/tilelens/core/simulation/nki.py b/tilelens/core/simulation/nki.py
index 30b35376..96eb838e 100644
--- a/tilelens/core/simulation/nki.py
+++ b/tilelens/core/simulation/nki.py
@@ -4,6 +4,7 @@
 try:
     import neuronxcc.nki.language as nl
     import neuronxcc.nki.isa as nisa
+    import neuronxcc.nki.compiler as ncc
 except (
     ModuleNotFoundError
 ) as exc:  # pragma: no cover - only hit when optional deps missing
@@ -24,10 +25,26 @@
 
 
 def _value(x):
-    """The numpy value of an NDArray operand; other operands (scalars) as is."""
+    """The numpy value of an NDArray operand; other operands (scalars, numpy arrays) as is."""
     return x.data if isinstance(x, NDArray) else x
 
 
+def _is_hbm(buffer):
+    return getattr(buffer, "name", None) in ("hbm", "shared_hbm", "private_hbm")
+
+
+def _tile_keys(data, keys):
+    """NKI treats a 1-D tile (a partition vector) as [P, 1]: `t[:, 0]` indexes it too
+    (`t[:, None]` still adds an axis)."""
+    if (
+        np.ndim(data) == 1
+        and len(keys) == 2
+        and isinstance(keys[1], (int, np.integer, slice))
+    ):
+        return keys[:1]
+    return keys
+
+
 def _axis(axis):
     """NKI accepts a list of axes; numpy needs an int or a tuple."""
     return tuple(axis) if isinstance(axis, list) else axis
@@ -142,6 +159,37 @@ def _cast(value, dtype):
     return np.asarray(value) if dtype is None else np.asarray(value).astype(dtype)
 
 
+def _bool_tile(pred, shape, reverse=False):
+    """A predicate or mask operand as a boolean array of `shape` (a 1-D predicate is a
+    partition vector)."""
+    value = np.asarray(_partition_aligned(pred, len(shape)), dtype=bool)
+    value = np.broadcast_to(value, shape)
+    return ~value if reverse else value
+
+
+def _per_partition(value):
+    """A [P, ...] array as [P, N]: ISA instructions match operands by partition and
+    elements per partition, not by the shape of the free dimensions."""
+    value = np.asarray(value)
+    return value.reshape(value.shape[0], -1) if value.ndim else value.reshape(1, 1)
+
+
+def _write_into(dst, value):
+    """Write `value` into the tile `dst` in place, as ISA instructions with a destination
+    operand do, through to the tensor `dst` was sliced from when the slice is a copy
+    (numpy copies on indexing with arrays)."""
+    target = dst.data
+    value = np.asarray(value)
+    if value.shape != target.shape and value.size == target.size:
+        value = value.reshape(target.shape)
+    target[...] = np.broadcast_to(value, target.shape)
+    while getattr(dst, "_parent", None) is not None:
+        base, keys = dst._parent
+        if not np.may_share_memory(base.data, dst.data):
+            base.data[keys] = dst.data.reshape(base.data[keys].shape)
+        dst = base
+
+
 class NDArray:
     def __init__(self, buffer=None, name="", **kwargs):
         self.buffer = buffer
@@ -209,23 +257,51 @@ def __getitem__(self, keys):
         if not isinstance(keys, tuple):
             keys = (keys,)
 
+        if (
+            len(keys) == 1
+            and np.asarray(_value(keys[0])).dtype == bool
+            and np.ndim(_value(keys[0])) == self.data.ndim
+        ):
+            # `t[mask]` attaches a predicate to the tile in NKI (it is not numpy's boolean
+            # selection); the interpreter does not model predicates of computations
+            return self
+        if _is_hbm(self.buffer) and not isinstance(self, HBMAccess):
+            # an access pattern, read by the load it is passed to (with that load's mask)
+            return HBMAccess(self, tuple(_value(k) for k in keys))
+
         # Apply the slicing to the underlying numpy array
-        new_keys = [k.data if isinstance(k, NDArray) else k for k in keys]
+        new_keys = _tile_keys(
+            self.data, [k.data if isinstance(k, NDArray) else k for k in keys]
+        )
         sliced_value = self.data[tuple(new_keys)]
 
         # Create a new NDArray with the sliced data
-        return NDArray(value=sliced_value, name=f"{self.name}_slice")
+        sliced = NDArray(value=sliced_value, name=f"{self.name}_slice")
+        sliced._parent = (
+            self,
+            tuple(new_keys),
+        )  # for ISA instructions that write into it
+        return sliced
 
     def __setitem__(self, keys, value):
         if not isinstance(keys, tuple):
             keys = (keys,)
 
         # Apply the slicing to the underlying numpy array
-        new_keys = tuple(k.data if isinstance(k, NDArray) else k for k in keys)
+        new_keys = _tile_keys(
+            self.data, tuple(k.data if isinstance(k, NDArray) else k for k in keys)
+        )
         target_shape = self.data[new_keys].shape
+        reduce = value if isinstance(value, LoopReduced) else None
         value = np.asarray(value.data if isinstance(value, NDArray) else value)
         if value.shape != target_shape and value.size == int(np.prod(target_shape)):
             value = value.reshape(target_shape)  # e.g. a [P, 1] result into a [P] slot
+        if reduce is not None:
+            # `dst[...] = nl.loop_reduce(x, op, ...)`: the destination is op's other input
+            current = self.data[new_keys]
+            value = reduce.op(current, np.broadcast_to(value, target_shape))
+            if reduce.mask is not None:
+                value = np.where(reduce.mask.reshape(target_shape), value, current)
         self.data[new_keys] = value
 
         return self
@@ -301,9 +377,22 @@ def __and__(self, other):
     def __or__(self, other):
         return self._binary_op(other, lambda a, b: a | b, "or", "|")
 
-    def reshape(self, *args, **kwargs):
+    def reshape(self, *args, shape=None, **kwargs):
+        # `t.reshape(shape)`, `t.reshape(*shape)`, or `t.reshape(shape=...)`; a view of the
+        # same memory, as reshaping an HBM tensor is in NKI.
+        if shape is None:
+            shape = args[0] if len(args) == 1 else args
+        return NDArray(
+            buffer=self.buffer,
+            value=self.data.reshape(_shape(shape)),
+            name=f"{self.name}_reshape",
+            **kwargs,
+        )
+
+    def view(self, dtype):
+        # the same memory reinterpreted as another dtype (e.g. fp32 values as 4 fp8 bytes)
         return NDArray(
-            value=self.data.reshape(*args), name=f"{self.name}_reshape", **kwargs
+            buffer=self.buffer, value=self.data.view(dtype), name=f"{self.name}_view"
         )
 
     def broadcast_to(self, *args, **kwargs):
@@ -314,6 +403,44 @@ def broadcast_to(self, *args, **kwargs):
         )
 
 
+class HBMAccess(NDArray):
+    """`t[keys]` of an HBM tensor outside a load or store, as in `a = t[i, j]` followed by
+    `nl.load(a, mask=m)`: in NKI an access pattern that the load transfers, applying its
+    mask. Load and store take `tensor` and `keys` from it, so clients see the access to
+    `tensor`; its own value is read only when used otherwise (an unmasked read)."""
+
+    def __init__(self, tensor, keys):
+        self._value = None
+        super().__init__(buffer=tensor.buffer, name=f"{tensor.name}_access")
+        self.tensor, self.keys = tensor, keys
+        self._parent = (tensor, nki_builder._convert_keys_to_numpy(keys))
+
+    @property
+    def data(self):
+        if self._value is None and getattr(self, "tensor", None) is not None:
+            self._value = self.tensor.data[self._parent[1]]
+        return self._value
+
+    @data.setter
+    def data(self, value):
+        self._value = value
+
+
+class LoopReduced(NDArray):
+    """The result of `nl.loop_reduce(x, op, ...)`. Assigned into a tensor slice, it combines
+    with the slice's current value: `dst[...] = op(dst[...], x)`, as in NKI's simulator.
+    Bound to a plain name it is just `x`."""
+
+    def __init__(self, value, op, mask=None, **kwargs):
+        super().__init__(value=value, **kwargs)
+        self.op = op
+        self.mask = (
+            None
+            if mask is None
+            else np.broadcast_to(np.asarray(mask, bool), value.shape)
+        )
+
+
 class GridIndices(NDArray):
     """The result of `nl.mgrid[...]`: unpacks into one index array per axis, and, as in NKI,
     also names them `.p` (the partition axis), `.x`, `.y`, and `.z` (the free axes)."""
@@ -352,6 +479,8 @@ def __init__(self, grid_dims=None):
         self.grid_z = None
         self.fn = None
         self.shared_hbm_arrays = {}
+        self.rng = np.random.default_rng(0)  # nl.dropout; reseeded by nl.random_seed
+        self.accumulator = None  # the Vector Engine's per-partition reduce registers
 
     def set_grid_dim(self, *grid_dims):
         self.grid_dims = grid_dims
@@ -424,11 +553,10 @@ def _convert_keys_to_numpy(self, keys):
 
     def load(self, src: NDArray, keys=None, *, mask=None, **kwargs):
         """Load array elements with masking for out-of-bounds errors."""
-        # Convert NDArray to numpy array
+        src, keys = access(src, keys)
         ndarray = src.data
-        mask_value = getattr(mask, "data", mask) if mask is not None else None
+        mask_value = _value(mask)
         dtype = kwargs.pop("dtype", None)
-        keys = full_keys(src) if keys is None else keys
 
         # Convert any NDArrays in keys to numpy arrays
         numpy_keys = self._convert_keys_to_numpy(keys)
@@ -444,13 +572,18 @@ def load(self, src: NDArray, keys=None, *, mask=None, **kwargs):
         )
 
     def store(
-        self, dst: NDArray, keys=None, value: NDArray = None, *, mask=None, **kwargs
+        self,
+        dst: NDArray,
+        keys=None,
+        value: NDArray | None = None,
+        *,
+        mask=None,
+        **kwargs,
     ):
         """Store array elements with masking for out-of-bounds errors."""
-        # Convert NDArrays to numpy arrays
+        dst, keys = access(dst, keys)
         ndarray = dst.data
-        mask_value = getattr(mask, "data", mask) if mask is not None else None
-        keys = full_keys(dst) if keys is None else keys
+        mask_value = _value(mask)
 
         # Convert any NDArrays in keys to numpy arrays
         numpy_keys = self._convert_keys_to_numpy(keys)
@@ -474,8 +607,6 @@ def store(
     def load_transpose2d(
         self, src: NDArray, keys=None, *, mask=None, dtype=None, **kwargs
     ):
-        if keys is None:
-            keys = tuple(slice(None) for _ in range(src.data.ndim))
         loaded = self.load(src, keys, mask=mask, **kwargs)
         value = loaded.data.astype(dtype) if dtype is not None else loaded.data
         return NDArray(value=value.T, name=f"{src.name}_load_transpose2d", **kwargs)
@@ -702,11 +833,14 @@ def _reduce_into(self, reduce_op, value, reduce_res):
         reduced = _ufunc(reduce_op).reduce(
             value, axis=_free_axes(value.ndim), keepdims=True
         )
-        reduce_res.data[...] = np.broadcast_to(
-            reduced.reshape(
-                reduce_res.data.shape[:1] + (1,) * (reduce_res.data.ndim - 1)
+        _write_into(
+            reduce_res,
+            np.broadcast_to(
+                reduced.reshape(
+                    reduce_res.data.shape[:1] + (1,) * (reduce_res.data.ndim - 1)
+                ),
+                reduce_res.data.shape,
             ),
-            reduce_res.data.shape,
         )
 
     def nc_matmul(self, stationary, moving, *, mask=None, **kwargs):
@@ -735,12 +869,13 @@ def activation(
         dtype=None,
         **kwargs,
     ):
-        nd = data.data.ndim
+        ndim = data.data.ndim
         value = data.data.astype(np.float32) * np.asarray(
-            _partition_aligned(scale, nd), dtype=np.float32
+            _partition_aligned(scale, ndim), dtype=np.float32
         )
         if bias is not None:
-            value = value + _partition_aligned(bias, nd)
+            value, bias_value = _aligned_pair(value, bias)  # either may be a [P] vector
+            value = value + bias_value
         value = self._activate(op, value)
         self._reduce_into(reduce_op, value, reduce_res)
         return NDArray(
@@ -791,6 +926,8 @@ def tensor_scalar(
         **kwargs,
     ):
         def apply(op, x, operand, reverse):
+            if getattr(op, "__name__", None) in ("abs", "absolute"):
+                return np.abs(x)  # NKI requires its operand to be zero
             u = _ufunc(op)
             operand = _partition_aligned(operand, np.ndim(x))
             return u(operand, x) if reverse else u(x, operand)
@@ -956,13 +1093,197 @@ def dma_transpose(
         return NDArray(value=_cast(value, dtype), name=f"{src.name}_dma_transpose")
 
     def get_nc_version(self):
-        # The interpreter models a Trn1 NeuronCore.
+        # The interpreter models a Trn1 NeuronCore (gen2), or a Trn2 one (gen3) when NKI's
+        # platform target is overridden to trn2, as for NKI's own simulator.
+        if os.environ.get("NEURON_PLATFORM_TARGET_OVERRIDE", "").startswith("trn2"):
+            return nisa.nc_version.gen3
         return nisa.nc_version.gen2
 
+    # ---- nl: selection, loop and program reductions, random ---------------
+    def where(self, condition, x, y, *, dtype=None, mask=None, **kwargs):
+        xv, yv, cv = _value(x), _value(y), _value(condition)
+        ndim = max(np.ndim(xv), np.ndim(yv), np.ndim(cv))
+        value = np.where(
+            np.asarray(_partition_aligned(cv, ndim), dtype=bool),
+            _partition_aligned(xv, ndim),
+            _partition_aligned(yv, ndim),
+        )
+        return NDArray(value=_cast(value, dtype), name="where")
+
+    def equal(self, x, y, *, dtype=bool, mask=None, **kwargs):
+        return self._binary_op(x, y, np.equal, "equal", dtype=dtype)
+
+    def loop_reduce(self, x, op, loop_indices, *, dtype=None, mask=None, **kwargs):
+        value = _cast(_value(x), dtype)
+        mask_value = None if mask is None else _value(mask)
+        return LoopReduced(value, _ufunc(op), mask_value, name="loop_reduce")
+
+    def all_reduce(self, x, op, program_axes, *, dtype=None, mask=None, **kwargs):
+        # A reduction across the programs of the grid. The interpreter runs the programs one
+        # after another, so it can only reduce over a grid of one program.
+        if int(np.prod(self.grid_dims)) != 1:
+            raise NotImplementedError(
+                "nl.all_reduce across several programs is not supported by the interpreter"
+            )
+        return NDArray(value=_cast(np.copy(_value(x)), dtype), name="all_reduce")
+
+    def random_seed(self, seed, *, mask=None, **kwargs):
+        self.rng = np.random.default_rng(int(np.asarray(_value(seed)).ravel()[0]))
+
+    def dropout(self, x, rate, *, dtype=None, mask=None, **kwargs):
+        # Zeroes each element with probability `rate` (no rescaling, as in NKI). Like NKI's
+        # simulator, the random numbers do not match the hardware's.
+        rate = _partition_aligned(_value(rate), x.data.ndim)
+        keep = self.rng.random(x.data.shape) >= np.asarray(rate, dtype=np.float32)
+        return NDArray(
+            value=_cast(np.where(keep, x.data, 0).astype(x.data.dtype), dtype),
+            name=f"{x.name}_dropout",
+        )
+
+    # ---- nisa: statistics, partition reduction, predicated copies -----------
+    def tensor_partition_reduce(self, op, data, *, mask=None, dtype=None, **kwargs):
+        value = _ufunc(op).reduce(data.data, axis=0, keepdims=True)
+        return NDArray(
+            value=_cast(value, dtype if dtype is not None else data.data.dtype),
+            name=f"{data.name}_partition_reduce",
+        )
+
+    def bn_stats(self, data, *, mask=None, dtype=None, **kwargs):
+        # Per partition: (count, mean, variance * count) of the elements at even positions,
+        # then of those at odd positions, over the elements the mask keeps; in fp32.
+        values = _per_partition(data.data).astype(np.float32)
+        keep = (
+            np.ones(values.shape, dtype=bool)
+            if mask is None
+            else _per_partition(_bool_tile(_value(mask), data.data.shape))
+        )
+        stats = []
+        for parity in (0, 1):
+            v, k = values[:, parity::2], keep[:, parity::2]
+            count = k.sum(axis=1, keepdims=True).astype(np.float32)
+            mean = np.where(k, v, 0).sum(axis=1, keepdims=True) / np.maximum(count, 1)
+            m2 = np.square(np.where(k, v - mean, 0)).sum(axis=1, keepdims=True)
+            stats += [count, mean, m2]
+        return NDArray(
+            value=_cast(
+                np.concatenate(stats, axis=1),
+                dtype if dtype is not None else data.data.dtype,
+            ),
+            name=f"{data.name}_bn_stats",
+        )
+
+    def bn_aggr(self, data, *, mask=None, dtype=None, **kwargs):
+        # Combines (count, mean, variance * count) triples per partition into (mean, variance).
+        triples = _per_partition(data.data).astype(np.float32)
+        triples = triples.reshape(triples.shape[0], -1, 3)
+        count, mean, m2 = triples[..., 0], triples[..., 1], triples[..., 2]
+        total = count.sum(axis=1, keepdims=True)
+        overall = (count * mean).sum(axis=1, keepdims=True) / np.maximum(total, 1)
+        m2_total = m2.sum(axis=1, keepdims=True) + (
+            count * np.square(mean - overall)
+        ).sum(axis=1, keepdims=True)
+        value = np.concatenate([overall, m2_total / np.maximum(total, 1)], axis=1)
+        return NDArray(
+            value=_cast(value, dtype if dtype is not None else data.data.dtype),
+            name=f"{data.name}_bn_aggr",
+        )
+
+    def select_reduce(
+        self,
+        *,
+        dst,
+        predicate,
+        on_true,
+        on_false,
+        reduce_res=None,
+        reduce_cmd=None,
+        reduce_op=None,
+        reverse_pred=False,
+        mask=None,
+        dtype=None,
+        **kwargs,
+    ):
+        # dst = where(predicate, on_true, on_false); with a reduce command, the per-partition
+        # max of the result also goes through the accumulator registers into `reduce_res`.
+        true_value = _per_partition(_value(on_true))
+        pred = _per_partition(_bool_tile(_value(predicate), np.shape(_value(on_true))))
+        if reverse_pred:
+            pred = ~pred
+        false_value = _partition_aligned(_value(on_false), 2)
+        result = np.where(pred, true_value, false_value)
+        _write_into(dst, _cast(result, dtype if dtype is not None else dst.data.dtype))
+        command = getattr(reduce_cmd, "name", "idle")
+        if command == "idle":
+            return
+        partial = np.max(result.astype(np.float32), axis=1)
+        if command == "reset" or self.accumulator is None:
+            self.accumulator = np.full(partial.shape, -np.inf, dtype=np.float32)
+        if command == "reset_reduce":
+            self.accumulator = partial
+        elif command in ("reduce", "load_reduce"):
+            self.accumulator = np.maximum(self.accumulator, partial)
+        if reduce_res is not None:
+            _write_into(reduce_res, self.accumulator.astype(reduce_res.data.dtype))
+
+    def tensor_copy_predicated(
+        self,
+        *,
+        src,
+        dst,
+        predicate,
+        reverse_pred=False,
+        mask=None,
+        dtype=None,
+        **kwargs,
+    ):
+        # dst = where(predicate, src, dst), matching operands by partition.
+        current = _per_partition(dst.data)
+        pred = _per_partition(
+            _bool_tile(_value(predicate), np.shape(_value(predicate)))
+        )
+        pred = np.broadcast_to(pred, current.shape)
+        if reverse_pred:
+            pred = ~pred
+        src_value = _value(src)
+        source = (
+            _per_partition(src_value)
+            if np.ndim(src_value) == dst.data.ndim or np.size(src_value) == current.size
+            else _partition_aligned(src_value, 2)
+        )
+        _write_into(dst, np.where(pred, source, current).astype(dst.data.dtype))
+
+    def tiled_offloaded_memcpy(
+        self, *, dst, src, dst_keys=None, src_keys=None, dtype=None, **kwargs
+    ):
+        # A DMA copy from one HBM tile to another, seen by clients as a load and a store.
+        self.store(dst, dst_keys, self.load(src, src_keys))
+
+    def tiled_offloaded_fma(
+        self,
+        *src,
+        dst,
+        scales,
+        src_keys=None,
+        dst_keys=None,
+        dtype=None,
+        mask=None,
+        **kwargs,
+    ):
+        # dst = src[0] * scales[0] + src[1] * scales[1] + ..., by the DMA engine, in HBM.
+        keys = src_keys if src_keys is not None else [None] * len(src)
+        acc = 0
+        for tensor, key, scale in zip(src, keys, scales):
+            acc = acc + self.load(tensor, key).data.astype(np.float32) * scale
+        self.store(dst, dst_keys, NDArray(value=np.asarray(acc)))
+
     def range(self, *args, **kwargs):
         # affine_range, sequential_range, and static_range all iterate like range()
         return range(*args)
 
+    def multi_buffer(self, *args, **kwargs):
+        # a scheduling directive for a range (`directives=`), which ranges ignore
+        return None
+
 
 nki_builder = Builder()
 
@@ -1009,6 +1330,12 @@ def _set_attr(obj, name, value):
         "program_ndim",
         "floor",
         "mod",
+        "where",
+        "equal",
+        "loop_reduce",
+        "all_reduce",
+        "random_seed",
+        "dropout",
     ):
         _set_attr(nl, name, getattr(nki_builder, name))
     for name in (
@@ -1029,9 +1356,17 @@ def _set_attr(obj, name, value):
         "iota",
         "tensor_tensor_scan",
         "get_nc_version",
+        "tensor_partition_reduce",
+        "bn_stats",
+        "bn_aggr",
+        "select_reduce",
+        "tensor_copy_predicated",
     ):
         _set_attr(nisa, name, getattr(nki_builder, name))
     _set_attr(nisa, "reciprocal", nki_builder.nisa_reciprocal)
+    _set_attr(ncc, "multi_buffer", nki_builder.multi_buffer)
+    _set_attr(nisa, "_tiled_offloaded_memcpy", nki_builder.tiled_offloaded_memcpy)
+    _set_attr(nisa, "_tiled_offloaded_fma", nki_builder.tiled_offloaded_fma)
     for name in (
         "max",
         "min",
@@ -1074,11 +1409,75 @@ def nki_unpatch_lang(scope=None):
         scope.restore()
 
 
+def nki_kernel_function(obj):
+    """The Python function of an NKI kernel object (`@nki.jit`, `@nki.baremetal`, ...), or
+    None for anything else."""
+    func = getattr(obj, "func", None)
+    if type(obj).__module__.startswith("neuronxcc") and inspect.isfunction(func):
+        return func
+    return None
+
+
+def _compile_rewritten(fn, namespace):
+    """Compile `fn` with the NKI frontend AST transformer applied (decorators dropped, loads
+    and stores given keys) into `namespace`, and return the new function."""
+    # source of the function, stripped of leading indents in case it was defined in scope
+    transformed_code = transform_code(textwrap.dedent(inspect.getsource(fn)))
+    import random
+    import string
+
+    rand_str = "".join(random.choices(string.ascii_letters + string.digits, k=16))
+    os.makedirs("/tmp/tilelens", exist_ok=True)
+    filename = f"/tmp/tilelens/{rand_str}.py"
+    with open(filename, "w") as f:
+        f.write(transformed_code)
+    exec(compile(transformed_code, filename=filename, mode="exec"), namespace)
+    new_fn = namespace[fn.__name__]
+    CODE_KEYS.add(get_code_key(new_fn))  # trace rewritten function for clients
+    return new_fn
+
+
+def _rewritten_namespace(module_globals, namespaces):
+    """A copy of a module's globals in which its NKI kernels (`@nki.jit` functions, which
+    kernels call as helpers) are rewritten, and NKI kernels it imports from other modules
+    are their rewrites there. One namespace per module, so kernels that call each other
+    find the rewrites. `namespaces` maps id(module globals) to the namespaces made so far."""
+    key = id(module_globals)
+    if key in namespaces:
+        return namespaces[key]
+    namespace = namespaces[key] = module_globals.copy()
+    for name, value in module_globals.items():
+        func = nki_kernel_function(value)
+        if func is None:
+            continue
+        if func.__globals__ is module_globals:
+            namespace[name] = _compile_rewritten(func, namespace)
+        else:
+            other = _rewritten_namespace(func.__globals__, namespaces)
+            namespace[name] = other.get(func.__name__, value)
+    return namespace
+
+
+def _rewrite_function(fn):
+    """`fn` rewritten by the NKI frontend AST transformer, in the rewritten namespace of its
+    module, so the NKI kernels it calls by name are rewritten too."""
+    return _compile_rewritten(fn, _rewritten_namespace(fn.__globals__, {}))
+
+
 def full_keys(tensor):
     """Keys of a whole-tensor access such as `nl.load(t)` or `nl.store(t, value=v)`."""
     return tuple(slice(None) for _ in range(tensor.data.ndim))
 
 
+def access(tensor, keys):
+    """The tensor and keys a load or store accesses: `nl.load(t[k])` arrives as (t, k),
+    `nl.load(t)` as (t, None), and `a = t[k]; nl.load(a)` as (a, None) with `a` an
+    HBMAccess."""
+    if keys is None and isinstance(tensor, HBMAccess):
+        return tensor.tensor, tensor.keys
+    return tensor, (full_keys(tensor) if keys is None else keys)
+
+
 def _to_numpy(ret):
     if isinstance(ret, NDArray):
         return ret.data
@@ -1107,6 +1506,8 @@ def run(self, *args, **kwargs):
             )
         nki_builder.set_grid_dim(*grid_dims)
         nki_builder.shared_hbm_arrays = {}
+        nki_builder.rng = np.random.default_rng(0)
+        nki_builder.accumulator = None
         nki_builder.fn = self.fn
 
         kwargs.pop("warmup", None)  # Remove warmup from kwargs if it exists
@@ -1118,36 +1519,27 @@ def run(self, *args, **kwargs):
         if client_manager is not None:
             client_manager.grid_callback(grid_dims)
 
-        # Apply the NKI frontend AST transformer.
-        if hasattr(self.fn, "__code__"):
-            # Get the source code of the function (stripped of leading indents in case it was defined in scope)
-            source_code = textwrap.dedent(inspect.getsource(self.fn))
-            # Transform the source code using the AST transformer
-            transformed_code = transform_code(source_code)
-            # Create a new function from the transformed code
-            exec_globals = self.fn.__globals__.copy()
-            import random
-            import string
-            import os
-
-            rand_str = "".join(
-                random.choices(string.ascii_letters + string.digits, k=16)
-            )
-            os.makedirs("/tmp/tilelens", exist_ok=True)
-            filename = f"/tmp/tilelens/{rand_str}.py"
-            with open(filename, "w") as f:
-                f.write(transformed_code)
-            code_obj = compile(transformed_code, filename=filename, mode="exec")
-            exec(code_obj, exec_globals)
-            self.fn = exec_globals[self.fn.__name__]
-            CODE_KEYS.add(get_code_key(self.fn))  # trace rewritten function for clients
-
-        # convert args to NDArray if they are not already
+        # Apply the NKI frontend AST transformer, to the kernel and to the NKI kernels it
+        # calls (`@nki.jit` helpers), once.
+        if not getattr(self, "_rewritten", False):
+            fn = nki_kernel_function(self.fn) or self.fn
+            if hasattr(fn, "__code__"):
+                self.fn = _rewrite_function(fn)
+            self._rewritten = True
+
+        # convert args to NDArray if they are not already, also those passed by keyword
         args = [
-            NDArray(value=arg) if isinstance(arg, np.ndarray) else arg for arg in args
+            NDArray(value=arg, buffer=nl.hbm) if isinstance(arg, np.ndarray) else arg
+            for arg in args
         ]
-
-        name_args = inspect.getcallargs(self.fn, *args)
+        kwargs = {
+            name: NDArray(value=arg, buffer=nl.hbm)
+            if isinstance(arg, np.ndarray)
+            else arg
+            for name, arg in kwargs.items()
+        }
+
+        name_args = inspect.getcallargs(self.fn, *args, **kwargs)
         call_args = {}
         for name, arg in name_args.items():
             call_args[name] = arg
@@ -1165,7 +1557,9 @@ def run(self, *args, **kwargs):
                         client_manager.grid_idx_callback((x, y, z))
 
                     if not client_manager.pre_run_callback(self.fn):
-                        return _to_numpy(ret)
+                        # a client that needs only some programs (the profiler's block
+                        # sampling) skips this one, as the Triton frontend does
+                        continue
                     ret = self.fn(*args, **kwargs)
                     if not client_manager.post_run_callback(self.fn):
                         # a client that needs no more programs stops the grid early
diff --git a/tilelens/core/simulation/nki_beta2.py b/tilelens/core/simulation/nki_beta2.py
index 5e9f79fc..9c69ad2c 100644
--- a/tilelens/core/simulation/nki_beta2.py
+++ b/tilelens/core/simulation/nki_beta2.py
@@ -1232,7 +1232,7 @@ def run(self, *args: Any, **kwargs: Any) -> None:
                     tuple(int(idx) for idx in grid_idx) + (0,) * (3 - len(grid_idx))
                 )
                 if not client_manager.pre_run_callback(self.fn):
-                    return
+                    continue  # a client that needs only some programs skips this one
             self.fn(*kernel_args, **kwargs)
             if client_manager is not None and not client_manager.post_run_callback(
                 self.fn
diff --git a/tilelens/core/trace.py b/tilelens/core/trace.py
index abdd11a1..16b9d036 100644
--- a/tilelens/core/trace.py
+++ b/tilelens/core/trace.py
@@ -433,6 +433,12 @@ def decorator(kernel) -> TritonTrace | NKITrace | GluonTrace | Any:
             trace.add_client(client)
             return trace
 
+        if frontend == "nki":
+            # an `@nki.jit` (or `@nki.baremetal`, ...) kernel runs as its Python function
+            from .simulation.nki import nki_kernel_function
+
+            kernel = nki_kernel_function(kernel) or kernel
+
         trace_source(kernel)
 
         # First-time wrapping

From 3f5bb1b0b712196547699e42108e32e81d2250b3 Mon Sep 17 00:00:00 2001
From: Hao Wu <hwu27@gmu.edu>
Date: Sun, 8 Nov 2026 18:22:35 -0500
Subject: [PATCH 2/6] [FEAT] Run the NKI beta2 matrix-multiplication tutorial
 kernels as written

- tilelens.trace(frontend="nki_beta2") accepts an `@nki.jit` kernel of the `nki` package,
  as frontend="nki" does for `neuronxcc.nki` (nki_kernel_function moves to trace.py).
- The beta2 runner returns the kernel's outputs (as numpy arrays), and a program a client
  declines is skipped.
- nl.shared_hbm is an HBM buffer; nisa.memset; nl.sequential_range and nl.static_range
  iterate like affine_range; nisa.tensor_copy accepts the `dtype` that NKI's tracer
  accepts and the tutorial passes.

Test: an `@nki.jit` kernel written like the tutorial (checked by NKI's tracer first)
returns its output. Full suite with -m "": 647 passed, 2 skipped.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0169QRxKtsivkLDCS5f7rrsm
---
 tests/nki/test_nki_beta2.py           | 36 +++++++++++++++++++++
 tilelens/core/simulation/nki.py       | 13 +++-----
 tilelens/core/simulation/nki_beta2.py | 45 +++++++++++++++++++++++++--
 tilelens/core/trace.py                | 15 +++++++--
 4 files changed, 95 insertions(+), 14 deletions(-)

diff --git a/tests/nki/test_nki_beta2.py b/tests/nki/test_nki_beta2.py
index 46d49e76..3e7f0175 100644
--- a/tests/nki/test_nki_beta2.py
+++ b/tests/nki/test_nki_beta2.py
@@ -2219,3 +2219,39 @@ def post_run_callback(self, fn):
 
     with pytest.raises(ValueError, match="1D SPMD"):
         fn.run(src, np.zeros_like(src), grid=(1,), client_manager=_RankMismatch())
+
+
+def test_jit_kernel_written_like_the_matmul_tutorial_returns_its_output():
+    """An `@nki.jit` kernel as the matrix-multiplication tutorial writes it: its output in
+    shared HBM, returned; `memset`; `tensor_copy` with a `dtype`; and a sequential range.
+    NKI's own tracer checks the kernel first (pre_trace)."""
+    import nki
+
+    @nki.jit
+    def matmul_kernel(lhsT, rhs):
+        K, M = lhsT.shape
+        _, N = rhs.shape
+        result = nl.ndarray((M, N), dtype=lhsT.dtype, buffer=nl.shared_hbm)
+        lhs_tile = nl.ndarray(lhsT.shape, dtype=lhsT.dtype, buffer=nl.sbuf)
+        rhs_tile = nl.ndarray(rhs.shape, dtype=rhs.dtype, buffer=nl.sbuf)
+        nisa.dma_copy(dst=lhs_tile, src=lhsT)
+        nisa.dma_copy(dst=rhs_tile, src=rhs)
+        acc = nl.ndarray((M, N), dtype=nl.float32, buffer=nl.sbuf)
+        nisa.memset(dst=acc, value=0.0)
+        for _ in nl.sequential_range(2):
+            psum = nl.ndarray((M, N), dtype=nl.float32, buffer=nl.psum)
+            nisa.nc_matmul(psum, lhs_tile, rhs_tile)
+            nisa.tensor_tensor(dst=acc, data1=acc, data2=psum, op=nl.add)
+        out = nl.ndarray((M, N), dtype=lhsT.dtype, buffer=nl.sbuf)
+        nisa.tensor_copy(dst=out, src=acc, dtype=lhsT.dtype)
+        nisa.dma_copy(dst=result, src=out)
+        return result
+
+    rng = np.random.default_rng(0)
+    lhsT = rng.random((128, 64)).astype(np.float32)
+    rhs = rng.random((128, 512)).astype(np.float32)
+    tilelens.clear()
+    got = tilelens.trace(client=Tracer(), frontend="nki_beta2")(matmul_kernel)[(1,)](
+        lhsT, rhs
+    )
+    np.testing.assert_allclose(got, 2 * lhsT.T @ rhs, rtol=1e-5)
diff --git a/tilelens/core/simulation/nki.py b/tilelens/core/simulation/nki.py
index 96eb838e..145b7d32 100644
--- a/tilelens/core/simulation/nki.py
+++ b/tilelens/core/simulation/nki.py
@@ -1409,15 +1409,6 @@ def nki_unpatch_lang(scope=None):
         scope.restore()
 
 
-def nki_kernel_function(obj):
-    """The Python function of an NKI kernel object (`@nki.jit`, `@nki.baremetal`, ...), or
-    None for anything else."""
-    func = getattr(obj, "func", None)
-    if type(obj).__module__.startswith("neuronxcc") and inspect.isfunction(func):
-        return func
-    return None
-
-
 def _compile_rewritten(fn, namespace):
     """Compile `fn` with the NKI frontend AST transformer applied (decorators dropped, loads
     and stores given keys) into `namespace`, and return the new function."""
@@ -1442,6 +1433,8 @@ def _rewritten_namespace(module_globals, namespaces):
     kernels call as helpers) are rewritten, and NKI kernels it imports from other modules
     are their rewrites there. One namespace per module, so kernels that call each other
     find the rewrites. `namespaces` maps id(module globals) to the namespaces made so far."""
+    from ..trace import nki_kernel_function
+
     key = id(module_globals)
     if key in namespaces:
         return namespaces[key]
@@ -1522,6 +1515,8 @@ def run(self, *args, **kwargs):
         # Apply the NKI frontend AST transformer, to the kernel and to the NKI kernels it
         # calls (`@nki.jit` helpers), once.
         if not getattr(self, "_rewritten", False):
+            from ..trace import nki_kernel_function
+
             fn = nki_kernel_function(self.fn) or self.fn
             if hasattr(fn, "__code__"):
                 self.fn = _rewrite_function(fn)
diff --git a/tilelens/core/simulation/nki_beta2.py b/tilelens/core/simulation/nki_beta2.py
index 9c69ad2c..61529b91 100644
--- a/tilelens/core/simulation/nki_beta2.py
+++ b/tilelens/core/simulation/nki_beta2.py
@@ -362,7 +362,18 @@ def view(
         return probe
 
 
+def _outputs_to_numpy(ret: Any) -> Any:
+    """A kernel's returned tensors (one, or a tuple or list of them) as numpy arrays."""
+    if isinstance(ret, NDArray):
+        return np.asarray(ret.data)
+    if isinstance(ret, (tuple, list)):
+        return type(ret)(_outputs_to_numpy(r) for r in ret)
+    return ret
+
+
 sbuf, hbm, psum = map(Buffer, ("sbuf", "hbm", "psum"))
+# kernels allocate the outputs they return in shared HBM: device memory like hbm
+shared_hbm = hbm
 
 
 def ndarray(
@@ -751,11 +762,33 @@ def dma_copy(
     return _store(dst, _array(dst_rmw_op.run(dst.data, src_value)))
 
 
+def memset(
+    dst: NDArray,
+    value: Any,
+    engine: Any = nisa.engine.unknown,
+    name: str | None = None,
+) -> NDArray:
+    """Fill an on-chip tensor with a constant.
+
+    Args:
+        dst: Destination tensor in SBUF or PSUM.
+        value: The constant.
+        engine: Requested execution engine token.
+        name: Optional kernel op name retained for API parity.
+    """
+    _require(
+        _buffer_name(dst) in ("sbuf", "psum"),
+        f"memset is on-chip only, received dst buffer {_buffer_name(dst)}",
+    )
+    return _store(dst, np.full(dst.shape, value))
+
+
 def tensor_copy(
     dst: NDArray,
     src: NDArray,
     engine: Any = nisa.engine.unknown,
     name: str | None = None,
+    dtype: Any = None,
 ) -> NDArray:
     """Copy tensor data within on-chip memory.
 
@@ -764,6 +797,8 @@ def tensor_copy(
         src: Source tensor in SBUF or PSUM.
         engine: Requested execution engine token.
         name: Optional kernel op name retained for API parity.
+        dtype: Output dtype, which NKI's tracer accepts (e.g. the matrix-multiplication
+            tutorial casts PSUM to the output dtype); the copy converts to ``dst``'s dtype.
     """
     dst_buffer, src_buffer = _buffer_name(dst), _buffer_name(src)
     _require(
@@ -1149,6 +1184,8 @@ def nki_patch_lang(scope: Any = None) -> None:
     set_attr(nl, "psum", psum)
     set_attr(nl, "program_id", program_id)
     set_attr(nl, "affine_range", affine_range)
+    set_attr(nl, "sequential_range", affine_range)  # all iterate like range()
+    set_attr(nl, "static_range", affine_range)
     set_attr(nl, "ds", ds)
 
     set_attr(nisa, "nc_matmul", nc_matmul)
@@ -1158,6 +1195,7 @@ def nki_patch_lang(scope: Any = None) -> None:
     set_attr(nisa, "reciprocal", reciprocal)
     set_attr(nisa, "dma_copy", dma_copy)
     set_attr(nisa, "tensor_copy", tensor_copy)
+    set_attr(nisa, "memset", memset)
     set_attr(nisa, "tensor_tensor", tensor_tensor)
     set_attr(nisa, "tensor_reduce", tensor_reduce)
     set_attr(nisa, "tensor_scalar", tensor_scalar)
@@ -1220,6 +1258,7 @@ def run(self, *args: Any, **kwargs: Any) -> None:
             assert arg is None or isinstance(arg, (bool, int, float, str, NDArray))
             if client_manager is not None:
                 client_manager.arg_callback(name, arg, arg)
+        ret = None
         for grid_idx in itertools.product(*(range(dim) for dim in grid_dims)):
             if len(grid_idx) != len(nki_builder.grid_dims):
                 raise ValueError(
@@ -1233,8 +1272,10 @@ def run(self, *args: Any, **kwargs: Any) -> None:
                 )
                 if not client_manager.pre_run_callback(self.fn):
                     continue  # a client that needs only some programs skips this one
-            self.fn(*kernel_args, **kwargs)
+            ret = self.fn(*kernel_args, **kwargs)
             if client_manager is not None and not client_manager.post_run_callback(
                 self.fn
             ):
-                return
+                break
+        # a kernel that allocates its outputs in HBM returns them, as numpy arrays
+        return _outputs_to_numpy(ret)
diff --git a/tilelens/core/trace.py b/tilelens/core/trace.py
index 16b9d036..32199c71 100644
--- a/tilelens/core/trace.py
+++ b/tilelens/core/trace.py
@@ -363,6 +363,17 @@ def warmup(self, *args, **kwargs):
         return self.run(*args, warmup=True, **kwargs)
 
 
+def nki_kernel_function(obj):
+    """The Python function of an NKI kernel object (`@nki.jit`, `@nki.baremetal`, ... of the
+    legacy `neuronxcc.nki` or of the newer `nki` package), or None for anything else."""
+    func = getattr(obj, "func", None)
+    if type(obj).__module__.split(".")[0] in ("neuronxcc", "nki") and isinstance(
+        func, types.FunctionType
+    ):
+        return func
+    return None
+
+
 def trace_source(kernel):
     """
     Add the kernel code to be traceable within stack traces for clients
@@ -433,10 +444,8 @@ def decorator(kernel) -> TritonTrace | NKITrace | GluonTrace | Any:
             trace.add_client(client)
             return trace
 
-        if frontend == "nki":
+        if frontend in ("nki", "nki_beta2"):
             # an `@nki.jit` (or `@nki.baremetal`, ...) kernel runs as its Python function
-            from .simulation.nki import nki_kernel_function
-
             kernel = nki_kernel_function(kernel) or kernel
 
         trace_source(kernel)

From 9f5af16160e0887a4b7825324feab192cdabeb2a Mon Sep 17 00:00:00 2001
From: Hao Wu <hwu27@gmu.edu>
Date: Sun, 8 Nov 2026 18:28:16 -0500
Subject: [PATCH 3/6] [FEAT] Interpret double-row nc_matmul
 (perf_mode="double_row_gen3")

Its operands are [K, 2, M] and [K, 2, N]: the two rows of tiles are contracted together, so
nc_matmul now contracts every axis but the last (the 2-D case is unchanged). The reference
double-row fp8 matmul of the NKI samples runs on Trn2 and matches NKI's simulator.

Test: a double-row matmul compared with nki.simulate_kernel on the trn2 target.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0169QRxKtsivkLDCS5f7rrsm
---
 tests/nki/test_nki_more_ops.py  | 23 +++++++++++++++++++++++
 tilelens/core/simulation/nki.py | 10 +++++++---
 2 files changed, 30 insertions(+), 3 deletions(-)

diff --git a/tests/nki/test_nki_more_ops.py b/tests/nki/test_nki_more_ops.py
index 35b77b43..262fd590 100644
--- a/tests/nki/test_nki_more_ops.py
+++ b/tests/nki/test_nki_more_ops.py
@@ -332,3 +332,26 @@ def pre_run_callback(self, fn):
         x
     )
     np.testing.assert_array_equal(got[128:], x[128:])
+
+
+def double_row_kernel(lhsT, rhs):
+    # nc_matmul with perf_mode="double_row_gen3" (Trn2): [K, 2, M] x [K, 2, N] -> [M, N]
+    out = nl.ndarray(
+        (lhsT.shape[2], rhs.shape[2]), dtype=nl.float32, buffer=nl.shared_hbm
+    )
+    res = nl.zeros((lhsT.shape[2], rhs.shape[2]), dtype=nl.float32, buffer=nl.psum)
+    res[...] += nisa.nc_matmul(nl.load(lhsT), nl.load(rhs), perf_mode="double_row_gen3")
+    nl.store(out, value=res)
+    return out
+
+
+def test_double_row_matmul_matches_nki_simulator(monkeypatch):
+    monkeypatch.setenv("NEURON_PLATFORM_TARGET_OVERRIDE", "trn2")
+    lhsT = rng.random((128, 2, 128)).astype(nl.float8_e4m3)
+    rhs = rng.random((128, 2, 512)).astype(nl.float8_e4m3)
+    tilelens.clear()
+    got = tilelens.trace(client=Tracer(), frontend="nki")(double_row_kernel)[(1,)](
+        lhsT, rhs
+    )
+    want = nki.simulate_kernel(nki.jit(double_row_kernel)[(1,)], lhsT, rhs)
+    np.testing.assert_allclose(got, want, rtol=1e-5)
diff --git a/tilelens/core/simulation/nki.py b/tilelens/core/simulation/nki.py
index 145b7d32..b0312aab 100644
--- a/tilelens/core/simulation/nki.py
+++ b/tilelens/core/simulation/nki.py
@@ -844,9 +844,13 @@ def _reduce_into(self, reduce_op, value, reduce_res):
         )
 
     def nc_matmul(self, stationary, moving, *, mask=None, **kwargs):
-        # stationary.T @ moving on the Tensor Engine, accumulated in fp32 (PSUM).
-        value = np.matmul(
-            stationary.data.astype(np.float32).T, moving.data.astype(np.float32)
+        # stationary.T @ moving on the Tensor Engine, accumulated in fp32 (PSUM). With
+        # perf_mode="double_row_*" the operands are [K, 2, M] and [K, 2, N]: the two rows of
+        # tiles are contracted together, so every axis but the last is contracted.
+        s = stationary.data.astype(np.float32)
+        m = moving.data.astype(np.float32)
+        value = np.tensordot(
+            s, m, axes=(list(range(s.ndim - 1)), list(range(m.ndim - 1)))
         )
         return NDArray(value=value, name=f"{stationary.name}_{moving.name}_nc_matmul")
 

From 5ef5eacaa8c8923941ab76f15a021b2f94149af6 Mon Sep 17 00:00:00 2001
From: Hao Wu <hwu27@gmu.edu>
Date: Sun, 8 Nov 2026 19:27:38 -0500
Subject: [PATCH 4/6] [FEAT] Model SBUF tensors placed by address, and start
 PSUM tensors from zero

- A tensor placed with ncc.sbuf.mod_alloc or ncc.sbuf.alloc is a view of the NeuronCore's
  physical SBUF (bytes per partition, zeroed per launch), at the addresses NKI's modulo
  allocation gives (block index modulo the tile counts, linearized, times the tile size,
  plus the base). Tensors and logical tiles at one address share their values, and a tile
  declared in a loop at a fixed address keeps its values across iterations. Logical tiles
  that outnumber the physical ones (num_free_tiles smaller than the blocks) map their block
  indices modulo the period (AddressedTiles). nl.par_dim marks the partition dimension.
- A PSUM tensor starts from zero, also when placed in a bank an earlier tensor used (its
  first `+=` of a matmul is that matmul, as in NKI's simulator). PSUM tensors were
  uninitialized numpy arrays, which held zeros only when fresh pages happened to back them.

The allocated SD attention and fused RMSNorm-QKV reference kernels and the attention
tutorial's v10 now match NKI's simulator.

Tests: a double-buffered kernel and a tile declared in a loop compared with
nki.simulate_kernel; tiles at one address share their values (the device semantics; NKI's
simulator gives each tensor its own storage); PSUM tensors start from zero when a bank is
reused.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0169QRxKtsivkLDCS5f7rrsm
---
 tests/nki/test_nki_more_ops.py  |  98 +++++++++++++++++
 tilelens/core/simulation/nki.py | 183 ++++++++++++++++++++++++++++++--
 2 files changed, 275 insertions(+), 6 deletions(-)

diff --git a/tests/nki/test_nki_more_ops.py b/tests/nki/test_nki_more_ops.py
index 262fd590..ec4bce07 100644
--- a/tests/nki/test_nki_more_ops.py
+++ b/tests/nki/test_nki_more_ops.py
@@ -355,3 +355,101 @@ def test_double_row_matmul_matches_nki_simulator(monkeypatch):
     )
     want = nki.simulate_kernel(nki.jit(double_row_kernel)[(1,)], lhsT, rhs)
     np.testing.assert_allclose(got, want, rtol=1e-5)
+
+
+def double_buffered_kernel(x):
+    # Four logical tiles on two physical tiles (ncc.sbuf.mod_alloc), and a tile declared in
+    # the loop at a fixed address: as NKI's simulator and the device compute it.
+    out = nl.ndarray(x.shape, dtype=x.dtype, buffer=nl.shared_hbm)
+    bufs = nl.ndarray(
+        (4, nl.par_dim(128), x.shape[1]),
+        dtype=x.dtype,
+        buffer=ncc.sbuf.mod_alloc(base_addr=0, num_free_tiles=(2,)),
+    )
+    for i in nl.affine_range(4):
+        rows = nl.ds(i * 128, 128)
+        bufs[i] = nl.load(x[rows, :])
+        res = nl.ndarray(
+            (nl.par_dim(128), x.shape[1]),
+            dtype=x.dtype,
+            buffer=ncc.sbuf.mod_alloc(base_addr=2 * x.shape[1] * 4),
+        )
+        res[...] = nl.exp(bufs[i])
+        nl.store(out[rows, :], value=res)
+    return out
+
+
+def test_address_placed_tiles_match_nki_simulator():
+    x = rng.standard_normal((512, 64)).astype(np.float32)
+    tilelens.clear()
+    got = tilelens.trace(client=Tracer(), frontend="nki")(double_buffered_kernel)[(1,)](
+        x
+    )
+    want = nki.simulate_kernel(nki.jit(double_buffered_kernel)[(1,)], x)
+    np.testing.assert_allclose(got, want, rtol=1e-6)
+
+
+def shared_address_kernel(x):
+    out = nl.ndarray((2, 128, x.shape[1]), dtype=x.dtype, buffer=nl.shared_hbm)
+    t = nl.ndarray(
+        (2, nl.par_dim(128), x.shape[1]),
+        dtype=x.dtype,
+        buffer=ncc.sbuf.mod_alloc(base_addr=0, num_free_tiles=(1,)),
+    )
+    u = nl.ndarray(
+        (nl.par_dim(128), x.shape[1]),
+        dtype=x.dtype,
+        buffer=ncc.sbuf.mod_alloc(base_addr=0),
+    )
+    t[0] = nl.load(x)
+    nl.store(out[0], value=nl.multiply(t[1], 2.0))  # logical tile 1 is tile 0's memory
+    nl.store(out[1], value=nl.multiply(u, 3.0))  # so is another tensor at address 0
+    return out
+
+
+def test_tiles_at_one_address_share_their_values():
+    # The device semantics of direct allocation (NKI's simulator gives each tensor its own
+    # storage, so it is not the reference here).
+    from tilelens.core.simulation.nki import AddressedTiles
+
+    x = rng.standard_normal((128, 64)).astype(np.float32)
+    tilelens.clear()
+    got = tilelens.trace(client=Tracer(), frontend="nki")(shared_address_kernel)[(1,)](
+        x
+    )
+    np.testing.assert_allclose(got[0], 2 * x, rtol=1e-6)
+    np.testing.assert_allclose(got[1], 3 * x, rtol=1e-6)
+    from tilelens.core.simulation.nki import nki_builder
+
+    t = nki_builder.ndarray(
+        (4, nl.par_dim(128), 8),
+        np.float32,
+        buffer=ncc.sbuf.mod_alloc(base_addr=0, num_free_tiles=(2,)),
+    )
+    assert (
+        isinstance(t, AddressedTiles)
+        and t.shape == (4, 128, 8)
+        and t.data.shape == (2, 128, 8)
+    )
+
+
+def psum_reuse_kernel(a, b):
+    out = nl.ndarray((2, 128, b.shape[1]), dtype=nl.float32, buffer=nl.shared_hbm)
+    x = nl.load(a)
+    y = nl.load(b)
+    for i in nl.affine_range(2):
+        p = nl.ndarray(
+            (128, b.shape[1]), dtype=nl.float32, buffer=ncc.psum.mod_alloc(base_bank=3)
+        )
+        p[...] += nisa.nc_matmul(x, y)  # starts from zero, though the bank was used
+        nl.store(out[i], value=nl.copy(p))
+    return out
+
+
+def test_psum_tensors_start_from_zero():
+    a = rng.standard_normal((128, 128)).astype(np.float32)
+    b = rng.standard_normal((128, 512)).astype(np.float32)
+    tilelens.clear()
+    got = tilelens.trace(client=Tracer(), frontend="nki")(psum_reuse_kernel)[(1,)](a, b)
+    want = nki.simulate_kernel(nki.jit(psum_reuse_kernel)[(1,)], a, b)
+    np.testing.assert_allclose(got, want, rtol=1e-4, atol=1e-4)
diff --git a/tilelens/core/simulation/nki.py b/tilelens/core/simulation/nki.py
index b0312aab..7f47da98 100644
--- a/tilelens/core/simulation/nki.py
+++ b/tilelens/core/simulation/nki.py
@@ -114,6 +114,71 @@ def _shape(shape):
     return int(getattr(shape, "value", shape))
 
 
+class _ParDim(int):
+    """`nl.par_dim(n)`: the size n, marking the partition dimension of a tensor's shape."""
+
+    @property
+    def value(self):
+        return int(self)
+
+
+def _partition_axis(shape):
+    """The partition dimension of a shape: the one marked `par_dim`, else the first (NKI's
+    blocks precede the partition dimension)."""
+    if isinstance(shape, (tuple, list)):
+        for axis, d in enumerate(shape):
+            if isinstance(d, _ParDim) or type(d).__name__ == "ParDim":
+                return axis
+    return 0
+
+
+# The physical SBUF of a NeuronCore, for tensors placed by address (ncc.sbuf.mod_alloc /
+# alloc): bytes per partition. SBUF is 192 KiB per partition on Trn1/Trn2; the model allows
+# more (pages that are not touched cost nothing), since NKI's own simulator does not check
+# the capacity either. PSUM tensors start from zero instead (see Builder.ndarray).
+_SBUF_BYTES = 16 * 1024 * 1024
+_PARTITIONS = 128
+
+
+def _linearize(index, tiles):
+    """NKI's modulo allocation: each block index modulo its tile count, linearized in
+    row-major order over the tile counts (all ones when `tiles` is empty)."""
+    tiles = tuple(tiles) or (1,) * len(index)
+    linear = 0
+    for i, n in zip(index, tiles):
+        linear = linear * n + i % n
+    return linear
+
+
+def _block_address(alloc, index, pdim, fdim_bytes):
+    """(start partition, byte offset in a partition) of a logical tile under a direct SBUF
+    allocation."""
+    if type(alloc).__name__ == "SBUFModuloAllocFunc":
+        return (
+            alloc.base_partition + _linearize(index, alloc.num_par_tiles) * pdim,
+            alloc.base_addr + _linearize(index, alloc.num_free_tiles) * fdim_bytes,
+        )
+    partition, offset = alloc.func(tuple(index), pdim, fdim_bytes)  # SBUFAllocFunc
+    return partition, offset
+
+
+def _block_periods(alloc, blocks):
+    """Per block dimension, after how many indices the addresses repeat (the logical tiles
+    of a modulo allocation can outnumber its physical tiles); a custom allocation function
+    is taken not to repeat."""
+    if type(alloc).__name__ != "SBUFModuloAllocFunc":
+        return tuple(blocks)
+    counts = [alloc.num_par_tiles, alloc.num_free_tiles]
+    periods = []
+    for d, b in enumerate(blocks):
+        period = 1
+        for tiles in counts:
+            if tiles:
+                period = np.lcm(period, int(tiles[d]))
+        periods.append(min(int(period), int(b)))
+    return tuple(periods)
+
+
 def _partition_aligned(operand, ndim):
     """NKI broadcasts a 1-D tile, which has only the partition dimension, along the free
     dimensions of an N-D tile: as an array it is [P, 1, ...], not [P] aligned to the last axis."""
@@ -403,6 +468,44 @@ def broadcast_to(self, *args, **kwargs):
         )
 
 
+class AddressedTiles(NDArray):
+    """A tensor placed by address whose logical tiles outnumber its physical tiles (e.g. two
+    blocks on one physical tile): its data are the physical tiles, and block indices are
+    taken modulo their period, so logical tiles that share memory share values."""
+
+    def __init__(self, value, logical_shape, periods, **kwargs):
+        super().__init__(value=value, **kwargs)
+        self._logical_shape = tuple(logical_shape)
+        self._periods = tuple(periods)
+
+    @property
+    def shape(self):
+        return self._logical_shape
+
+    def _physical_keys(self, keys):
+        keys = list(keys if isinstance(keys, tuple) else (keys,))
+        for d, period in enumerate(self._periods[: len(keys)]):
+            k = _value(keys[d])
+            if isinstance(k, slice):
+                start, stop, step = k.indices(self._logical_shape[d])
+                if max(start, stop - 1) >= period:
+                    k = np.arange(start, stop, step) % period
+            elif k is not None and k is not Ellipsis:
+                k = (
+                    np.asarray(k) % period
+                    if isinstance(k, np.ndarray)
+                    else int(k) % period
+                )
+            keys[d] = k
+        return tuple(keys)
+
+    def __getitem__(self, keys):
+        return NDArray.__getitem__(self, self._physical_keys(keys))
+
+    def __setitem__(self, keys, value):
+        return NDArray.__setitem__(self, self._physical_keys(keys), value)
+
+
 class HBMAccess(NDArray):
     """`t[keys]` of an HBM tensor outside a load or store, as in `a = t[i, j]` followed by
     `nl.load(a, mask=m)`: in NKI an access pattern that the load transfers, applying its
@@ -480,6 +583,7 @@ def __init__(self, grid_dims=None):
         self.fn = None
         self.shared_hbm_arrays = {}
         self.rng = np.random.default_rng(0)  # nl.dropout; reseeded by nl.random_seed
+        self.memory = {}  # "sbuf" / "psum": physical memory for tensors placed by address
         self.accumulator = None  # the Vector Engine's per-partition reduce registers
 
     def set_grid_dim(self, *grid_dims):
@@ -490,7 +594,75 @@ def set_grid_idx(self, x, y, z):
         self.grid_y = y
         self.grid_z = z
 
+    def _physical_memory(self, kind):
+        if kind not in self.memory:
+            self.memory[kind] = np.zeros((_PARTITIONS, _SBUF_BYTES), dtype=np.uint8)
+        return self.memory[kind]
+
+    def _placed_ndarray(self, shape, dtype, buffer, name, kwargs):
+        """An SBUF tensor placed by address (`buffer=ncc.sbuf.mod_alloc(...)` or
+        `ncc.sbuf.alloc(...)`): a view of the NeuronCore's physical SBUF, so tensors (and
+        logical tiles) at the same address share their values, as on the device."""
+        alloc = buffer.allocation
+        p_axis = _partition_axis(shape)
+        dims = _shape(shape)
+        blocks, pdim, free = dims[:p_axis], dims[p_axis], dims[p_axis + 1 :]
+        dtype = np.dtype(dtype)
+        fdim_bytes = int(np.prod(free, dtype=np.int64)) * dtype.itemsize
+        memory = self._physical_memory("sbuf")
+        row = memory.strides[0]
+        periods = _block_periods(alloc, blocks)
+
+        def offset(index):
+            partition, byte = _block_address(alloc, index, pdim, fdim_bytes)
+            return partition * row + byte
+
+        origin = offset((0,) * len(blocks))
+        block_strides = []
+        for d in range(len(blocks)):
+            unit = [0] * len(blocks)
+            unit[d] = 1
+            block_strides.append(offset(unit) - origin if periods[d] > 1 else 0)
+        for index in np.ndindex(*periods):  # the layout must be strided to be a view
+            expected = origin + sum(i * st for i, st in zip(index, block_strides))
+            if offset(index) != expected:
+                raise NotImplementedError(
+                    "the interpreter supports direct allocations whose tiles are evenly spaced"
+                )
+        free_strides = []
+        step = dtype.itemsize
+        for f in reversed(free):
+            free_strides.insert(0, step)
+            step *= f
+        view_shape = tuple(periods) + (pdim,) + tuple(free)
+        strides = tuple(block_strides) + (row,) + tuple(free_strides)
+        last = origin + sum((n - 1) * st for n, st in zip(view_shape, strides) if n)
+        if origin < 0 or last + dtype.itemsize > memory.nbytes:
+            raise NotImplementedError(
+                "a direct allocation outside the on-chip memory the interpreter models"
+            )
+        view = np.ndarray(
+            view_shape, dtype=dtype, buffer=memory, offset=origin, strides=strides
+        )
+        value = kwargs.pop("value", None)
+        if value is not None:
+            view[...] = (
+                value
+                if tuple(periods) == tuple(blocks)
+                else value[tuple(slice(0, n) for n in periods)]
+            )
+        if tuple(periods) == tuple(blocks):
+            return NDArray(buffer=buffer, name=name, value=view, **kwargs)
+        return AddressedTiles(view, dims, periods, buffer=buffer, name=name, **kwargs)
+
     def ndarray(self, shape, dtype, *, buffer=None, name=None, **kwargs):
+        if getattr(buffer, "name", None) == "psum" and "value" not in kwargs:
+            # A PSUM tensor starts from zero, also when placed in a bank an earlier tensor
+            # used: its first `+=` of a matmul result is that result, as in NKI's simulator.
+            kwargs["value"] = np.zeros(_shape(shape), dtype=dtype)
+            buffer = nl.psum
+        elif getattr(buffer, "allocation", None) is not None:
+            return self._placed_ndarray(shape, dtype, buffer, name, kwargs)
         shape = _shape(shape)
         if buffer == nl.shared_hbm:
             if name is None:
@@ -521,8 +693,7 @@ def ndarray(self, shape, dtype, *, buffer=None, name=None, **kwargs):
         return ret
 
     def zeros(self, shape, dtype, *, buffer=None, name=None, **kwargs):
-        shape = _shape(shape)
-        value = np.zeros(shape, dtype=dtype)
+        value = np.zeros(_shape(shape), dtype=dtype)
         return self.ndarray(
             shape, dtype, buffer=buffer, name=name, value=value, **kwargs
         )
@@ -752,8 +923,8 @@ def static_cast(self, x, dtype):
         )
 
     def par_dim(self, value):
-        # Marks a partition dimension for the compiler; the size is all the simulator needs.
-        return value
+        # Marks the partition dimension; tensors placed by address need to know which it is.
+        return _ParDim(value)
 
     def square(self, x: NDArray, **kwargs):
         return self._unary_op(x, np.square, "square", **kwargs)
@@ -766,8 +937,7 @@ def multiply(self, x, y, **kwargs):
 
     # ---- nl: tensor creation, slicing, grid ------------------------------
     def full(self, shape, fill_value, dtype, *, buffer=None, name=None, **kwargs):
-        shape = _shape(shape)
-        value = np.full(shape, fill_value, dtype=dtype)
+        value = np.full(_shape(shape), fill_value, dtype=dtype)
         return self.ndarray(
             shape, dtype, buffer=buffer, name=name, value=value, **kwargs
         )
@@ -1505,6 +1675,7 @@ def run(self, *args, **kwargs):
         nki_builder.shared_hbm_arrays = {}
         nki_builder.rng = np.random.default_rng(0)
         nki_builder.accumulator = None
+        nki_builder.memory = {}
         nki_builder.fn = self.fn
 
         kwargs.pop("warmup", None)  # Remove warmup from kwargs if it exists

From 3430dd062c1777797d3caa9375c9ea484a289ba0 Mon Sep 17 00:00:00 2001
From: Hao Wu <hwu27@gmu.edu>
Date: Sun, 8 Nov 2026 20:48:05 -0500
Subject: [PATCH 5/6] [FEAT] Run the paged-attention kernel of the NKI samples:
 skip mode, dynamic loops, helpers

- Loads, stores, and DMA with oob_mode.skip: out-of-bounds lanes are not moved, and a tile
  assigned a skipping load keeps its values there (SkippedLanes). Clients see the access as
  performed (skipped_access: lanes masked off, keys clipped); in_bounds_lanes computes which
  lanes of t[keys] are in bounds from zero-stride views.
- `scalar(t)` (also imported by name from neuronxcc.nki.typing) gives a tile's value, e.g.
  the trip count of a dynamic loop.
- Plain helper functions that call NKI kernels through their own module's globals reach the
  rewrites: during a launch the kernels in user modules are swapped for their rewrites, and
  restored afterwards. nki_kernel_function checks an object's type before its attributes.
- `nl.store(t[k], v)` left unrewritten (in a helper) takes its second argument as the value.
- NKI tiles start with the partition dimension: an operand of fewer dimensions (a [P, 1]
  bias, a [P, A] predicate of a [P, A, B] tile) aligns from the left.
- A scalar index tile (a 0-d integer array) slices like an int, so writes through the slice
  reach the tile (numpy treated it as an advanced index and copied).
- NDArray.itemsize; tile_shape clips out-of-bounds int keys.

With these, flash_paged_attention_varlen runs and matches its test's reference (NKI's
simulator cannot run it).

Tests: skip-mode loads and a plain helper compared with nki.simulate_kernel; a dynamic trip
count; left-aligned operands; a scalar index keeping a view.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0169QRxKtsivkLDCS5f7rrsm
---
 tests/nki/test_nki_more_ops.py     | 105 +++++++++++++++
 tilelens/core/frontend/nki.py      |  24 +++-
 tilelens/core/masked_load_store.py |  49 ++++++-
 tilelens/core/simulation/nki.py    | 209 +++++++++++++++++++++++++----
 tilelens/core/trace.py             |  10 +-
 5 files changed, 355 insertions(+), 42 deletions(-)

diff --git a/tests/nki/test_nki_more_ops.py b/tests/nki/test_nki_more_ops.py
index ec4bce07..0a21e8bb 100644
--- a/tests/nki/test_nki_more_ops.py
+++ b/tests/nki/test_nki_more_ops.py
@@ -453,3 +453,108 @@ def test_psum_tensors_start_from_zero():
     got = tilelens.trace(client=Tracer(), frontend="nki")(psum_reuse_kernel)[(1,)](a, b)
     want = nki.simulate_kernel(nki.jit(psum_reuse_kernel)[(1,)], a, b)
     np.testing.assert_allclose(got, want, rtol=1e-4, atol=1e-4)
+
+
+from neuronxcc.nki.isa.constants import oob_mode  # noqa: E402
+from neuronxcc.nki.typing import scalar  # noqa: E402 - kernels import it by name
+
+
+@nki.jit(mode="trace")
+def _scale_rows(t, factor):
+    return nl.multiply(t, factor)
+
+
+def _plain_helper(x_hbm, rows):
+    # not an NKI kernel: calls one through this module's globals, and stores positionally
+    return _scale_rows(nl.load(x_hbm[rows, :]), 2.0)
+
+
+def skip_and_helper_kernel(x, blocks, idx):
+    # loads with oob_mode.skip at a loaded index, one past the tensor (the tile keeps its
+    # values there); a helper that is not an NKI kernel, calling one through this module's
+    # globals, whose result is stored positionally
+    out = nl.ndarray((128, x.shape[1]), dtype=x.dtype, buffer=nl.shared_hbm)
+    picked = nl.ndarray(
+        (2, 128, blocks.shape[2]), dtype=blocks.dtype, buffer=nl.shared_hbm
+    )
+    nl.store(out, _plain_helper(x, nl.ds(128, 128)))
+    i = nl.load(idx)
+    tile = nl.full((128, blocks.shape[2]), -1.0, dtype=blocks.dtype)
+    tile[...] = nl.load(blocks[i[0, 0]], mode=oob_mode.skip)  # in bounds: copied
+    nl.store(picked[0], value=tile)
+    tile[...] = nl.load(blocks[i[0, 1]], mode=oob_mode.skip)  # past the end: kept
+    nl.store(picked[1], value=tile)
+    return out, picked
+
+
+def test_skip_mode_and_plain_helpers_match_nki_simulator():
+    x = rng.standard_normal((256, 32)).astype(np.float32)
+    blocks = rng.standard_normal((4, 128, 16)).astype(np.float32)
+    idx = np.array([[2, 9]], dtype=np.int32)
+    tilelens.clear()
+    got = tilelens.trace(client=Tracer(), frontend="nki")(skip_and_helper_kernel)[(1,)](
+        x, blocks, idx
+    )
+    want = nki.simulate_kernel(nki.jit(skip_and_helper_kernel)[(1,)], x, blocks, idx)
+    for g, w in zip(got, want):
+        np.testing.assert_allclose(g, w, rtol=1e-5, atol=1e-5)
+    np.testing.assert_array_equal(got[1][0], blocks[2])
+    np.testing.assert_array_equal(got[1][1], blocks[2])
+
+
+def dynamic_loop_kernel(x, steps):
+    # range(scalar(t)): a trip count loaded from HBM (a dynamic loop; NKI's CPU simulator
+    # does not run these, so the expected values are written out)
+    out = nl.ndarray(x.shape, dtype=x.dtype, buffer=nl.shared_hbm)
+    acc = nl.zeros(x.shape, dtype=nl.float32, buffer=nl.sbuf)
+    for _ in range(scalar(nl.load(steps))):
+        acc[...] = nl.add(acc, nl.load(x))
+    nl.store(out, value=acc)
+    return out
+
+
+def test_dynamic_trip_count():
+    x = rng.standard_normal((128, 32)).astype(np.float32)
+    tilelens.clear()
+    got = tilelens.trace(client=Tracer(), frontend="nki")(dynamic_loop_kernel)[(1,)](
+        x, np.array([[3]], dtype=np.int32)
+    )
+    np.testing.assert_allclose(got, 3 * x, rtol=1e-6)
+
+
+def left_aligned_kernel(x, pred):
+    # a [P, 1] bias of a [P, A, B] activation, and a [P, A] predicate of a [P, A, B] copy
+    out = nl.ndarray(x.shape, dtype=x.dtype, buffer=nl.shared_hbm)
+    t = nl.load(x)
+    bias = nl.full((x.shape[0], 1), 0.5, dtype=nl.float32)
+    a = nisa.activation(op=nl.copy, data=t, bias=bias, scale=1.0)
+    dst = nl.full(x.shape, -1.0, dtype=x.dtype)
+    nisa.tensor_copy_predicated(src=a, dst=dst, predicate=nl.load(pred))
+    nl.store(out, value=dst)
+    return out
+
+
+def test_lower_rank_operands_align_from_the_partition_dimension():
+    x = rng.standard_normal((128, 4, 8)).astype(np.float32)
+    pred = (rng.random((128, 4)) < 0.5).astype(np.uint8)
+    tilelens.clear()
+    got = tilelens.trace(client=Tracer(), frontend="nki")(left_aligned_kernel)[(1,)](
+        x, pred
+    )
+    want = np.where(pred[:, :, None].astype(bool), x + 0.5, -1.0)
+    np.testing.assert_allclose(got, want, rtol=1e-6)
+
+
+def test_scalar_index_tiles_keep_slices_as_views():
+    # A 0-d index (a loaded value, or a loop index the sanitizer tracks) must slice like an
+    # int: a write through the slice reaches the tile.
+    from tilelens.core.masked_load_store import tile_shape
+    from tilelens.core.simulation.nki import NDArray
+
+    t = NDArray(value=np.zeros((4, 3), dtype=np.float32))
+    row = t[NDArray(value=np.array(2))]
+    row[...] = 5.0
+    assert np.all(t.data[2] == 5.0) and np.all(t.data[[0, 1, 3]] == 0.0)
+    assert tile_shape(np.zeros((4, 3)), (7, slice(None))) == (
+        3,
+    )  # out-of-bounds int key
diff --git a/tilelens/core/frontend/nki.py b/tilelens/core/frontend/nki.py
index 50dd70eb..21b8a957 100644
--- a/tilelens/core/frontend/nki.py
+++ b/tilelens/core/frontend/nki.py
@@ -19,7 +19,12 @@
 HAS_NKI = False
 nki_builder = None
 try:
-    from tilelens.core.simulation.nki import NDArray, access, nki_builder  # type: ignore
+    from tilelens.core.simulation.nki import (  # type: ignore
+        NDArray,
+        access,
+        nki_builder,
+        skipped_access,
+    )
 
     HAS_NKI = True
 except ModuleNotFoundError:
@@ -33,10 +38,17 @@ def _nki_dot_adapter(x: Any, y: Any, *_args: Any, **_kwargs: Any) -> AdapterResu
     return AdapterResult(x, y)
 
 
-def _nki_access_adapter(tensor: Any, keys: Any, mask: Any) -> AdapterResult:
+def _nki_access_adapter(
+    tensor: Any, keys: Any, mask: Any, kwargs: dict | None = None
+) -> AdapterResult:
     # `nl.load(t)` and `nl.store(t, value=v)` access the whole tensor, and
-    # `a = t[k]; nl.load(a)` accesses t[k].
+    # `a = t[k]; nl.load(a)` accesses t[k]. With `mode=oob_mode.skip`, clients see the
+    # access as it is performed: out-of-bounds lanes masked off, keys clipped.
     tensor, keys = access(tensor, keys)
+    kwargs = kwargs or {}
+    mode = kwargs.get("oob_mode", kwargs.get("mode"))
+    if getattr(mode, "name", None) == "skip":
+        keys, mask, _ = skipped_access(tensor, keys, mask)
     return AdapterResult(tensor, mask, keys)
 
 
@@ -70,14 +82,16 @@ def _nki_reduce_sum_adapter(
         ProgramId: lambda axis, *_args, **_kwargs: AdapterResult(axis),
         Allocate: lambda *_args, **_kwargs: AdapterResult(),
         Load: lambda src, keys=None, *, mask=None, **_kwargs: _nki_access_adapter(
-            src, keys, mask
+            src, keys, mask, _kwargs
         ),
         Store: lambda dst,
         keys=None,
         value=None,
         *,
         mask=None,
-        **_kwargs: _nki_access_adapter(dst, keys, mask),
+        **_kwargs: _nki_access_adapter(
+            dst, None if value is None else keys, mask, _kwargs
+        ),
         Dot: _nki_dot_adapter,
         ReduceSum: _nki_reduce_sum_adapter,
     }
diff --git a/tilelens/core/masked_load_store.py b/tilelens/core/masked_load_store.py
index 5d089721..5a027526 100644
--- a/tilelens/core/masked_load_store.py
+++ b/tilelens/core/masked_load_store.py
@@ -96,21 +96,58 @@ def tile_shape(ndarray: np.ndarray, keys) -> tuple:
     """Shape of ``ndarray[keys]``, also when some integer-array keys are out of bounds."""
     if not isinstance(keys, tuple):
         keys = (keys,)
-    clipped = []
-    for dim, key in enumerate(keys):
-        if (
-            isinstance(key, np.ndarray)
-            and key.dtype.kind in "iu"
-            and dim < ndarray.ndim
+    clipped: list = []
+    dim = 0
+    for key in keys:
+        if key is None:
+            clipped.append(key)
+            continue
+        if dim < ndarray.ndim and (
+            isinstance(key, (int, np.integer))
+            or (isinstance(key, np.ndarray) and key.dtype.kind in "iu")
         ):
             key = np.clip(key, 0, max(ndarray.shape[dim] - 1, 0))
+            key = key if isinstance(key, np.ndarray) and key.ndim else int(key)
         clipped.append(key)
+        dim += 1
     # index a one-byte, zero-stride view: only the shape is needed, not the values
     return np.broadcast_to(np.zeros((), dtype=bool), ndarray.shape)[
         tuple(clipped)
     ].shape
 
 
+def in_bounds_lanes(ndarray: np.ndarray, keys) -> np.ndarray:
+    """Which lanes of ``ndarray[keys]`` index inside the array (a boolean array of the tile's
+    shape), also when some keys are out of bounds. Each dimension's validity is a
+    zero-stride view indexed with the keys, so no array of the tensor's size is built."""
+    keys = keys if isinstance(keys, tuple) else (keys,)
+    norm, singleton_dims = normalize_slice(ndarray, keys)
+    target = list(_calculate_target_shape(norm, ndarray.shape))
+    fixed = []
+    for dim, key in enumerate(norm):
+        if isinstance(key, (int, np.integer)) or (
+            isinstance(key, np.ndarray) and key.dtype.kind in "iu"
+        ):
+            value = np.asarray(key)
+            if (
+                value.size and value.min() < 0
+            ):  # negative indices: one more, invalid index
+                target[dim] += 1
+                value = np.where(value < 0, target[dim] - 1, value)
+            fixed.append(value if isinstance(key, np.ndarray) else int(value))
+        else:
+            fixed.append(key)
+    lanes = None
+    for dim, size in enumerate(ndarray.shape):
+        valid = (np.arange(target[dim]) < size).reshape(
+            [-1 if d == dim else 1 for d in range(len(target))]
+        )
+        part = np.broadcast_to(valid, target)[tuple(fixed)]
+        lanes = part if lanes is None else lanes & part
+    lanes = np.ones((), dtype=bool) if lanes is None else lanes
+    return np.expand_dims(lanes, singleton_dims)
+
+
 def tile_mask(ndarray: np.ndarray, keys, mask) -> np.ndarray | None:
     """``mask`` broadcast to the shape of ``ndarray[keys]``, as NKI broadcasts masks;
     None stays None."""
diff --git a/tilelens/core/simulation/nki.py b/tilelens/core/simulation/nki.py
index 7f47da98..94a7bf48 100644
--- a/tilelens/core/simulation/nki.py
+++ b/tilelens/core/simulation/nki.py
@@ -5,6 +5,7 @@
     import neuronxcc.nki.language as nl
     import neuronxcc.nki.isa as nisa
     import neuronxcc.nki.compiler as ncc
+    import neuronxcc.nki.typing as nki_typing
 except (
     ModuleNotFoundError
 ) as exc:  # pragma: no cover - only hit when optional deps missing
@@ -13,9 +14,16 @@
     ) from exc
 import inspect
 import os
+import sys
 import textwrap
 from ..frontend.nki_transform import transform_code
-from ..masked_load_store import masked_load, masked_store, tile_mask, tile_shape
+from ..masked_load_store import (
+    in_bounds_lanes,
+    masked_load,
+    masked_store,
+    tile_mask,
+    tile_shape,
+)
 
 
 _PACKAGE_DIR = (
@@ -33,6 +41,17 @@ def _is_hbm(buffer):
     return getattr(buffer, "name", None) in ("hbm", "shared_hbm", "private_hbm")
 
 
+def _index_key(key):
+    """A subscript key as numpy takes it. A scalar index tile (a 0-d integer array, e.g. a
+    value loaded from SBUF or a loop index the sanitizer tracks) becomes an int: numpy would
+    treat a 0-d array as an advanced index and return a copy, so writes through the slice
+    would be lost; in NKI the slice is an access pattern of the same memory."""
+    key = key.data if isinstance(key, NDArray) else key
+    if isinstance(key, np.ndarray) and key.ndim == 0 and key.dtype.kind in "iu":
+        return int(key)
+    return key
+
+
 def _tile_keys(data, keys):
     """NKI treats a 1-D tile (a partition vector) as [P, 1]: `t[:, 0]` indexes it too
     (`t[:, None]` still adds an axis)."""
@@ -180,11 +199,13 @@ def _block_periods(alloc, blocks):
 
 
 def _partition_aligned(operand, ndim):
-    """NKI broadcasts a 1-D tile, which has only the partition dimension, along the free
-    dimensions of an N-D tile: as an array it is [P, 1, ...], not [P] aligned to the last axis."""
+    """NKI tiles start with the partition dimension, so an operand of fewer dimensions aligns
+    from the left and broadcasts along the trailing ones: a 1-D tile (a partition vector) or a
+    [P, 1] tile with an N-D tile is [P, 1, ...], not aligned to the last axis as in numpy."""
     value = _value(operand)
-    if np.ndim(value) == 1 and ndim >= 2:
-        return np.reshape(value, (np.shape(value)[0],) + (1,) * (ndim - 1))
+    rank = np.ndim(value)
+    if 1 <= rank < ndim:
+        return np.reshape(value, np.shape(value) + (1,) * (ndim - rank))
     return value
 
 
@@ -232,6 +253,18 @@ def _bool_tile(pred, shape, reverse=False):
     return ~value if reverse else value
 
 
+def _left_aligned(operand, shape):
+    """A predicate as a boolean array of `shape`: one of as many elements per partition is
+    taken as it is; one of fewer dimensions aligns from the partition dimension and
+    broadcasts along the trailing ones (a [P, A] predicate of a [P, A, B] tile)."""
+    value = np.asarray(operand)
+    if value.size == int(np.prod(shape)):
+        return value.reshape(shape).astype(bool)
+    if value.ndim < len(shape):
+        value = value.reshape(value.shape + (1,) * (len(shape) - value.ndim))
+    return np.broadcast_to(value.astype(bool), shape)
+
+
 def _per_partition(value):
     """A [P, ...] array as [P, N]: ISA instructions match operands by partition and
     elements per partition, not by the shape of the free dimensions."""
@@ -281,6 +314,10 @@ def shape(self):
     def dtype(self):
         return self.data.dtype if self.data is not None else None
 
+    @property
+    def itemsize(self):
+        return self.data.itemsize
+
     def data_ptr(self):
         if self._data_ptr is None:
             self._data_ptr = self.data.ctypes.data
@@ -335,9 +372,7 @@ def __getitem__(self, keys):
             return HBMAccess(self, tuple(_value(k) for k in keys))
 
         # Apply the slicing to the underlying numpy array
-        new_keys = _tile_keys(
-            self.data, [k.data if isinstance(k, NDArray) else k for k in keys]
-        )
+        new_keys = _tile_keys(self.data, [_index_key(k) for k in keys])
         sliced_value = self.data[tuple(new_keys)]
 
         # Create a new NDArray with the sliced data
@@ -353,11 +388,10 @@ def __setitem__(self, keys, value):
             keys = (keys,)
 
         # Apply the slicing to the underlying numpy array
-        new_keys = _tile_keys(
-            self.data, tuple(k.data if isinstance(k, NDArray) else k for k in keys)
-        )
+        new_keys = _tile_keys(self.data, tuple(_index_key(k) for k in keys))
         target_shape = self.data[new_keys].shape
         reduce = value if isinstance(value, LoopReduced) else None
+        skipped = value.valid if isinstance(value, SkippedLanes) else None
         value = np.asarray(value.data if isinstance(value, NDArray) else value)
         if value.shape != target_shape and value.size == int(np.prod(target_shape)):
             value = value.reshape(target_shape)  # e.g. a [P, 1] result into a [P] slot
@@ -367,6 +401,8 @@ def __setitem__(self, keys, value):
             value = reduce.op(current, np.broadcast_to(value, target_shape))
             if reduce.mask is not None:
                 value = np.where(reduce.mask.reshape(target_shape), value, current)
+        if skipped is not None:  # lanes a skipping load did not read keep their values
+            value = np.where(skipped.reshape(target_shape), value, self.data[new_keys])
         self.data[new_keys] = value
 
         return self
@@ -529,6 +565,57 @@ def data(self, value):
         self._value = value
 
 
+class SkippedLanes(NDArray):
+    """The result of a load with `mode=oob_mode.skip`: its out-of-bounds lanes are not read.
+    Assigned into a tile, it leaves those lanes of the tile as they were, as the DMA that
+    writes the tile skips them."""
+
+    def __init__(self, value, valid, **kwargs):
+        super().__init__(value=value, **kwargs)
+        self.valid = np.broadcast_to(np.asarray(valid, dtype=bool), value.shape)
+
+
+def _clip_keys(keys, shape):
+    """Integer and integer-array keys clipped into the array, for an access whose
+    out-of-bounds lanes are masked off (oob_mode.skip): those lanes are neither read nor
+    written, and the clipped keys keep the access in bounds."""
+    keys = keys if isinstance(keys, tuple) else (keys,)
+    clipped, dim = [], 0
+    for key in keys:
+        if key is None:
+            clipped.append(key)
+            continue
+        if dim < len(shape) and (
+            isinstance(key, (int, np.integer))
+            or (isinstance(key, np.ndarray) and key.dtype.kind in "iu")
+        ):
+            value = np.clip(key, 0, max(shape[dim] - 1, 0))
+            key = value if isinstance(key, np.ndarray) else int(value)
+        clipped.append(key)
+        dim += 1
+    return tuple(clipped)
+
+
+def skipped_access(tensor, keys, mask):
+    """An access with `mode=oob_mode.skip` as it is performed: its out-of-bounds lanes
+    masked off and its keys clipped into the tensor. Returns (numpy keys, mask, lanes in
+    bounds)."""
+    numpy_keys = _expand_oob_slices(
+        nki_builder._convert_keys_to_numpy(keys), tensor.data.shape
+    )
+    valid = in_bounds_lanes(tensor.data, numpy_keys)
+    mask_value = tile_mask(tensor.data, numpy_keys, _value(mask))
+    mask_value = valid if mask_value is None else mask_value & valid
+    return _clip_keys(numpy_keys, tensor.data.shape), mask_value, valid
+
+
+def _skips_oob(kwargs):
+    """Whether a load/store/DMA has `mode=oob_mode.skip` (or `oob_mode=`); removes it."""
+    mode = kwargs.pop("mode", None)
+    mode = kwargs.pop("oob_mode", mode)
+    return getattr(mode, "name", None) == "skip"
+
+
 class LoopReduced(NDArray):
     """The result of `nl.loop_reduce(x, op, ...)`. Assigned into a tensor slice, it combines
     with the slice's current value: `dst[...] = op(dst[...], x)`, as in NKI's simulator.
@@ -728,16 +815,23 @@ def load(self, src: NDArray, keys=None, *, mask=None, **kwargs):
         ndarray = src.data
         mask_value = _value(mask)
         dtype = kwargs.pop("dtype", None)
+        skip = _skips_oob(kwargs)
 
-        # Convert any NDArrays in keys to numpy arrays
-        numpy_keys = self._convert_keys_to_numpy(keys)
-
-        # Call the actual masked_load function
-        numpy_keys = _expand_oob_slices(numpy_keys, ndarray.shape)
-        mask_value = tile_mask(ndarray, numpy_keys, mask_value)
+        valid = None
+        if skip:  # out-of-bounds lanes are not read (oob_mode.skip)
+            numpy_keys, mask_value, valid = skipped_access(src, keys, mask)
+        else:
+            # Convert any NDArrays in keys to numpy arrays
+            numpy_keys = self._convert_keys_to_numpy(keys)
+            numpy_keys = _expand_oob_slices(numpy_keys, ndarray.shape)
+            mask_value = tile_mask(ndarray, numpy_keys, mask_value)
         result = masked_load(ndarray, numpy_keys, mask=mask_value)
 
         # Convert result back to NDArray
+        if valid is not None:
+            return SkippedLanes(
+                _cast(result, dtype), valid, name=f"{src.name}_masked_load", **kwargs
+            )
         return NDArray(
             value=_cast(result, dtype), name=f"{src.name}_masked_load", **kwargs
         )
@@ -752,9 +846,14 @@ def store(
         **kwargs,
     ):
         """Store array elements with masking for out-of-bounds errors."""
+        if (
+            value is None
+        ):  # `nl.store(t[k], v)` not rewritten: the value is the second argument
+            keys, value = None, keys
         dst, keys = access(dst, keys)
         ndarray = dst.data
         mask_value = _value(mask)
+        skip = _skips_oob(kwargs)
 
         # Convert any NDArrays in keys to numpy arrays
         numpy_keys = self._convert_keys_to_numpy(keys)
@@ -769,6 +868,10 @@ def store(
             value_array = value_array.reshape(tile_shape_)
         value_array = np.broadcast_to(value_array, tile_shape_).astype(ndarray.dtype)
         mask_value = tile_mask(ndarray, numpy_keys, mask_value)
+        if skip:  # out-of-bounds lanes are not written (oob_mode.skip)
+            valid = in_bounds_lanes(ndarray, numpy_keys)
+            mask_value = valid if mask_value is None else mask_value & valid
+            numpy_keys = _clip_keys(numpy_keys, ndarray.shape)
 
         # Call the actual masked_store function
         masked_store(ndarray, numpy_keys, value_array, mask=mask_value)
@@ -1178,15 +1281,16 @@ def dma_copy(
         src_keys=None,
         mask=None,
         dst_rmw_op=None,
+        oob_mode=None,
         **kwargs,
     ):
         # A DMA transfer reads `src[src_keys]` and writes `dst[dst_keys]`; both go through
         # load/store, so analysis clients see them as a load and a store.
-        value = self.load(src, src_keys, mask=mask)
+        value = self.load(src, src_keys, mask=mask, oob_mode=oob_mode)
         if dst_rmw_op is not None:
-            current = self.load(dst, dst_keys, mask=mask)
+            current = self.load(dst, dst_keys, mask=mask, oob_mode=oob_mode)
             value = NDArray(value=_ufunc(dst_rmw_op)(current.data, value.data))
-        self.store(dst, dst_keys, value, mask=mask)
+        self.store(dst, dst_keys, value, mask=mask, oob_mode=oob_mode)
 
     def scalar_tensor_tensor(
         self,
@@ -1412,9 +1516,7 @@ def tensor_copy_predicated(
     ):
         # dst = where(predicate, src, dst), matching operands by partition.
         current = _per_partition(dst.data)
-        pred = _per_partition(
-            _bool_tile(_value(predicate), np.shape(_value(predicate)))
-        )
+        pred = _per_partition(_left_aligned(_value(predicate), dst.data.shape))
         pred = np.broadcast_to(pred, current.shape)
         if reverse_pred:
             pred = ~pred
@@ -1454,6 +1556,12 @@ def range(self, *args, **kwargs):
         # affine_range, sequential_range, and static_range all iterate like range()
         return range(*args)
 
+    def scalar(self, value, *, dtype=None):
+        # A tile's single value as a scalar, e.g. the trip count of a dynamic loop
+        # (`for _ in range(scalar(t))`).
+        item = np.asarray(_value(value)).reshape(-1)[0]
+        return _cast(item, dtype).item()
+
     def multi_buffer(self, *args, **kwargs):
         # a scheduling directive for a range (`directives=`), which ranges ignore
         return None
@@ -1539,6 +1647,13 @@ def _set_attr(obj, name, value):
         _set_attr(nisa, name, getattr(nki_builder, name))
     _set_attr(nisa, "reciprocal", nki_builder.nisa_reciprocal)
     _set_attr(ncc, "multi_buffer", nki_builder.multi_buffer)
+    # `scalar` is usually imported by name (from neuronxcc.nki.typing import scalar): patch the
+    # modules that hold it too
+    real_scalar = nki_typing.scalar
+    for module in list(sys.modules.values()):
+        namespace = getattr(module, "__dict__", None)
+        if isinstance(namespace, dict) and namespace.get("scalar") is real_scalar:
+            _set_attr(module, "scalar", nki_builder.scalar)
     _set_attr(nisa, "_tiled_offloaded_memcpy", nki_builder.tiled_offloaded_memcpy)
     _set_attr(nisa, "_tiled_offloaded_fma", nki_builder.tiled_offloaded_fma)
     for name in (
@@ -1625,10 +1740,41 @@ def _rewritten_namespace(module_globals, namespaces):
     return namespace
 
 
-def _rewrite_function(fn):
+def _rewrite_function(fn, namespaces):
     """`fn` rewritten by the NKI frontend AST transformer, in the rewritten namespace of its
     module, so the NKI kernels it calls by name are rewritten too."""
-    return _compile_rewritten(fn, _rewritten_namespace(fn.__globals__, {}))
+    return _compile_rewritten(fn, _rewritten_namespace(fn.__globals__, namespaces))
+
+
+_LIBRARY_PATHS = ("site-packages", "dist-packages")
+
+
+def _swap_in_rewritten_kernels(namespaces):
+    """Replace, in user modules, the NKI kernel objects that plain helper functions call
+    through their module's globals with their rewrites; return what to restore."""
+    from ..trace import nki_kernel_function
+
+    swapped = []
+    for module in list(sys.modules.values()):
+        globals_ = getattr(module, "__dict__", None)
+        path = getattr(module, "__file__", None) or ""
+        if (
+            not isinstance(globals_, dict)
+            or not path
+            or any(p in path for p in _LIBRARY_PATHS)
+        ):
+            continue
+        names = [
+            n for n, v in list(globals_.items()) if nki_kernel_function(v) is not None
+        ]
+        if not names:
+            continue
+        rewritten = _rewritten_namespace(globals_, namespaces)
+        for name in names:
+            if rewritten.get(name) is not globals_[name]:
+                swapped.append((globals_, name, globals_[name]))
+                globals_[name] = rewritten[name]
+    return swapped
 
 
 def full_keys(tensor):
@@ -1693,8 +1839,9 @@ def run(self, *args, **kwargs):
             from ..trace import nki_kernel_function
 
             fn = nki_kernel_function(self.fn) or self.fn
+            self._namespaces = {}
             if hasattr(fn, "__code__"):
-                self.fn = _rewrite_function(fn)
+                self.fn = _rewrite_function(fn, self._namespaces)
             self._rewritten = True
 
         # convert args to NDArray if they are not already, also those passed by keyword
@@ -1716,6 +1863,16 @@ def run(self, *args, **kwargs):
             ret = arg
             client_manager.arg_callback(name, arg, ret)
 
+        # Plain helper functions call NKI kernels through their own module's globals: run
+        # with the rewrites in place there, and restore the kernels afterwards.
+        swapped = _swap_in_rewritten_kernels(getattr(self, "_namespaces", {}))
+        try:
+            return self._run_grid(grid_dims, args, kwargs, client_manager)
+        finally:
+            for globals_, name, kernel in reversed(swapped):
+                globals_[name] = kernel
+
+    def _run_grid(self, grid_dims, args, kwargs, client_manager):
         ret = None
         for x in range(grid_dims[0]):
             for y in range(grid_dims[1]):
diff --git a/tilelens/core/trace.py b/tilelens/core/trace.py
index 32199c71..435f9ca4 100644
--- a/tilelens/core/trace.py
+++ b/tilelens/core/trace.py
@@ -366,12 +366,12 @@ def warmup(self, *args, **kwargs):
 def nki_kernel_function(obj):
     """The Python function of an NKI kernel object (`@nki.jit`, `@nki.baremetal`, ... of the
     legacy `neuronxcc.nki` or of the newer `nki` package), or None for anything else."""
+    if type(obj).__module__.split(".")[0] not in ("neuronxcc", "nki"):
+        return (
+            None  # (checked first: some objects, e.g. proxies, fail on any attribute)
+        )
     func = getattr(obj, "func", None)
-    if type(obj).__module__.split(".")[0] in ("neuronxcc", "nki") and isinstance(
-        func, types.FunctionType
-    ):
-        return func
-    return None
+    return func if isinstance(func, types.FunctionType) else None
 
 
 def trace_source(kernel):

From a036fa02d9a12f8c63bd46ef068b05730c9dd847 Mon Sep 17 00:00:00 2001
From: Hao Wu <hwu27@gmu.edu>
Date: Sun, 8 Nov 2026 21:28:39 -0500
Subject: [PATCH 6/6] [PERF] Cache the module scans of NKI launches

Finding the user modules whose kernels to swap for their rewrites, and the modules that
import `scalar` by name, scanned every loaded module on every launch: about 9 ms with torch
loaded (2,700 modules). The scans are now cached until more modules are loaded; at launch only
the user modules (outside Python's library and the installed packages) are checked for
kernels, so kernels defined later are still found.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0169QRxKtsivkLDCS5f7rrsm
---
 tilelens/core/simulation/nki.py | 53 +++++++++++++++++++++++++--------
 1 file changed, 41 insertions(+), 12 deletions(-)

diff --git a/tilelens/core/simulation/nki.py b/tilelens/core/simulation/nki.py
index 94a7bf48..9d92152d 100644
--- a/tilelens/core/simulation/nki.py
+++ b/tilelens/core/simulation/nki.py
@@ -1650,9 +1650,11 @@ def _set_attr(obj, name, value):
     # `scalar` is usually imported by name (from neuronxcc.nki.typing import scalar): patch the
     # modules that hold it too
     real_scalar = nki_typing.scalar
-    for module in list(sys.modules.values()):
-        namespace = getattr(module, "__dict__", None)
-        if isinstance(namespace, dict) and namespace.get("scalar") is real_scalar:
+    for namespace in _modules_with(
+        "scalar", lambda module, globals_: globals_.get("scalar") is real_scalar
+    ):
+        module = sys.modules.get(namespace.get("__name__"))
+        if module is not None and namespace.get("scalar") is real_scalar:
             _set_attr(module, "scalar", nki_builder.scalar)
     _set_attr(nisa, "_tiled_offloaded_memcpy", nki_builder.tiled_offloaded_memcpy)
     _set_attr(nisa, "_tiled_offloaded_fma", nki_builder.tiled_offloaded_fma)
@@ -1746,7 +1748,42 @@ def _rewrite_function(fn, namespaces):
     return _compile_rewritten(fn, _rewritten_namespace(fn.__globals__, namespaces))
 
 
+def _library_paths():
+    import sysconfig
+
+    paths = sysconfig.get_paths()
+    roots = {paths.get(k) for k in ("stdlib", "platstdlib", "purelib", "platlib")}
+    return tuple(sorted(r for r in roots if r)) + (_PACKAGE_DIR,)
+
+
 _LIBRARY_PATHS = ("site-packages", "dist-packages")
+_module_scans: dict = {}  # what -> (len(sys.modules), [module globals]), rescanned on growth
+
+
+def _modules_with(what, has):
+    """Globals of the loaded modules for which `has(module, globals)` holds, cached until
+    more modules are loaded (scanning thousands of modules on every launch costs ms)."""
+    count = len(sys.modules)
+    cached = _module_scans.get(what)
+    if cached is None or cached[0] != count:
+        found = []
+        for module in list(sys.modules.values()):
+            globals_ = getattr(module, "__dict__", None)
+            if isinstance(globals_, dict) and has(module, globals_):
+                found.append(globals_)
+        cached = _module_scans[what] = (count, found)
+    return cached[1]
+
+
+def _user_module(module, globals_):
+    """A module of the user's code (a file outside Python's library and the installed
+    packages, TileLens's included)."""
+    path = getattr(module, "__file__", None) or ""
+    return (
+        bool(path)
+        and not any(p in path for p in _LIBRARY_PATHS)
+        and not path.startswith(_library_paths())
+    )
 
 
 def _swap_in_rewritten_kernels(namespaces):
@@ -1755,15 +1792,7 @@ def _swap_in_rewritten_kernels(namespaces):
     from ..trace import nki_kernel_function
 
     swapped = []
-    for module in list(sys.modules.values()):
-        globals_ = getattr(module, "__dict__", None)
-        path = getattr(module, "__file__", None) or ""
-        if (
-            not isinstance(globals_, dict)
-            or not path
-            or any(p in path for p in _LIBRARY_PATHS)
-        ):
-            continue
+    for globals_ in _modules_with("user", _user_module):
         names = [
             n for n, v in list(globals_.items()) if nki_kernel_function(v) is not None
         ]