K4 (`mlp`) and K1/K3 (`res_layernorm_fold`): -4.8% end-to-end encoder, bitwise identical

by jlongen-aws - opened 7 days ago

base: refs/heads/main

←

from: refs/pr/2

Discussion Files changed

+179

-292

Files changed (5) hide show

CHANGELOG.md +29 -0
build/torch-neuron/layers.py +2 -1
build/torch-neuron/nki_kernels/__init__.py +2 -2
build/torch-neuron/nki_kernels/mlp.py +97 -235
build/torch-neuron/nki_kernels/res_layernorm.py +49 -54

CHANGELOG.md CHANGED Viewed

@@ -1,5 +1,34 @@
 # Changelog
 ## v1.1 (2026-10)
 Documentation only: README points to

 # Changelog
+## v1.2 (2026-10)
+K4 (`mlp`) and K1/K3 (`res_layernorm_fold`) reworked: **-4.8% end-to-end encoder**, output
+bitwise identical to v1.1.
+* Full encoder (`torch.compile(backend="neuron")`, 24 layers, one 16-frame clip, trn2.3xlarge):
+  26.635 -> 25.353 ms/clip wall (3 rounds each), 26.83 -> 25.50 ms device time,
+  `mfu_estimated_percent` 39.5% -> 41.7%. Encoder output `torch.equal` to v1.1.
+* Standalone at the encoder shape (2 rounds each, all bitwise identical):
+  * K4: 295.3 -> 248.6 us (-15.8%), MFU 70.3% -> 82.8%.
+  * K1 (1 branch): 48.0 -> 41.9 us (-12.7%), HBM bandwidth 73% -> 84% of the 716 GB/s ceiling.
+  * K3 (2 branches): 57.7 -> 48.2 us (-16.3%), HBM bandwidth 71% -> 85%.
+* K4 schedule: x for row block 0 is staged before the weights (was queued behind 16 MB of
+  W1/W2); fc1 bias loads in one 16 KB DMA from a host-packed [128, H/128] layout (was 32 x 512 B
+  DMAs that stalled each gelu, and through PSUM the fc1 matmuls); row block 0's x is transposed on
+  the idle Tensor engine; b1 rides the Scalar engine's HW-DGE queue. Code: 3D SBUF tiles,
+  `nl.tile_size` constants, exact-multiple shapes asserted.
+* K1/K3: every DMA uses `dge_mode.none` (descriptors generated by the runtime before execution).
+  Before, each of the ~40 DMAs was issued by GpSimd's software DGE at ~0.75 us apiece, one after
+  another, and the first add waited ~8 us for them. Hardware DGE on the Sync queue was as fast
+  standalone but slower in the encoder, so it is not used. All token tiles are also resident in SBUF
+  and loaded before any compute. Same instructions in the same order, so results are unchanged.
+* `pack_b1` is exported; `mlp` accepts b1 either as [H] (packed on the fly) or pre-packed.
+* K4 now requires K % 512 == 0, H % 128 == 0, R % 1024 == 0 (the encoder shape satisfies all three).
+* Toolchain: everything above was measured on torch-neuronx 2.14.3 with nki 0.7.0b1. Standalone K4
+  and K1/K3 were also verified on the latest public NKI release (nki 0.6.0, neuronx-cc 2.27.5334):
+  K4 248.6 / 247.7 us vs 291.6 / 294.0 us for v1.1, K1 42.6 us, K3 48.1 us, all bitwise identical.
+  The full encoder could not be run on nki 0.6.0 with torch-neuronx 2.14.3 (v1.1 fails there too).
 ## v1.1 (2026-10)
 Documentation only: README points to

build/torch-neuron/layers.py CHANGED Viewed

@@ -40,6 +40,7 @@ from .nki_kernels import (
     attention_block,
     build_rope_tables,
     mlp,
     pack_out_weights,
     pack_tables,
     pack_v_weights,
@@ -85,7 +86,7 @@ def _pack_layer(lyr, tab):
         "wqk_f": wqk_f.to(dev), "bqk_f": bqk_f.to(dev), "tab": tab,
         "wvt_f": wvt_f.to(dev), "bv_f": bv_f.to(dev), "wot": wot.to(dev), "bo": bo.to(dev),
         "w1t_f": (w1 * g2[None, :]).t().contiguous().to(torch.bfloat16).to(dev),
-        "b1_f": (b1 + w1 @ t2).contiguous().to(dev),
         "w2t": lyr.mlp.fc2.weight.detach().t().contiguous().to(torch.bfloat16).to(dev),
         "b2": lyr.mlp.fc2.bias.detach().float().contiguous().to(dev),
     }

     attention_block,
     build_rope_tables,
     mlp,
+    pack_b1,
     pack_out_weights,
     pack_tables,
     pack_v_weights,
         "wqk_f": wqk_f.to(dev), "bqk_f": bqk_f.to(dev), "tab": tab,
         "wvt_f": wvt_f.to(dev), "bv_f": bv_f.to(dev), "wot": wot.to(dev), "bo": bo.to(dev),
         "w1t_f": (w1 * g2[None, :]).t().contiguous().to(torch.bfloat16).to(dev),
+        "b1_f": pack_b1(b1 + w1 @ t2).to(dev),          # [128, H/128], one 16 KB DMA in K4
         "w2t": lyr.mlp.fc2.weight.detach().t().contiguous().to(torch.bfloat16).to(dev),
         "b2": lyr.mlp.fc2.bias.detach().float().contiguous().to(dev),
     }

build/torch-neuron/nki_kernels/__init__.py CHANGED Viewed

@@ -8,12 +8,12 @@ from .attention_block import (
     pack_v_weights,
     pack_weights,
 )
