v1.1: K4 (mlp/mlp_res_ln) + K1/K3 (res_layernorm) DMA/engine-schedule rework

Bitwise-identical to v1 at LNC=1 and LNC=2. Packed fc1 bias loaded in one DMA on the
Scalar HW-DGE queue; row-block-0 x transposed on the Tensor engine before the weights;
res_layernorm uses dge_mode.none with all token tiles SBUF-resident. Based on an external
contribution by jlongen-aws (PR #2 on jburtoft/vjepa2-neuron-kernels), re-implemented with
the LNC-parametric path preserved. Verified on trn2.3xlarge: real-weight V-JEPA2 output
bitwise identical to v1, -4.0% ms/clip.

Files changed (6) hide show

CHANGELOG.md +35 -0
build/torch-neuron/encoder.py +3 -1
build/torch-neuron/nki_kernels/__init__.py +2 -1
build/torch-neuron/nki_kernels/mlp.py +103 -65
build/torch-neuron/nki_kernels/mlp_res_ln.py +46 -104
build/torch-neuron/nki_kernels/res_layernorm.py +60 -56

CHANGELOG.md CHANGED Viewed

@@ -1,5 +1,40 @@
 # Changelog
 ## v1 (2026-10)
 First release.

 # Changelog
+## v1.1 (2026-10)
+Reworked DMA/engine schedule for `mlp` (K4), `mlp_res_ln` (K4 + next LN) and `res_layernorm_fold`
+(K1/K3). Output is **bitwise identical** to v1 at both LNC=1 and LNC=2. Based on an external
+contribution by [jlongen-aws](https://huggingface.co/jlongen-aws) (PR #2 on
+`jburtoft/vjepa2-neuron-kernels`), re-implemented here with the LNC-parametric path and the `min()`
+remainder handling preserved.
+Verified on trn2.3xlarge (A/B on one consistent stack: old vs new compiled and run in the same
+container + inputs, device cycles from `neuron-explorer show-session`). LNC=2, V-JEPA2 ViT-L shape
+(R=2048, H=1024, MLP hidden 4096):
+* K4 `mlp`: -7.5% device cycles (reproducible, -7.8% +/- ~1% over 3 captures).
+* K4+LN `mlp_res_ln`: -6.8%.
+* K1 `res_layernorm_fold` (1 branch): -16.2%; K3 (2 branches): -18.3%.
+* Full 24-layer encoder (synthetic ViT-L weights, includes the unchanged attention): -3.7% device
+  cycles, output `torch.equal` to v1.
+* Real-weight V-JEPA2 (`facebook/vjepa2-vitl-fpc64-256`, 16 frames x 256 px, LNC=2): output bitwise
+  identical to v1 (`new_vs_v1_rel` 0.0), rel_l2 0.0389 vs fp32 CPU (unchanged from v1), -4.0% ms/clip
+  (speedup vs stock fp32 3.19x -> 3.33x).
+* Per-kernel output is `torch.equal` to v1 at LNC=1 and LNC=2.
+Changes:
+* `mlp` / `mlp_res_ln`: fc1 bias is host-packed to `[128, H/128]` (`pack_b1`, now exported) and
+  loaded in one DMA on the Scalar engine's HW-DGE queue instead of `H/128` column DMAs; row block 0's
+  `x` is staged and transposed on the (idle) Tensor engine before the weights, later row blocks stream
+  under row block 0's matmuls. `mlp` still accepts an unpacked `[H]` bias (packed on the fly).
+* `res_layernorm_fold`: every DMA uses `dge_mode.none` (the runtime pre-generates descriptors) and all
+  token tiles are made SBUF-resident before any compute, replacing the depth-4 ring. Per-tile statistic
+  columns (`[128, n_t]`) so one tile's stat chain does not wait on the previous tile's.
+* `pack_block` now stores the GELU fc1 bias pre-packed; `pack_b1` is exported from `nki_kernels`.
+The compute instructions and their order are unchanged, so results are bit-identical to v1.
 ## v1 (2026-10)
 First release.

build/torch-neuron/encoder.py CHANGED Viewed

@@ -34,6 +34,7 @@ _ab1 = _il.import_module(__package__ + ".nki_kernels.attention_block_l1")
 _ab1ln = _il.import_module(__package__ + ".nki_kernels.attention_block_l1_ln")
 from .nki_kernels.attention_block import LNC
 from .nki_kernels.mlp import nki_mlp_k4_xp
 from .nki_kernels.mlp_res_ln import nki_mlp_res_ln
 from .nki_kernels.res_layernorm import nki_res_layernorm_fold
 from .nki_kernels.swiglu_stream import nki_swiglu_stream, nki_swiglu_stream_ln
@@ -132,7 +133,8 @@ def pack_block(bw: BlockWeights, num_heads: int, rope: str, device) -> dict:
         w1, w2 = _f(bw.fc1_w), _f(bw.fc2_w)
         p["swiglu"] = False
         p["w1t"] = (w1 * g2[None]).t().contiguous().to(torch.bfloat16)
-        p["b1"] = (_z(_f(bw.fc1_b), w1.shape[0]) + w1 @ t2).contiguous()
         p["w2t"] = (w2 * ls2[:, None]).t().contiguous().to(torch.bfloat16)
         p["b2"] = (_z(_f(bw.fc2_b), H) * ls2).contiguous()
     return {k: (v.to(device) if torch.is_tensor(v) else v) for k, v in p.items()}

 _ab1ln = _il.import_module(__package__ + ".nki_kernels.attention_block_l1_ln")
 from .nki_kernels.attention_block import LNC
 from .nki_kernels.mlp import nki_mlp_k4_xp
+from .nki_kernels.mlp import pack_b1 as _mlp_pack_b1
 from .nki_kernels.mlp_res_ln import nki_mlp_res_ln
 from .nki_kernels.res_layernorm import nki_res_layernorm_fold
 from .nki_kernels.swiglu_stream import nki_swiglu_stream, nki_swiglu_stream_ln
         w1, w2 = _f(bw.fc1_w), _f(bw.fc2_w)
         p["swiglu"] = False
         p["w1t"] = (w1 * g2[None]).t().contiguous().to(torch.bfloat16)
+        # Pre-pack fc1 bias to [128, H_mlp/128] so K4/K4+LN load it in one DMA (see mlp.pack_b1).
+        p["b1"] = _mlp_pack_b1((_z(_f(bw.fc1_b), w1.shape[0]) + w1 @ t2).contiguous())
         p["w2t"] = (w2 * ls2[:, None]).t().contiguous().to(torch.bfloat16)
         p["b2"] = (_z(_f(bw.fc2_b), H) * ls2).contiguous()
     return {k: (v.to(device) if torch.is_tensor(v) else v) for k, v in p.items()}

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

@@ -11,6 +11,7 @@ from .attention_block import (
 from .attention_block_l1 import attention_block_l1
 from .attention_block_l1_ln import attention_block_l1_ln
 from .mlp import nki_mlp_k4_xp as mlp
 from .mlp_res_ln import nki_mlp_res_ln as mlp_res_ln
 from .res_layernorm import nki_res_layernorm_fold as res_layernorm_fold
 from .swiglu_stream import nki_swiglu_stream as swiglu_stream
@@ -18,6 +19,6 @@ from .swiglu_stream import nki_swiglu_stream_ln as swiglu_stream_ln
 __all__ = [
     "attention_block", "attention_block_l1", "attention_block_l1_ln", "res_layernorm_fold",
-    "mlp", "mlp_res_ln", "swiglu_stream", "swiglu_stream_ln",
     "pack_weights", "pack_tables", "pack_v_weights", "pack_out_weights", "lane_permutation",
 ]

 from .attention_block_l1 import attention_block_l1
 from .attention_block_l1_ln import attention_block_l1_ln
 from .mlp import nki_mlp_k4_xp as mlp
+from .mlp import pack_b1
 from .mlp_res_ln import nki_mlp_res_ln as mlp_res_ln
 from .res_layernorm import nki_res_layernorm_fold as res_layernorm_fold
 from .swiglu_stream import nki_swiglu_stream as swiglu_stream
 __all__ = [
     "attention_block", "attention_block_l1", "attention_block_l1_ln", "res_layernorm_fold",
+    "mlp", "mlp_res_ln", "pack_b1", "swiglu_stream", "swiglu_stream_ln",
     "pack_weights", "pack_tables", "pack_v_weights", "pack_out_weights", "lane_permutation",
 ]

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

@@ -7,35 +7,46 @@ 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
@@ -50,8 +61,6 @@ _F_MOV_MAX = 512
 _BF16 = nl.bfloat16
-#: Logical-to-physical NeuronCore degree. Must match the shipping shard's.
-import os as _os
 #: logical-core degree the kernels are built for; must match NEURON_LOGICAL_NC_CONFIG.
 LNC = int(_os.environ.get("NEURON_LOGICAL_NC_CONFIG", "2"))
@@ -69,20 +78,23 @@ def _kernel_assert(condition: bool, error_text: str):
     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.
@@ -102,8 +114,15 @@ def _stage_xt(x_hbm, r_abs, pr, n_k, K, rb_size):
         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
@@ -114,7 +133,7 @@ def _mlp_k4_xp_kernel(x_hbm, w1t_hbm, b1_hbm, w2t_hbm, b2_hbm):
     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.
@@ -128,6 +147,9 @@ def _mlp_k4_xp_kernel(x_hbm, w1t_hbm, b1_hbm, w2t_hbm, b2_hbm):
     _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}. "
@@ -144,17 +166,26 @@ def _mlp_k4_xp_kernel(x_hbm, w1t_hbm, b1_hbm, w2t_hbm, b2_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
@@ -163,6 +194,13 @@ def _mlp_k4_xp_kernel(x_hbm, w1t_hbm, b1_hbm, w2t_hbm, b2_hbm):
         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
@@ -171,21 +209,6 @@ def _mlp_k4_xp_kernel(x_hbm, w1t_hbm, b1_hbm, w2t_hbm, b2_hbm):
         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)
@@ -212,10 +235,10 @@ def _mlp_k4_xp_kernel(x_hbm, w1t_hbm, b1_hbm, w2t_hbm, b2_hbm):
                     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] ----
