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/tests/nki/test_nki_more_ops.py b/tests/nki/test_nki_more_ops.py
new file mode 100644
index 00000000..0a21e8bb
--- /dev/null
+++ b/tests/nki/test_nki_more_ops.py
@@ -0,0 +1,560 @@
+"""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:])
+
+
+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)
+
+
+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)
+
+
+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/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..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, full_keys, nki_builder  # type: ignore
+    from tilelens.core.simulation.nki import (  # type: ignore
+        NDArray,
+        access,
+        nki_builder,
+        skipped_access,
+    )
 
     HAS_NKI = True
 except ModuleNotFoundError:
@@ -33,6 +38,20 @@ 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, 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]. 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)
+
+
 def _nki_reduce_sum_adapter(
     input_tensor: Any, *args: Any, mask: Any = None, **kwargs: Any
 ) -> AdapterResult:
@@ -62,21 +81,16 @@ 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, _kwargs
         ),
         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, None if value is None else keys, mask, _kwargs
         ),
         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..5a027526 100644
--- a/tilelens/core/masked_load_store.py
+++ b/tilelens/core/masked_load_store.py
@@ -67,20 +67,85 @@ 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):
         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)
-    return ndarray[tuple(clipped)].shape
+        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:
@@ -107,19 +172,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 +218,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 +254,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..9d92152d 100644
--- a/tilelens/core/simulation/nki.py
+++ b/tilelens/core/simulation/nki.py
@@ -4,6 +4,8 @@
 try:
     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
@@ -12,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 = (
@@ -24,10 +33,37 @@
 
 
 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 _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)."""
+    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
@@ -97,12 +133,79 @@ 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."""
+    """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
 
 
@@ -142,6 +245,49 @@ 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 _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."""
+    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
@@ -168,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
@@ -209,23 +359,50 @@ 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, [_index_key(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(_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
+        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)
+        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
@@ -301,9 +478,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 +504,133 @@ 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
+    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 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.
+    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 +669,9 @@ 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.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):
         self.grid_dims = grid_dims
@@ -361,7 +681,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:
@@ -392,8 +780,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
         )
@@ -424,33 +811,49 @@ 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
+        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
         )
 
     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
+        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 = 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)
+        skip = _skips_oob(kwargs)
 
         # Convert any NDArrays in keys to numpy arrays
         numpy_keys = self._convert_keys_to_numpy(keys)
@@ -465,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)
@@ -474,8 +881,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)
@@ -621,8 +1026,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)
@@ -635,8 +1040,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
         )
@@ -702,17 +1106,24 @@ 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):
-        # 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")
 
@@ -735,12 +1146,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 +1203,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)
@@ -867,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,
@@ -956,13 +1371,201 @@ 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(_left_aligned(_value(predicate), dst.data.shape))
+        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 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
+
 
 nki_builder = Builder()
 
@@ -1009,6 +1612,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 +1638,26 @@ 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)
+    # `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 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)
     for name in (
         "max",
         "min",
@@ -1074,11 +1700,126 @@ def nki_unpatch_lang(scope=None):
         scope.restore()
 
 
+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."""
+    from ..trace import nki_kernel_function
+
+    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, 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__, 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):
+    """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 globals_ in _modules_with("user", _user_module):
+        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):
     """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 +1848,9 @@ 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.memory = {}
         nki_builder.fn = self.fn
 
         kwargs.pop("warmup", None)  # Remove warmup from kwargs if it exists
@@ -1118,42 +1862,46 @@ 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):
+            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._namespaces)
+            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
             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]):
@@ -1165,7 +1913,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..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(
@@ -1232,9 +1271,11 @@ 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
-            self.fn(*kernel_args, **kwargs)
+                    continue  # a client that needs only some programs skips this one
+            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 abdd11a1..435f9ca4 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."""
+    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)
+    return func if isinstance(func, types.FunctionType) else None
+
+
 def trace_source(kernel):
     """
     Add the kernel code to be traceable within stack traces for clients
@@ -433,6 +444,10 @@ def decorator(kernel) -> TritonTrace | NKITrace | GluonTrace | Any:
             trace.add_client(client)
             return trace
 
+        if frontend in ("nki", "nki_beta2"):
+            # an `@nki.jit` (or `@nki.baremetal`, ...) kernel runs as its Python function
+            kernel = nki_kernel_function(kernel) or kernel
+
         trace_source(kernel)
 
         # First-time wrapping