-from .mlp import mlp
 from .res_layernorm import res_layernorm_fold
 from .rope_tables import build_rope_tables
 __all__ = [
-    "res_layernorm_fold", "attention_block", "mlp",
     "pack_weights", "pack_tables", "pack_v_weights", "pack_out_weights",
     "lane_permutation", "build_rope_tables",
 ]

     pack_v_weights,
     pack_weights,
 )
+from .mlp import mlp, pack_b1
 from .res_layernorm import res_layernorm_fold
 from .rope_tables import build_rope_tables
 __all__ = [
+    "res_layernorm_fold", "attention_block", "mlp", "pack_b1",
     "pack_weights", "pack_tables", "pack_v_weights", "pack_out_weights",
     "lane_permutation", "build_rope_tables",
 ]

build/torch-neuron/nki_kernels/mlp.py CHANGED Viewed

@@ -1,42 +1,21 @@
-"""K4 -- fused MLP (fc1 + gelu + fc2), x-prestage ("XP") in prologue mode.
-WHAT IT COMPUTES
-----------------
-``out = gelu(x @ w1.T + b1) @ w2.T + b2``, with the rows sharded across LNC cores and
-the weights pre-transposed to bf16. Returns ``[R, K]`` bf16.
-WHY IT IS SHAPED THIS WAY
--------------------------
-This is the effective shipping kernel: it is the ``b1pre`` MLP (fc1 bias staged in the
-prologue so every gelu's bias is resident from t=0) with the x-prestage applied. Two
-design choices carry the performance:
-* b1 first: the 32 per-block ``b1`` bias tiles are DMA'd before the 16 MB of weight
-  staging, so they land within ~1 us of kernel start and no gelu ever waits on its bias.
-* x prestaged in the prologue: BOTH row blocks' ``x`` load + dma_transpose staging is
-  issued before the ``rb`` loop, in row-block order. rb=0's staging already hid in the
-  pre-K4 seam (where the Tensor engine is dark for ~148 us because K1/K3 emit no Tensor
-  instructions); prestaging rb=1 gives it the same cover, filling the ~26.6 us Tensor
-  idle window that used to sit at the rb0->rb1 boundary. Appending rb=1's stages after
-  rb=0's, in order, means rb=0's readiness cannot regress -- the first matmul cannot be
-  delayed.
-The x staging is emitted through a single module-level function (``_stage_xt``) rather
-than inlined per row block, so the ``x`` load and the dma_transpose each keep exactly one
-source site. The two-row-block staging is preserved exactly: ``xt_sets`` holds one
-transposed-tile set per row block, and the loop consumes ``xt_sets[rb]``.
-PSUM: fc1 uses 1 bank, fc2 uses up to 4 banks; the ``y1_psum`` / ``out_psums`` tiles are
-declared inside their loops so the ring never accumulates onto a stale PSUM value. The
-b2 bias broadcast is kept (removing it regressed the sibling kernel by 23-30%).
-There are no ``engine=`` pins anywhere in this kernel. INFERENCE ONLY: no autograd rule.
 """
 from __future__ import annotations
-import math
 import torch
 from torch_neuronx import nki_op
@@ -44,251 +23,134 @@ import nki
 import nki.isa as nisa
 import nki.language as nl
-#: PE array partition limit, and the moving-operand free-axis limit.
-_P_MAX = 128
-_F_MOV_MAX = 512
 _BF16 = nl.bfloat16
 #: Logical-to-physical NeuronCore degree. Must match the shipping shard's.
 LNC = 2
-_INV_SQRT2 = 1.0 / math.sqrt(2.0)
-_INV_SQRT_2PI = 1.0 / math.sqrt(2.0 * math.pi)
-def _div_ceil(n: int, d: int) -> int:
-    """Ceiling division (never write (n + d - 1) // d inline)."""
-    return (n + d - 1) // d
 def _kernel_assert(condition: bool, error_text: str):
-    """Assert with an NKI-formatted error message."""
     assert condition, f"[INTERNAL_ERROR] [NCC_INKI016] Kernel validation exception: {error_text}"
-def _stage_xt(x_hbm, r_abs, pr, n_k, K, rb_size):
-    """Emit one row block's x staging and return its ``n_k`` transposed bf16 tiles.
-    ``n_rs`` loads of ``[<=128, K]`` bf16, each followed by ``n_k`` DMA transposes into
-    the row block's ``xt_tiles``. Kept as a module-level emitter so the ``x`` load and
-    the dma_transpose each map to a single source site regardless of where it is called.
-    Args:
-        x_hbm:   [R, K] @ HBM bf16, the MLP input.
-        r_abs:   absolute starting row of this row block.
-        pr:      number of rows in this row block (<= rb_size).
-        n_k:     number of K tiles (fc1 contraction dimension tiles).
-        K:       hidden-in dimension.
-        rb_size: row-block width (free axis of the transposed tiles).
-    Returns:
-        list of ``n_k`` transposed ``[128, rb_size]`` bf16 SBUF tiles.
-    """
-    n_rs = _div_ceil(pr, _P_MAX)
-    xt_tiles = []
-    for _ in range(n_k):
-        xt_tiles.append(nl.ndarray((_P_MAX, rb_size), dtype=_BF16, buffer=nl.sbuf))
     for rs in range(n_rs):