@@ -259,6 +282,17 @@ def _mlp_k4_xp_kernel(x_hbm, w1t_hbm, b1_hbm, w2t_hbm, b2_hbm):
     return out_hbm
 def _no_grad(*tensors) -> None:
     """Refuse gradient-tracking inputs -- this op registers no autograd rule."""
     for t in tensors:
@@ -281,13 +315,17 @@ def nki_mlp_k4_xp(
     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)

 WHY IT IS SHAPED THIS WAY
 -------------------------
+This is the effective shipping kernel: the ``b1pre`` MLP (fc1 bias staged in the prologue
+so every gelu's bias is resident from t=0) with the x-prestage applied. The schedule was
+rebuilt from the device profile (external contribution by jlongen-aws, PR #2 on
+``jburtoft/vjepa2-neuron-kernels``). Four design choices carry the performance; the
+instruction order and numerics are unchanged, so the output is bit-identical to the
+previous version:
+* x BEFORE weights: row-block 0's ``x`` load + transpose is issued first, then W1, then the
+  later row blocks' ``x``, then W2. The first matmul no longer waits on ~16 MB of weights it
+  does not need yet (the first fc1 matmul used to retire at ~57 us).
+* b1 in ONE DMA on the Scalar engine's HW-DGE queue: ``b1`` is host-packed to ``[128, H/128]``
+  (:func:`pack_b1`) and lands in a single ~16 KB DMA instead of ``n_h`` ``[128, 1]`` fp32 DMAs
+  that trickled in and stalled each gelu (and, through its PSUM bank, the fc1 matmuls behind
+  it). Routing it on ``engine=scalar`` HW-DGE keeps gelu's wait off the W1/W2 traffic.
+* row-block 0's ``x`` is transposed on the TENSOR engine (``nc_transpose``) in the prologue,
+  while the Tensor engine is otherwise dark in the pre-K4 seam. Later row blocks keep
+  ``dma_transpose`` (they stream under rb0's matmuls).
+* x prestaged in the prologue: EVERY row block's ``x`` staging is issued before the ``rb``
+  loop, in row-block order, so a later block's serialised DMA chain runs inside the pre-K4
+  seam instead of in the idle window at the rb boundary. Appending later blocks after rb=0's
+  stages means rb=0's readiness cannot regress -- the first matmul cannot be delayed.
+The x staging is emitted through a module-level emitter (``_stage_xt``) so the ``x`` load,
+the ``nc_transpose`` and the ``dma_transpose`` each keep exactly one source site.
 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%). Weight DMAs stay
+one-per-tile: merging them into a single access-pattern DMA lost ~8 us (one completion signal
+per 8 MB weight).
+This kernel is LNC-parametric (``LNC`` from ``NEURON_LOGICAL_NC_CONFIG``) and keeps the
+``min()`` remainder handling, so it is NOT restricted to exact-multiple shapes. INFERENCE
+ONLY: no autograd rule.
 """
 from __future__ import annotations
 import math
+import os as _os
 import torch
 from torch_neuronx import nki_op
 _BF16 = nl.bfloat16
 #: logical-core degree the kernels are built for; must match NEURON_LOGICAL_NC_CONFIG.
 LNC = int(_os.environ.get("NEURON_LOGICAL_NC_CONFIG", "2"))
     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, tensor_transpose):
     """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`` transposes into the row
+    block's ``xt_tiles``. The transpose runs on the Tensor engine (``nc_transpose``, via a PSUM
+    tile) when ``tensor_transpose`` is set -- used for row block 0 while the Tensor engine is
+    dark in the pre-K4 seam -- and on the DMA engine (``dma_transpose``) otherwise. Kept as a
+    module-level emitter so the ``x`` load and each transpose keep a single source site.
     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).
+        tensor_transpose: transpose on the Tensor engine (True) or the DMA engine (False).
     Returns:
         list of ``n_k`` transposed ``[128, rb_size]`` bf16 SBUF tiles.
         for kt in range(n_k):
             k0 = kt * _P_MAX
             pk = min(_P_MAX, K - k0)
+            if tensor_transpose:
+                # Transpose on the (idle) Tensor engine via PSUM, then evict to the xt tile.
+                tp = nl.ndarray((_P_MAX, _P_MAX), dtype=_BF16, buffer=nl.psum)
+                nisa.nc_transpose(dst=tp[0:pk, 0:prs], data=x_b16[0:prs, k0:k0 + pk],
+                                  engine=nisa.engine.tensor)
+                nisa.tensor_copy(dst=xt_tiles[kt][0:pk, rr0:rr0 + prs], src=tp[0:pk, 0:prs])
+            else:
+                nisa.dma_transpose(dst=xt_tiles[kt][0:pk, rr0:rr0 + prs],
+                                   src=x_b16[0:prs, k0:k0 + pk])
     return xt_tiles
     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): [128, H/128] @ HBM, float32, == pack_b1(fc1 bias).
         w2t_hbm (nl.ndarray): [H, K] @ HBM, **bf16**, == fc2.weight.T.
         b2_hbm  (nl.ndarray): [K]    @ HBM, float32.
     _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(H % _P_MAX == 0, f"H={H} must be a multiple of {_P_MAX} for the packed b1 layout")