-        rr0 = rs * _P_MAX
-        prs = min(_P_MAX, pr - rr0)
-        x_b16 = nl.ndarray((_P_MAX, K), dtype=_BF16, buffer=nl.sbuf)
-        nisa.dma_copy(
-            dst=x_b16[0:prs, 0:K],
-            src=x_hbm[r_abs + rr0:r_abs + rr0 + prs, 0:K],
-        )
-        for kt in range(n_k):
-            k0 = kt * _P_MAX
-            pk = min(_P_MAX, K - k0)
-            nisa.dma_transpose(dst=xt_tiles[kt][0:pk, rr0:rr0 + prs],
-                               src=x_b16[0:prs, k0:k0 + pk])
-    return xt_tiles
 @nki.jit
 def _mlp_k4_xp_kernel(x_hbm, w1t_hbm, b1_hbm, w2t_hbm, b2_hbm):
-    """gelu(x @ w1.T + b1) @ w2.T + b2, rows sharded, b1 AND both x blocks resident early.
-    Args:
-        x_hbm   (nl.ndarray): [R, K] @ HBM, **bfloat16** (kernel 1 emits bf16).
-        w1t_hbm (nl.ndarray): [K, H] @ HBM, **bf16**, == fc1.weight.T.
-        b1_hbm  (nl.ndarray): [H]    @ HBM, float32.
-        w2t_hbm (nl.ndarray): [H, K] @ HBM, **bf16**, == fc2.weight.T.
-        b2_hbm  (nl.ndarray): [K]    @ HBM, float32.
-    Returns:
-        nl.ndarray: [R, K] @ HBM, **bfloat16**.
-    """
     R, K = x_hbm.shape
-    K_w1t, H = w1t_hbm.shape
-    H_w2t, K_w2t = w2t_hbm.shape
-    _kernel_assert(R >= 1 and K >= 1 and H >= 1, f"R({R}) K({K}) H({H}) must be >= 1")
-    _kernel_assert(K_w1t == K, f"w1t in-dim {K_w1t} != K {K}; expected w1t = w1.T of shape [K, H]")
-    _kernel_assert(K_w2t == K, f"w2t out-dim {K_w2t} != K {K}; expected w2t = w2.T of shape [H, K]")
-    _kernel_assert(H_w2t == H, f"w2t in-dim {H_w2t} != H {H}")
-    _kernel_assert(
-        R % LNC == 0,
-        f"the {LNC}-core kernel needs R divisible by {LNC}, got R={R}. "
-        f"program_id is a symbolic value, so a partial final row block cannot be "
-        f"clamped; use the one-core pre-transposed kernel for uneven R",
-    )
     rows_per = R // LNC
-    pid = nl.program_id(0)
-    r_base = pid * rows_per
     out_hbm = nl.ndarray((R, K), dtype=_BF16, buffer=nl.shared_hbm)