+    _kernel_assert(tuple(b1_hbm.shape) == (_P_MAX, H // _P_MAX),
+                   f"b1 must be pack_b1(fc1 bias): [{_P_MAX}, {H // _P_MAX}], got {tuple(b1_hbm.shape)}")
     _kernel_assert(
         R % LNC == 0,
         f"the {LNC}-core kernel needs R divisible by {LNC}, got R={R}. "
     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, in ONE DMA on the Scalar engine's HW-DGE queue: the packed [128, H/128] bias
+    # lands in a single ~16 KB transfer so every gelu's bias is resident long before the first
+    # matmul retires, and gelu's wait no longer counts the W1/W2 traffic. b1_sb[:, ht] is the
+    # bias column for H tile ht (b1_sb[p, ht] == fc1_bias[ht*128 + p]).
+    b1_sb = nl.ndarray((_P_MAX, n_h), dtype=nl.float32, buffer=nl.sbuf)
+    nisa.dma_copy(dst=b1_sb, src=b1_hbm, dge_mode=nisa.dge_mode.hwdge, engine=nisa.engine.scalar)
+    # Main fused loop geometry (needed to size the x prestage).
+    rb_size = _F_MOV_MAX
+    n_rb = _div_ceil(rows_per, rb_size)
+    n_ko = _div_ceil(K, _F_MOV_MAX)
+    # x BEFORE weights: stage row block 0's x first (Tensor-engine transpose, filling the dark
+    # pre-K4 seam), then W1, then later row blocks' x (DMA transpose, streaming under rb0's
+    # matmuls), then W2. Appending later blocks after rb=0's stages leaves rb=0's readiness fixed.
+    xt_sets = [None] * n_rb
+    xt_sets[0] = _stage_xt(x_hbm, r_base, min(rb_size, rows_per), n_k, K, rb_size,
+                           tensor_transpose=True)
+    # W1 staging (pre-transposed bf16, one DMA per K tile).
     w1t_tiles = []
     for kt in range(n_k):
         k0 = kt * _P_MAX
         nisa.dma_copy(dst=t[0:pk, :], src=w1t_hbm[k0:k0 + pk, :])
         w1t_tiles.append(t)
+    # Later row blocks' x via DMA transpose, under rb0's matmuls.
+    for rb_pre in range(1, n_rb):
+        r_pre = rb_pre * rb_size
+        xt_sets[rb_pre] = _stage_xt(x_hbm, r_base + r_pre, min(rb_size, rows_per - r_pre),
+                                    n_k, K, rb_size, tensor_transpose=False)
+    # W2 staging (one DMA per H tile, so fc2 can start on the first strips).
     w2t_tiles = []
     for ht in range(n_h):
         h0 = ht * _P_MAX
         nisa.dma_copy(dst=t[0:ph, :], src=w2t_hbm[h0:h0 + ph, :])
         w2t_tiles.append(t)
     for rb in range(n_rb):
         r_local = rb * rb_size
         pr = min(rb_size, rows_per - r_local)
                     stationary=w1t_tiles[kt][0:pk, h0:h0 + ph],
                     moving=xt_tiles[kt][0:pk, 0:pr],
                 )
+            # b1 comes from the prologue-resident packed tile, column ht.
             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_sb[0:ph, ht:ht + 1])
             gt_tiles.append(gt)
         # ---- fc2: out[r,k] = sum_h g[r,h] * w2[k,h] + b2[k] ----
     return out_hbm
+def pack_b1(b1: torch.Tensor) -> torch.Tensor:
+    """Pack fc1 bias ``[H]`` fp32 into ``[128, H/128]`` so ``b1p[p, ht] == b1[ht*128 + p]``.
+    Lets the kernel load the whole bias in one DMA instead of ``H/128`` column DMAs. ``H`` must
+    be a multiple of 128.
+    """
+    H = b1.shape[0]
+    assert H % _P_MAX == 0, f"pack_b1 needs H % {_P_MAX} == 0, got H={H}"
+    return b1.reshape(-1, _P_MAX).t().contiguous()
 def _no_grad(*tensors) -> None:
     """Refuse gradient-tracking inputs -- this op registers no autograd rule."""
     for t in tensors:
     Args:
         x   : [R, K] **bf16** on 'neuron'. ``R`` must be divisible by :data:`LNC`.
+        w1t : [K, H] bf16 == fc1.weight.T.
+        b1  : fc1 bias, either ``[H]`` float32 (packed on the fly) or pre-packed
+              ``[128, H/128]`` from :func:`pack_b1`.
         w2t : [H, K] bf16 == fc2.weight.T.   b2 : [K] float32.
     Returns:
         [R, K] bf16.
     """
     _no_grad(x, w1t, b1, w2t, b2)
+    if b1.dim() == 1:                       # accept the pre-pack [H] bias; pack it on the fly
+        b1 = pack_b1(b1)
     return _mlp_k4_xp_kernel[LNC](x, w1t, b1, w2t, b2)

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

@@ -1,50 +1,25 @@
 """GELU MLP with the residual add + NEXT layer's affine-free LN1 in the epilogue.