-    n_k = _div_ceil(K, _P_MAX)   # K tiles (fc1 contraction)
-    n_h = _div_ceil(H, _P_MAX)   # H tiles (fc1 out-rows / fc2 contraction)
-    # b1 FIRST, before the 16 MB of weight DMA: 32 [128,1] tiles issue at t=0 and land in
-    # ~1 us, so every gelu's bias is resident long before the first matmul retires.
-    b1_tiles = []
-    for ht in range(n_h):
-        h0 = ht * _P_MAX
-        ph = min(_P_MAX, H - h0)
-        t = nl.ndarray((_P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf)
-        nisa.dma_copy(dst=t[0:ph, 0:1], src=b1_hbm[h0:h0 + ph])
-        b1_tiles.append(t)
-    # Weight staging (pre-transposed bf16, one DMA per tile).
-    w1t_tiles = []
     for kt in range(n_k):
-        k0 = kt * _P_MAX
-        pk = min(_P_MAX, K - k0)
-        t = nl.ndarray((_P_MAX, H), dtype=_BF16, buffer=nl.sbuf)
-        nisa.dma_copy(dst=t[0:pk, :], src=w1t_hbm[k0:k0 + pk, :])
-        w1t_tiles.append(t)
-    w2t_tiles = []
     for ht in range(n_h):
-        h0 = ht * _P_MAX
-        ph = min(_P_MAX, H - h0)
-        t = nl.ndarray((_P_MAX, K), dtype=_BF16, buffer=nl.sbuf)
-        nisa.dma_copy(dst=t[0:ph, :], src=w2t_hbm[h0:h0 + ph, :])
-        w2t_tiles.append(t)
-    # Main fused loop.
-    rb_size = _F_MOV_MAX
-    n_rb = _div_ceil(rows_per, rb_size)
-    n_ko = _div_ceil(K, _F_MOV_MAX)
-    # THE X-PRESTAGE: stage EVERY row block's x here, in rb order, so rb=1's serialised
-    # DMA chain runs inside the pre-K4 seam (Tensor dark) instead of inside the idle
-    # window at the rb0->rb1 boundary. Appending rb=1's stages after rb=0's leaves rb=0's
-    # stages exactly where they are on the Sync engine's sub-stream.
-    xt_sets = []
-    for rb_pre in range(n_rb):
-        xt_sets.append(_stage_xt(
-            x_hbm, r_base + rb_pre * rb_size,
-            min(rb_size, rows_per - rb_pre * rb_size), n_k, K, rb_size))
     for rb in range(n_rb):
-        r_local = rb * rb_size
-        pr = min(rb_size, rows_per - r_local)
-        n_rs = _div_ceil(pr, _P_MAX)
-        r_abs = r_base + r_local
-        # x_T bf16 tiles: consume the prologue-staged set for this row block.
-        xt_tiles = xt_sets[rb]
-        # ---- fc1: gt_tiles[ht] = gelu(w1 @ x.T + b1), bf16, SBUF-resident ----
-        gt_tiles = []
         for ht in range(n_h):
-            h0 = ht * _P_MAX
-            ph = min(_P_MAX, H - h0)
-            # Declared INSIDE the ht loop, so the ring never reuses a written PSUM tile.
-            # nc_matmul's default accumulate=None auto-detects "previously written" and
-            # would accumulate onto a stale value.
-            y1_psum = nl.ndarray((_P_MAX, rb_size), dtype=nl.float32, buffer=nl.psum)
             for kt in nl.affine_range(n_k):
-                k0 = kt * _P_MAX
-                pk = min(_P_MAX, K - k0)
-                nisa.nc_matmul(
-                    dst=y1_psum[0:ph, 0:pr],
-                    stationary=w1t_tiles[kt][0:pk, h0:h0 + ph],
-                    moving=xt_tiles[kt][0:pk, 0:pr],
-                )
-            # b1 comes from the prologue-resident tile.
-            gt = nl.ndarray((_P_MAX, rb_size), dtype=_BF16, buffer=nl.sbuf)
-            nisa.activation(dst=gt[0:ph, 0:pr], op=nl.gelu,
-                            data=y1_psum[0:ph, 0:pr], bias=b1_tiles[ht][0:ph, 0:1])
-            gt_tiles.append(gt)
-        # ---- fc2: out[r,k] = sum_h g[r,h] * w2[k,h] + b2[k] ----
         for ko in range(n_ko):
-            ko0 = ko * _F_MOV_MAX
-            nk = min(_F_MOV_MAX, K - ko0)
-            out_psums = []
             for _ in range(n_rs):
-                out_psums.append(
-                    nl.ndarray((_P_MAX, _F_MOV_MAX), dtype=nl.float32, buffer=nl.psum)
-                )
             for ht in nl.affine_range(n_h):
-                h0 = ht * _P_MAX
-                ph = min(_P_MAX, H - h0)
                 for rs in range(n_rs):
-                    rr0 = rs * _P_MAX
-                    prs = min(_P_MAX, pr - rr0)
-                    nisa.nc_matmul(
-                        dst=out_psums[rs][0:prs, 0:nk],
-                        stationary=gt_tiles[ht][0:ph, rr0:rr0 + prs],
-                        moving=w2t_tiles[ht][0:ph, ko0:ko0 + nk],
-                    )
-            # b2 broadcast, hoisted out of the rs loop (removing the broadcast itself
-            # cost 23% and 30% in two separate attempts).
-            b2_bc = nl.ndarray((_P_MAX, _F_MOV_MAX), dtype=nl.float32, buffer=nl.sbuf)
-            nisa.dma_copy(dst=b2_bc[0:_P_MAX, 0:nk],
-                          src=b2_hbm.ap(pattern=[[0, _P_MAX], [1, nk]], offset=ko0))
             for rs in range(n_rs):
-                rr0 = rs * _P_MAX
-                prs = min(_P_MAX, pr - rr0)
-                out_sb = nl.ndarray((_P_MAX, _F_MOV_MAX), dtype=_BF16, buffer=nl.sbuf)
-                nisa.tensor_tensor(dst=out_sb[0:prs, 0:nk],
-                                   data1=out_psums[rs][0:prs, 0:nk],
-                                   data2=b2_bc[0:prs, 0:nk],
-                                   op=nl.add)
-                nisa.dma_copy(
-                    dst=out_hbm[r_abs + rr0:r_abs + rr0 + prs, ko0:ko0 + nk],
-                    src=out_sb[0:prs, 0:nk],
-                )
     return out_hbm
 def _no_grad(*tensors) -> None:
     """Refuse gradient-tracking inputs -- this op registers no autograd rule."""
     for t in tensors:
         if t is not None and t.requires_grad:
-            raise RuntimeError(
-                "nki_mlp_k4_xp registers no autograd rule (INFERENCE ONLY). "
-                "Use under torch.no_grad()/eval()."
-            )
 @nki_op("vjepa2::nki_mlp_k4_xp", mutates_args={})
-def nki_mlp_k4_xp(
-    x: torch.Tensor,
-    w1t: torch.Tensor,
-    b1: torch.Tensor,
-    w2t: torch.Tensor,
-    b2: torch.Tensor,
-) -> torch.Tensor:
-    """Fused MLP on LNC=2, pre-transposed bf16 weights, x prestaged in the prologue.
-    Args:
-        x   : [R, K] **bf16** on 'neuron'. ``R`` must be divisible by :data:`LNC`.
-        w1t : [K, H] bf16 == fc1.weight.T.   b1 : [H] float32.
-        w2t : [H, K] bf16 == fc2.weight.T.   b2 : [K] float32.
-    Returns:
-        [R, K] bf16.
-    """
     _no_grad(x, w1t, b1, w2t, b2)
     return _mlp_k4_xp_kernel[LNC](x, w1t, b1, w2t, b2)
-#: Export alias. Same object as the nki_op above, so the runtime
-#: op-name string ("vjepa2::nki_mlp_k4_xp") is unchanged whichever name is used.
 mlp = nki_mlp_k4_xp

+"""K4 -- fused MLP (fc1 + gelu + fc2).
+``out = gelu(x @ w1.T + b1) @ w2.T + b2``, rows sharded across LNC cores, bf16 matmuls with
+fp32 PSUM accumulation. Same instruction order as v10 (bit-exact with mlp.py), with:
+* 3D SBUF tiles instead of lists of tiles: W1 [P, n_k, H], W2 [P, n_h, K], x per row block
+  [P, n_rs, K], fc1 output [P, n_h, FMOV]. DMAs stay per tile (v14 merged them into single
+  access-pattern DMAs and lost ~8 us: one completion signal per 8 MB weight), b2 is one
+  partition-broadcast access-pattern load [P, K];
+* hardware sizes from ``nl.tile_size`` instead of literals;
+* exact-multiple shapes asserted up front instead of ``min()`` remainder handling.
+Prologue order (v3/v8/v10): b1 (Scalar HW-DGE queue) -> x(rb0) + Tensor-engine transposes
+-> W1 -> x(rb1..) + DMA transposes -> W2. INFERENCE ONLY: no autograd rule.
 """
 from __future__ import annotations
 import torch
 from torch_neuronx import nki_op
 import nki.isa as nisa
 import nki.language as nl
+P = nl.tile_size.pmax                 # partitions / stationary width (128)
+FMOV = nl.tile_size.gemm_moving_fmax  # moving free-axis limit = one PSUM bank of fp32 (512)
 _BF16 = nl.bfloat16
 #: Logical-to-physical NeuronCore degree. Must match the shipping shard's.
 LNC = 2
 def _kernel_assert(condition: bool, error_text: str):
     assert condition, f"[INTERNAL_ERROR] [NCC_INKI016] Kernel validation exception: {error_text}"
+def _load_rows(x_hbm, r0, n_rs, K):
+    """x[r0 : r0 + n_rs*P, :] into [P, n_rs, K], one DMA per row chunk so consumers start early."""
+    x_sb = nl.ndarray((P, n_rs, K), dtype=_BF16, buffer=nl.sbuf)
     for rs in range(n_rs):
+        nisa.dma_copy(dst=x_sb[:, rs, :], src=x_hbm[r0 + rs * P:r0 + (rs + 1) * P, :])
+    return x_sb
 @nki.jit
 def _mlp_k4_xp_kernel(x_hbm, w1t_hbm, b1_hbm, w2t_hbm, b2_hbm):
+    """Args: x [R, K] bf16, w1t [K, H] bf16 (= fc1.weight.T), b1 [P, H/P] fp32 (pack_b1),
+    w2t [H, K] bf16 (= fc2.weight.T), b2 [K] fp32. Returns [R, K] bf16."""
     R, K = x_hbm.shape
+    H = w1t_hbm.shape[1]
     rows_per = R // LNC
+    _kernel_assert(tuple(w1t_hbm.shape) == (K, H) and tuple(w2t_hbm.shape) == (H, K),
+                   f"expected w1t [K, H] and w2t [H, K], got {w1t_hbm.shape}, {w2t_hbm.shape}")
+    _kernel_assert(tuple(b1_hbm.shape) == (P, H // P), f"b1 must be pack_b1(b1): [{P}, {H // P}]")
+    _kernel_assert(K % FMOV == 0 and H % P == 0 and R % (LNC * FMOV) == 0,
+                   f"needs K % {FMOV} == 0, H % {P} == 0, R % {LNC * FMOV} == 0; got R={R} K={K} H={H}")
+    n_k, n_h, n_ko, n_rs, n_rb = K // P, H // P, K // FMOV, FMOV // P, rows_per // FMOV
+    r_base = nl.program_id(0) * rows_per
     out_hbm = nl.ndarray((R, K), dtype=_BF16, buffer=nl.shared_hbm)
+    # b1 on the Scalar engine's HW-DGE queue, so gelu's wait doesn't count W1/W2 traffic.
+    b1 = nl.ndarray((P, n_h), dtype=nl.float32, buffer=nl.sbuf)
+    nisa.dma_copy(dst=b1, src=b1_hbm, dge_mode=nisa.dge_mode.hwdge, engine=nisa.engine.scalar)
+    # x for rb0, transposed on the (idle) Tensor engine: xt[rb][:, kt, :] is x.T's K tile kt.
+    xt = []
+    for _ in range(n_rb):
+        xt.append(nl.ndarray((P, n_k, FMOV), dtype=_BF16, buffer=nl.sbuf))
+    x0 = _load_rows(x_hbm, r_base, n_rs, K)
     for kt in range(n_k):
+        tp = nl.ndarray((P, FMOV), dtype=_BF16, buffer=nl.psum)
+        for rs in range(n_rs):
+            nisa.nc_transpose(dst=tp[:, rs * P:(rs + 1) * P], data=x0[:, rs, kt * P:(kt + 1) * P],
+                              engine=nisa.engine.tensor)
+        nisa.tensor_copy(dst=xt[0][:, kt, :], src=tp)
+    # W1 as [P, n_k, H]: w1[p, kt, h] = w1t[kt*P + p, h]. One DMA per K tile.
+    w1 = nl.ndarray((P, n_k, H), dtype=_BF16, buffer=nl.sbuf)
+    for kt in range(n_k):
+        nisa.dma_copy(dst=w1[:, kt, :], src=w1t_hbm[kt * P:(kt + 1) * P, :])
+    # Later row blocks' x via DMA transposes, streaming under rb0's matmuls.
+    for rb in range(1, n_rb):
+        xr = _load_rows(x_hbm, r_base + rb * FMOV, n_rs, K)
+        for rs in range(n_rs):
+            for kt in range(n_k):
+                nisa.dma_transpose(dst=xt[rb][:, kt, rs * P:(rs + 1) * P],
+                                   src=xr[:, rs, kt * P:(kt + 1) * P])
+    # W2 as [P, n_h, K]: w2[p, ht, k] = w2t[ht*P + p, k]. One DMA per H tile, so fc2 can
+    # start on the first strips (a single 8 MB DMA signals completion only at the end).
+    w2 = nl.ndarray((P, n_h, K), dtype=_BF16, buffer=nl.sbuf)
     for ht in range(n_h):
+        nisa.dma_copy(dst=w2[:, ht, :], src=w2t_hbm[ht * P:(ht + 1) * P, :])
     for rb in range(n_rb):
+        r_abs = r_base + rb * FMOV
+        # fc1 + gelu: g[:, ht, :] = gelu(w1[:, :, ht].T @ x.T + b1[:, ht]), kept in SBUF.
+        g = nl.ndarray((P, n_h, FMOV), dtype=_BF16, buffer=nl.sbuf)
         for ht in range(n_h):
+            # Declared per ht so the PSUM ring never accumulates onto a stale value.
+            y1 = nl.ndarray((P, FMOV), dtype=nl.float32, buffer=nl.psum)
             for kt in nl.affine_range(n_k):
+                nisa.nc_matmul(dst=y1, stationary=w1[:, kt, ht * P:(ht + 1) * P], moving=xt[rb][:, kt, :])
+            nisa.activation(dst=g[:, ht, :], op=nl.gelu, data=y1, bias=b1[:, ht:ht + 1])
+        # fc2 + b2: out[rows, ko chunk] = sum_ht g[:, ht, rows].T @ w2[:, ht, chunk] + b2.
         for ko in range(n_ko):
+            k0 = ko * FMOV
+            acc = []
             for _ in range(n_rs):
+                acc.append(nl.ndarray((P, FMOV), dtype=nl.float32, buffer=nl.psum))
             for ht in nl.affine_range(n_h):
                 for rs in range(n_rs):
+                    nisa.nc_matmul(dst=acc[rs], stationary=g[:, ht, rs * P:(rs + 1) * P],
+                                   moving=w2[:, ht, k0:k0 + FMOV])
+            # b2 chunk broadcast across partitions (stride 0), loaded per (rb, ko) as in v10.
+            b2 = nl.ndarray((P, FMOV), dtype=nl.float32, buffer=nl.sbuf)
+            nisa.dma_copy(dst=b2, src=b2_hbm.ap(pattern=[[0, P], [1, FMOV]], offset=k0))
+            o = nl.ndarray((P, n_rs, FMOV), dtype=_BF16, buffer=nl.sbuf)
             for rs in range(n_rs):
+                nisa.tensor_tensor(dst=o[:, rs, :], data1=acc[rs], data2=b2, op=nl.add)
+                nisa.dma_copy(dst=out_hbm[r_abs + rs * P:r_abs + (rs + 1) * P, k0:k0 + FMOV], src=o[:, rs, :])
     return out_hbm
+def pack_b1(b1: torch.Tensor) -> torch.Tensor:
+    """Pack fc1.bias [H] fp32 into [P, H/P] so b1p[p, ht] == b1[ht*P + p]."""
+    return b1.reshape(-1, P).T.contiguous()
 def _no_grad(*tensors) -> None:
     """Refuse gradient-tracking inputs -- this op registers no autograd rule."""
     for t in tensors:
         if t is not None and t.requires_grad:
+            raise RuntimeError("nki_mlp_k4_xp registers no autograd rule (INFERENCE ONLY). "
+                               "Use under torch.no_grad()/eval().")
 @nki_op("vjepa2::nki_mlp_k4_xp", mutates_args={})
+def nki_mlp_k4_xp(x: torch.Tensor, w1t: torch.Tensor, b1: torch.Tensor,
+                  w2t: torch.Tensor, b2: torch.Tensor) -> torch.Tensor:
+    """Fused MLP on LNC=2. x [R, K] bf16, w1t [K, H] bf16, b1 [H] or pack_b1(b1) [P, H/P] fp32,
+    w2t [H, K] bf16, b2 [K] fp32. Returns [R, K] bf16."""
     _no_grad(x, w1t, b1, w2t, b2)
+    if b1.dim() == 1:  # accept the pre-v1.2 [H] bias; NeuronVJEPA2Encoder passes it pre-packed
+        b1 = pack_b1(b1)
     return _mlp_k4_xp_kernel[LNC](x, w1t, b1, w2t, b2)
+#: Export alias, so callers using ``mlp`` (e.g. layers.py) pick up this op unchanged.
 mlp = nki_mlp_k4_xp

build/torch-neuron/nki_kernels/res_layernorm.py CHANGED Viewed

@@ -19,9 +19,15 @@ WHY IT IS SHAPED THIS WAY
 -------------------------
 * LNC = 2 shard: the kernel is launched on both NeuronCores; each owns ``R // 2`` rows
   selected by ``program_id(0)``, so the two cores never touch the same tokens.
-* Ring of depth 4 over the token tiles: four rotating sets of SBUF tiles let the DMA
-  loads of tile ``t+1`` overlap the compute of tile ``t`` without a write-after-read
-  hazard on the buffers.
 * fp32 residual / bf16 branch: ``a`` (the residual stream) stays fp32 so ``x_new`` is
   bitwise-invariant math and can be carried forward losslessly; the branch inputs and
   the normalized output ``z`` are bf16, which is the precision boundary the downstream
@@ -54,9 +60,6 @@ _BF16 = nl.bfloat16
 #: Logical-to-physical NeuronCore degree. Rows are split evenly across LNC cores.
 LNC = 2
-#: Ring depth: number of rotating SBUF tile sets that overlap load with compute.
-_RING = 4
 def _kernel_assert(condition: bool, error_text: str):
     """Assert with an NKI-formatted error message."""
@@ -68,8 +71,8 @@ def _res_layernorm_fold_kernel(a_hbm, b_hbm, eps):
     """Emit the fold kernel body: ``(a + sum(b), LN_no_affine(a + sum(b)))``, token-major.
     gamma/beta live in the consumers' packed weights, so this kernel takes no affine
-    parameters. Rows are sharded across ``LNC`` cores; each token tile is loaded through
-    a depth-4 ring, summed in fp32, normalized with the thinned stat chain, and the
     normalized result is cast to bf16 on write.
     Args:
@@ -97,8 +100,7 @@ def _res_layernorm_fold_kernel(a_hbm, b_hbm, eps):
     rows_per = R // LNC
     n_t = rows_per // _P_MAX
-    pid = nl.program_id(0)
-    r_base = pid * rows_per
     x_new_hbm = nl.ndarray((R, H), dtype=nl.float32, buffer=nl.shared_hbm)
     z_hbm = nl.ndarray((R, H), dtype=_BF16, buffer=nl.shared_hbm)
@@ -106,69 +108,62 @@ def _res_layernorm_fold_kernel(a_hbm, b_hbm, eps):
     inv_h = 1.0 / float(H)
     neg_inv_h = -inv_h
-    # Ring buffers: _RING rotating sets of (a tile, B branch tiles, sum tile, z tile).
-    a_sb = []
-    b_sb = []
-    s_sb = []
-    z_sb = []
-    for _i in range(_RING):
-        a_sb.append(nl.ndarray((_P_MAX, H), dtype=nl.float32, buffer=nl.sbuf))
-        _slot = []
-        for _j in range(n_b):
-            _slot.append(nl.ndarray((_P_MAX, H), dtype=_BF16, buffer=nl.sbuf))
-        b_sb.append(_slot)
-        s_sb.append(nl.ndarray((_P_MAX, H), dtype=nl.float32, buffer=nl.sbuf))
-        z_sb.append(nl.ndarray((_P_MAX, H), dtype=_BF16, buffer=nl.sbuf))
     # Scratch destination for the reduce activations (its full tile is discarded; only
     # the reduce result is consumed).
     sink = nl.ndarray((_P_MAX, H), dtype=nl.float32, buffer=nl.sbuf)
-    # Per-lane [128, 1] statistics.
-    acc = nl.ndarray((_P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf)
-    neg_shift = nl.ndarray((_P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf)
-    neg_mu = nl.ndarray((_P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf)
-    var = nl.ndarray((_P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf)
-    rstd = nl.ndarray((_P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf)
     for t in range(n_t):
-        k = t % _RING
         r0 = r_base + t * _P_MAX
-        # ---- load residual + branches for this tile ----
-        nisa.dma_copy(dst=a_sb[k], src=a_hbm[r0:r0 + _P_MAX, 0:H])
-        for j in range(n_b):
-            nisa.dma_copy(dst=b_sb[k][j], src=b_hbm[j, r0:r0 + _P_MAX, 0:H])
         # ---- x_new = a + sum_j b[j] in fp32, then write the residual stream out ----
-        nisa.tensor_tensor(dst=s_sb[k], data1=a_sb[k], data2=b_sb[k][0], op=nl.add)
         for j in range(1, n_b):
-            nisa.tensor_tensor(dst=s_sb[k], data1=s_sb[k], data2=b_sb[k][j], op=nl.add)
-        nisa.dma_copy(dst=x_new_hbm[r0:r0 + _P_MAX, 0:H], src=s_sb[k])
         # ---- thinned stat chain: mean via shift-and-reduce, variance, rstd ----
-        nisa.tensor_scalar(dst=neg_shift, data=s_sb[k][0:_P_MAX, 0:1],
                            op0=nl.multiply, operand0=-1.0)
-        nisa.activation_reduce(dst=sink, op=nl.copy, data=s_sb[k],
-                               reduce_op=nl.add, reduce_res=acc,
-                               bias=neg_shift, scale=1.0)
-        nisa.tensor_scalar(dst=neg_mu, data=acc, op0=nl.multiply, operand0=neg_inv_h,
-                           op1=nl.add, operand1=neg_shift)
-        nisa.activation_reduce(dst=sink, op=nl.square, data=s_sb[k],
-                               reduce_op=nl.add, reduce_res=acc, bias=neg_mu, scale=1.0)
-        nisa.tensor_scalar(dst=var, data=acc, op0=nl.multiply, operand0=inv_h,
                            op1=nl.add, operand1=eps)
-        nisa.activation(dst=rstd, op=nl.rsqrt, data=var)
         # ---- one-op affine-free apply: z = (x_new + neg_mu) * rstd, bf16 cast on write.
-        # Both operands are per-lane [128, 1] vectors; the fused two-op tensor_scalar
-        # costs the same as one op.
-        nisa.tensor_scalar(dst=z_sb[k], data=s_sb[k],
-                           op0=nl.add, operand0=neg_mu,
-                           op1=nl.multiply, operand1=rstd)
-        nisa.dma_copy(dst=z_hbm[r0:r0 + _P_MAX, 0:H], src=z_sb[k])
     return x_new_hbm, z_hbm

 -------------------------
 * LNC = 2 shard: the kernel is launched on both NeuronCores; each owns ``R // 2`` rows
   selected by ``program_id(0)``, so the two cores never touch the same tokens.
+* Every token tile has its own SBUF slot ([P, n_t, H]; ~116 KB/partition at H=1024, B=2)
+  and all loads are issued before any compute, so tile ``t+1``'s loads never wait behind
+  tile ``t``'s stores. One DMA per tile and operand: a single whole-shard DMA would signal
+  completion only at the end.
+* ``dge_mode.none`` on every DMA: the runtime generates the descriptors before execution,
+  so no engine spends time issuing them. The default (software DGE) issued each of the ~40
+  DMAs on GpSimd at ~0.75 us apiece, one after another. Hardware DGE on the Sync queue was
+  as fast standalone, but inside the encoder it delayed K2/K4's loads on that queue by
+  about as much as it saved.
 * fp32 residual / bf16 branch: ``a`` (the residual stream) stays fp32 so ``x_new`` is
   bitwise-invariant math and can be carried forward losslessly; the branch inputs and
   the normalized output ``z`` are bf16, which is the precision boundary the downstream
 #: Logical-to-physical NeuronCore degree. Rows are split evenly across LNC cores.
 LNC = 2
 def _kernel_assert(condition: bool, error_text: str):
     """Assert with an NKI-formatted error message."""
     """Emit the fold kernel body: ``(a + sum(b), LN_no_affine(a + sum(b)))``, token-major.
     gamma/beta live in the consumers' packed weights, so this kernel takes no affine
+    parameters. Rows are sharded across ``LNC`` cores; every token tile is loaded up
+    front into its own SBUF slot, summed in fp32, normalized with the thinned stat chain, and the
     normalized result is cast to bf16 on write.
     Args:
     rows_per = R // LNC
     n_t = rows_per // _P_MAX
+    r_base = nl.program_id(0) * rows_per
     x_new_hbm = nl.ndarray((R, H), dtype=nl.float32, buffer=nl.shared_hbm)
     z_hbm = nl.ndarray((R, H), dtype=_BF16, buffer=nl.shared_hbm)
     inv_h = 1.0 / float(H)
     neg_inv_h = -inv_h
+    # One SBUF slot per tile; all loads issued before any compute or store (see module doc).
+    a_sb = nl.ndarray((_P_MAX, n_t, H), dtype=nl.float32, buffer=nl.sbuf)
+    b_sb = nl.ndarray((_P_MAX, n_b, n_t, H), dtype=_BF16, buffer=nl.sbuf)
+    for t in range(n_t):
+        r0 = r_base + t * _P_MAX
+        nisa.dma_copy(dst=a_sb[:, t, :], src=a_hbm[r0:r0 + _P_MAX, 0:H], dge_mode=nisa.dge_mode.none)
+        for j in range(n_b):
+            nisa.dma_copy(dst=b_sb[:, j, t, :], src=b_hbm[j, r0:r0 + _P_MAX, 0:H], dge_mode=nisa.dge_mode.none)
+    s_sb = nl.ndarray((_P_MAX, n_t, H), dtype=nl.float32, buffer=nl.sbuf)
+    z_sb = nl.ndarray((_P_MAX, n_t, H), dtype=_BF16, buffer=nl.sbuf)
     # Scratch destination for the reduce activations (its full tile is discarded; only
     # the reduce result is consumed).
     sink = nl.ndarray((_P_MAX, H), dtype=nl.float32, buffer=nl.sbuf)
+    # Per-lane statistics, one column per tile so tile t+1's chain doesn't wait on tile t's.
+    acc = nl.ndarray((_P_MAX, n_t), dtype=nl.float32, buffer=nl.sbuf)
+    acc2 = nl.ndarray((_P_MAX, n_t), dtype=nl.float32, buffer=nl.sbuf)
+    neg_shift = nl.ndarray((_P_MAX, n_t), dtype=nl.float32, buffer=nl.sbuf)
+    neg_mu = nl.ndarray((_P_MAX, n_t), dtype=nl.float32, buffer=nl.sbuf)
+    var = nl.ndarray((_P_MAX, n_t), dtype=nl.float32, buffer=nl.sbuf)
+    rstd = nl.ndarray((_P_MAX, n_t), dtype=nl.float32, buffer=nl.sbuf)
     for t in range(n_t):
         r0 = r_base + t * _P_MAX
+        s = s_sb[:, t, :]
+        c = slice(t, t + 1)
         # ---- x_new = a + sum_j b[j] in fp32, then write the residual stream out ----
+        nisa.tensor_tensor(dst=s, data1=a_sb[:, t, :], data2=b_sb[:, 0, t, :], op=nl.add)
         for j in range(1, n_b):
+            nisa.tensor_tensor(dst=s, data1=s, data2=b_sb[:, j, t, :], op=nl.add)
+        nisa.dma_copy(dst=x_new_hbm[r0:r0 + _P_MAX, 0:H], src=s, dge_mode=nisa.dge_mode.none)
         # ---- thinned stat chain: mean via shift-and-reduce, variance, rstd ----
+        nisa.tensor_scalar(dst=neg_shift[:, c], data=s_sb[:, t, 0:1],
                            op0=nl.multiply, operand0=-1.0)
+        nisa.activation_reduce(dst=sink, op=nl.copy, data=s,
+                               reduce_op=nl.add, reduce_res=acc[:, c],
+                               bias=neg_shift[:, c], scale=1.0)
+        nisa.tensor_scalar(dst=neg_mu[:, c], data=acc[:, c], op0=nl.multiply, operand0=neg_inv_h,
+                           op1=nl.add, operand1=neg_shift[:, c])
+        nisa.activation_reduce(dst=sink, op=nl.square, data=s,
+                               reduce_op=nl.add, reduce_res=acc2[:, c], bias=neg_mu[:, c], scale=1.0)
+        nisa.tensor_scalar(dst=var[:, c], data=acc2[:, c], op0=nl.multiply, operand0=inv_h,
                            op1=nl.add, operand1=eps)
+        nisa.activation(dst=rstd[:, c], op=nl.rsqrt, data=var[:, c])
         # ---- one-op affine-free apply: z = (x_new + neg_mu) * rstd, bf16 cast on write.
+        nisa.tensor_scalar(dst=z_sb[:, t, :], data=s,
+                           op0=nl.add, operand0=neg_mu[:, c],
+                           op1=nl.multiply, operand1=rstd[:, c])
+        nisa.dma_copy(dst=z_hbm[r0:r0 + _P_MAX, 0:H], src=z_sb[:, t, :], dge_mode=nisa.dge_mode.none)
     return x_new_hbm, z_hbm