-Same fc1/gelu as ``mlp.py``. fc2 is reordered to (row subtile rs, out block ko) so each 128-row tile's
-full [128, K] fc2 output lands in one fp32 SBUF row, then:
     x_new = x + (g @ w2.T + b2)            fp32 (MLP output never rounded to bf16)
     z1    = LN_no_affine(x_new; eps)        bf16 (same stat chain as res_layernorm.py)
 Returns (x_new [R,K] fp32, z1 [R,K] bf16). Replaces K4 + the next layer's K1.
---- original mlp.py notes ---
-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
@@ -53,14 +28,14 @@ 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.
-import os as _os
 #: logical-core degree the kernels are built for; must match NEURON_LOGICAL_NC_CONFIG.
 LNC = int(_os.environ.get("NEURON_LOGICAL_NC_CONFIG", "2"))
@@ -78,44 +53,6 @@ def _kernel_assert(condition: bool, error_text: str):
     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_res_ln_kernel(x_hbm, w1t_hbm, b1_hbm, w2t_hbm, b2_hbm, a_hbm, eps):
     """gelu(x @ w1.T + b1) @ w2.T + b2, rows sharded, b1 AND both x blocks resident early.
@@ -123,7 +60,7 @@ def _mlp_res_ln_kernel(x_hbm, w1t_hbm, b1_hbm, w2t_hbm, b2_hbm, a_hbm, eps):
     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.
@@ -137,6 +74,9 @@ def _mlp_res_ln_kernel(x_hbm, w1t_hbm, b1_hbm, w2t_hbm, b2_hbm, a_hbm, eps):
     _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}. "
@@ -157,17 +97,25 @@ def _mlp_res_ln_kernel(x_hbm, w1t_hbm, b1_hbm, w2t_hbm, b2_hbm, a_hbm, eps):
     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
@@ -176,6 +124,12 @@ def _mlp_res_ln_kernel(x_hbm, w1t_hbm, b1_hbm, w2t_hbm, b2_hbm, a_hbm, eps):
         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
@@ -193,21 +147,6 @@ def _mlp_res_ln_kernel(x_hbm, w1t_hbm, b1_hbm, w2t_hbm, b2_hbm, a_hbm, eps):
     var = nl.ndarray((_P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf)
     rstd = nl.ndarray((_P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf)
-    # 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)
@@ -234,10 +173,10 @@ def _mlp_res_ln_kernel(x_hbm, w1t_hbm, b1_hbm, w2t_hbm, b2_hbm, a_hbm, eps):
                     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 + b2 + residual -> x_new (fp32), LN -> z1 (bf16), per 128-row subtile ----
@@ -303,6 +242,9 @@ def nki_mlp_res_ln(
     eps: float,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """MLP(x) + residual a, then affine-free LN. x [R,K] bf16, a [R,K] fp32.
     Returns (x_new [R,K] fp32, z [R,K] bf16)."""
     _no_grad(x, w1t, b1, w2t, b2, a)
     return _mlp_res_ln_kernel[LNC](x, w1t, b1, w2t, b2, a, eps)

 """GELU MLP with the residual add + NEXT layer's affine-free LN1 in the epilogue.
+Same fc1/gelu as ``mlp.py`` (including its reworked prologue, external contribution by
+jlongen-aws, PR #2 on ``jburtoft/vjepa2-neuron-kernels``). fc2 is reordered to (row subtile
+rs, out block ko) so each 128-row tile's full [128, K] fc2 output lands in one fp32 SBUF row,
+then:
     x_new = x + (g @ w2.T + b2)            fp32 (MLP output never rounded to bf16)
     z1    = LN_no_affine(x_new; eps)        bf16 (same stat chain as res_layernorm.py)
 Returns (x_new [R,K] fp32, z1 [R,K] bf16). Replaces K4 + the next layer's K1.
+Prologue (shared with mlp.py): b1 packed to [128, H/128] and loaded in one DMA on the Scalar
+engine's HW-DGE queue; row block 0's x transposed on the (idle) Tensor engine, later blocks on
+the DMA engine, all x staged before the weights so the first fc1 matmul does not wait on W1/W2.
+Instruction order and numerics are unchanged from the previous version.
+There are no ``engine=`` pins in the matmul/epilogue. INFERENCE ONLY: no autograd rule.
 """
 from __future__ import annotations
 import math
+import os as _os
 import torch
 from torch_neuronx import nki_op
 import nki.isa as nisa
 import nki.language as nl
+from .mlp import _stage_xt, pack_b1
 #: PE array partition limit, and the moving-operand free-axis limit.
 _P_MAX = 128
 _F_MOV_MAX = 512
 _BF16 = nl.bfloat16
 #: logical-core degree the kernels are built for; must match NEURON_LOGICAL_NC_CONFIG.
 LNC = int(_os.environ.get("NEURON_LOGICAL_NC_CONFIG", "2"))
     assert condition, f"[INTERNAL_ERROR] [NCC_INKI016] Kernel validation exception: {error_text}"
 @nki.jit
 def _mlp_res_ln_kernel(x_hbm, w1t_hbm, b1_hbm, w2t_hbm, b2_hbm, a_hbm, eps):
     """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): [128, H/128] @ HBM, float32, == pack_b1(fc1 bias).
         w2t_hbm (nl.ndarray): [H, K] @ HBM, **bf16**, == fc2.weight.T.
         b2_hbm  (nl.ndarray): [K]    @ HBM, float32.
     _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(H % _P_MAX == 0, f"H={H} must be a multiple of {_P_MAX} for the packed b1 layout")
+    _kernel_assert(tuple(b1_hbm.shape) == (_P_MAX, H // _P_MAX),
+                   f"b1 must be pack_b1(fc1 bias): [{_P_MAX}, {H // _P_MAX}], got {tuple(b1_hbm.shape)}")
     _kernel_assert(
         R % LNC == 0,
         f"the {LNC}-core kernel needs R divisible by {LNC}, got R={R}. "
     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)
+    # Main fused loop geometry (needed to size the x prestage).
+    rb_size = _F_MOV_MAX
+    n_rb = _div_ceil(rows_per, rb_size)
+    n_ko = _div_ceil(K, _F_MOV_MAX)
+    # b1 FIRST, in ONE DMA on the Scalar engine's HW-DGE queue (packed [128, H/128]): the bias
+    # is resident long before the first matmul retires and gelu's wait no longer counts W1/W2
+    # traffic. b1_sb[:, ht] is the bias column for H tile ht.
+    b1_sb = nl.ndarray((_P_MAX, n_h), dtype=nl.float32, buffer=nl.sbuf)
+    nisa.dma_copy(dst=b1_sb, src=b1_hbm, dge_mode=nisa.dge_mode.hwdge, engine=nisa.engine.scalar)
+    # x BEFORE weights: rb0's x first (Tensor-engine transpose, filling the dark pre-K4 seam),
+    # then W1, then later row blocks' x (DMA transpose), then W2. Appending later blocks after
+    # rb0 leaves rb0's readiness fixed, so the first matmul cannot be delayed.
+    xt_sets = [None] * n_rb
+    xt_sets[0] = _stage_xt(x_hbm, r_base, min(rb_size, rows_per), n_k, K, rb_size,
+                           tensor_transpose=True)
+    # W1 staging (pre-transposed bf16, one DMA per K tile).
     w1t_tiles = []
     for kt in range(n_k):
         k0 = kt * _P_MAX
         nisa.dma_copy(dst=t[0:pk, :], src=w1t_hbm[k0:k0 + pk, :])
         w1t_tiles.append(t)
+    # Later row blocks' x via DMA transpose, under rb0's matmuls.
+    for rb_pre in range(1, n_rb):
+        r_pre = rb_pre * rb_size
+        xt_sets[rb_pre] = _stage_xt(x_hbm, r_base + r_pre, min(rb_size, rows_per - r_pre),
+                                    n_k, K, rb_size, tensor_transpose=False)
     w2t_tiles = []
     for ht in range(n_h):
         h0 = ht * _P_MAX
     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 rb in range(n_rb):
         r_local = rb * rb_size
         pr = min(rb_size, rows_per - r_local)
                     stationary=w1t_tiles[kt][0:pk, h0:h0 + ph],
                     moving=xt_tiles[kt][0:pk, 0:pr],
                 )
+            # b1 comes from the prologue-resident packed tile, column ht.
             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_sb[0:ph, ht:ht + 1])
             gt_tiles.append(gt)
         # ---- fc2 + b2 + residual -> x_new (fp32), LN -> z1 (bf16), per 128-row subtile ----
     eps: float,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """MLP(x) + residual a, then affine-free LN. x [R,K] bf16, a [R,K] fp32.
+    b1 is fc1 bias, either [H] float32 (packed on the fly) or pre-packed [128, H/128].
     Returns (x_new [R,K] fp32, z [R,K] bf16)."""
     _no_grad(x, w1t, b1, w2t, b2, a)
+    if b1.dim() == 1:
+        b1 = pack_b1(b1)
     return _mlp_res_ln_kernel[LNC](x, w1t, b1, w2t, b2, a, eps)

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

@@ -19,9 +19,16 @@ 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
@@ -30,10 +37,11 @@ WHY IT IS SHAPED THIS WAY
   (neg_shift -> copy-reduce -> neg_mu -> square-reduce -> var -> rsqrt) that keeps the
   per-lane statistics in fp32; the final normalize is a single fused ``tensor_scalar``
   (``z = (x_new + neg_mu) * rstd``) whose two per-lane operands cost the same as one op
-  and whose bf16 cast rides the write.
-There are no ``engine=`` pins anywhere in this kernel; nothing engine-related to
-preserve. INFERENCE ONLY: no autograd rule is registered.
 """
 from __future__ import annotations
@@ -56,9 +64,6 @@ import os as _os
 #: logical-core degree the kernels are built for; must match NEURON_LOGICAL_NC_CONFIG.
 LNC = int(_os.environ.get("NEURON_LOGICAL_NC_CONFIG", "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."""
@@ -70,9 +75,10 @@ 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:
         a_hbm (nl.ndarray): [R, H] @ HBM float32 -- the residual stream.
@@ -108,69 +114,67 @@ 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]) 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. (External
+  contribution by jlongen-aws, PR #2 on ``jburtoft/vjepa2-neuron-kernels``; replaces the
+  depth-4 ring, same instructions in the same order so results are unchanged.)
+* ``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, and the first add waited ~8 us for
+  them. 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, so it is not used.
 * 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
   (neg_shift -> copy-reduce -> neg_mu -> square-reduce -> var -> rsqrt) that keeps the
   per-lane statistics in fp32; the final normalize is a single fused ``tensor_scalar``
   (``z = (x_new + neg_mu) * rstd``) whose two per-lane operands cost the same as one op
+  and whose bf16 cast rides the write. Each per-lane statistic has its own column per tile
+  (``[128, n_t]``) so tile ``t+1``'s chain does not wait on tile ``t``'s.
+This kernel is LNC-parametric (``LNC`` from ``NEURON_LOGICAL_NC_CONFIG``). There are no
+``engine=`` pins. INFERENCE ONLY: no autograd rule is registered.
 """
 from __future__ import annotations
 #: logical-core degree the kernels are built for; must match NEURON_LOGICAL_NC_CONFIG.
 LNC = int(_os.environ.get("NEURON_LOGICAL_NC_CONFIG", "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 (all loads issued before any compute), summed in fp32, normalized
+    with the thinned stat chain, and the normalized result is cast to bf16 on write. Every DMA
+    uses ``dge_mode.none`` so the runtime pre-generates the descriptors.
     Args:
         a_hbm (nl.ndarray): [R, H] @ HBM float32 -- the residual stream.
     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[:, t, :] is tile t's residual, b_sb[:, j, t, :] is branch j of tile t.
+    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.
+        # Both operands are per-lane vectors; the fused two-op tensor_scalar costs one op.
+        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