v1.2: packed-bias one-DMA for SwiGLU (swiglu_stream/_ln)

Silu bias b1 and up bias bg host-packed to [128, I/128] and loaded in one DMA each on the
Scalar HW-DGE queue (same technique as v1.1 K4). Bitwise-identical to v1.1 at LNC=1 and LNC=2;
swiglu_stream -2.5% device cycles, swiglu_stream_ln within noise. GELU path unchanged.

Files changed (3) hide show

CHANGELOG.md +11 -0
build/torch-neuron/encoder.py +2 -2
build/torch-neuron/nki_kernels/swiglu_stream.py +25 -16

CHANGELOG.md CHANGED Viewed

@@ -1,5 +1,16 @@
 # Changelog
 ## v1.1 (2026-10)
 Reworked DMA/engine schedule for `mlp` (K4), `mlp_res_ln` (K4 + next LN) and `res_layernorm_fold`

 # Changelog
+## v1.2 (2026-10)
+Applied the v1.1 packed-bias one-DMA technique to the SwiGLU FFN kernels (`swiglu_stream`,
+`swiglu_stream_ln`). The silu bias `b1` and the up-projection bias `bg` are now host-packed to
+`[128, I/128]` (`pack_b1`) and each loaded in a single DMA on the Scalar engine's HW-DGE queue,
+replacing `I/128` `[128, 1]` column DMAs apiece. Output is **bitwise identical** to v1.1 at LNC=1 and
+LNC=2 (`torch.equal`, verified on trn2.3xlarge). Device cycles (same stack, A/B): `swiglu_stream`
+-2.5% (reproducible); `swiglu_stream_ln` within measurement noise (the LN epilogue dominates and the
+bias DMAs already overlap compute). The wrappers still accept an unpacked `[I]` bias (packed on the
+fly). Affects only the gated-MLP (SwiGLU) models; the GELU-MLP path is unchanged from v1.1.
 ## v1.1 (2026-10)
 Reworked DMA/engine schedule for `mlp` (K4), `mlp_res_ln` (K4 + next LN) and `res_layernorm_fold`

build/torch-neuron/encoder.py CHANGED Viewed

@@ -124,9 +124,9 @@ def pack_block(bw: BlockWeights, num_heads: int, rope: str, device) -> dict:
         pi = lambda t: torch.cat([t, t.new_zeros(Ip - t.shape[0], *t.shape[1:])], 0)
         p["swiglu"] = True
         p["s_w1t"] = pi(gw * g2[None]).t().contiguous().to(torch.bfloat16)          # gate (silu)
-        p["s_b1"] = pi(_z(_f(bw.gate_b), I) + gw @ t2).contiguous()
         p["s_wgt"] = pi(uw * g2[None]).t().contiguous().to(torch.bfloat16)          # up
-        p["s_bg"] = pi(_z(_f(bw.up_b), I) + uw @ t2).contiguous()
         p["s_w2t"] = pi((dw * ls2[:, None]).t().contiguous()).contiguous().to(torch.bfloat16)
         p["s_b2"] = (_z(_f(bw.down_b), H) * ls2).contiguous()
     else:                                                                     # GELU MLP

         pi = lambda t: torch.cat([t, t.new_zeros(Ip - t.shape[0], *t.shape[1:])], 0)
         p["swiglu"] = True
         p["s_w1t"] = pi(gw * g2[None]).t().contiguous().to(torch.bfloat16)          # gate (silu)
+        p["s_b1"] = _mlp_pack_b1(pi(_z(_f(bw.gate_b), I) + gw @ t2).contiguous())   # [128, Ip/128], one-DMA
         p["s_wgt"] = pi(uw * g2[None]).t().contiguous().to(torch.bfloat16)          # up
+        p["s_bg"] = _mlp_pack_b1(pi(_z(_f(bw.up_b), I) + uw @ t2).contiguous())     # [128, Ip/128], one-DMA
         p["s_w2t"] = pi((dw * ls2[:, None]).t().contiguous()).contiguous().to(torch.bfloat16)
         p["s_b2"] = (_z(_f(bw.down_b), H) * ls2).contiguous()
     else:                                                                     # GELU MLP

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

@@ -28,6 +28,8 @@ import nki
 import nki.isa as nisa
 import nki.language as nl
 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"))
@@ -108,18 +110,17 @@ def _body(x_hbm, w1t_hbm, b1_hbm, wgt_hbm, bg_hbm, w2t_hbm, b2_hbm, a_hbm, eps,
     else:
         out_hbm = nl.ndarray((R, K), dtype=_BF16, buffer=nl.shared_hbm)
-    # biases (tiny, resident): b1/bg per 128-row tile of I, b2 broadcast over partitions.
-    b1_t = []
-    bg_t = []
-    for ht in range(n_h):
-        h0 = ht * _P_MAX
-        ph = min(_P_MAX, I - h0)
-        t1 = nl.ndarray((_P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf)
-        nisa.dma_copy(dst=t1[0:ph, 0:1], src=b1_hbm[h0:h0 + ph])
-        b1_t.append(t1)
-        tg = nl.ndarray((_P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf)
-        nisa.dma_copy(dst=tg[0:ph, 0:1], src=bg_hbm[h0:h0 + ph])
-        bg_t.append(tg)
     b2_bc = nl.ndarray((_P_MAX, K), dtype=nl.float32, buffer=nl.sbuf)
     nisa.dma_copy(dst=b2_bc[0:_P_MAX, 0:K], src=b2_hbm.ap(pattern=[[0, _P_MAX], [1, K]], offset=0))
@@ -187,11 +188,11 @@ def _body(x_hbm, w1t_hbm, b1_hbm, wgt_hbm, bg_hbm, w2t_hbm, b2_hbm, a_hbm, eps,
                                        moving=xt[kt][0:pk, q0:q0 + rw])
                     sa = nl.ndarray((_P_MAX, _F_MOV_MAX), dtype=nl.float32, buffer=nl.sbuf)
                     nisa.activation(dst=sa[0:ph, 0:rw], op=nl.silu, data=y1[0:ph, 0:rw],
-                                    bias=b1_t[hg][0:ph, 0:1])
                     g = nl.ndarray((_P_MAX, _F_MOV_MAX), dtype=_BF16, buffer=nl.sbuf)
                     # g = (yg + bg) * silu(y1 + b1)   -- one fused op, PSUM operand
                     nisa.scalar_tensor_tensor(dst=g[0:ph, 0:rw], data=yg[0:ph, 0:rw],
-                                              op0=nl.add, operand0=bg_t[hg][0:ph, 0:1],
                                               op1=nl.multiply, operand1=sa[0:ph, 0:rw])
                     g_t.append(g)
                 # ---- down-projection partial for this chunk, accumulated into acc (fp32 SBUF) ----
@@ -259,6 +260,14 @@ def _swiglu_stream_ln_kernel(x_hbm, w1t_hbm, b1_hbm, wgt_hbm, bg_hbm, w2t_hbm, b
     return _body(x_hbm, w1t_hbm, b1_hbm, wgt_hbm, bg_hbm, w2t_hbm, b2_hbm, a_hbm, eps, S, CI, 1)
 def _no_grad(*ts):
     for t in ts:
         if t is not None and t.requires_grad:
@@ -278,7 +287,7 @@ def nki_swiglu_stream(x: torch.Tensor, w1t: torch.Tensor, b1: torch.Tensor, wgt:
     """x [R,K] bf16; w1t/wgt [K,I] bf16; b1/bg [I] f32; w2t [I,K] bf16; b2 [K] f32 -> [R,K] bf16."""
     _no_grad(x, w1t, b1, wgt, bg, w2t, b2)
     S, CI = _plan_for(x, w1t)
-    return _swiglu_stream_kernel[LNC](x, w1t, b1, wgt, bg, w2t, b2, S, CI)
 @nki_op("vit_enc::nki_swiglu_stream_ln", mutates_args={})
@@ -288,4 +297,4 @@ def nki_swiglu_stream_ln(x: torch.Tensor, w1t: torch.Tensor, b1: torch.Tensor, w
     """SwiGLU(x) + residual a, then affine-free LN. Returns (x_new [R,K] fp32, z [R,K] bf16)."""
     _no_grad(x, w1t, b1, wgt, bg, w2t, b2, a)
     S, CI = _plan_for(x, w1t)
-    return _swiglu_stream_ln_kernel[LNC](x, w1t, b1, wgt, bg, w2t, b2, a, eps, S, CI)

 import nki.isa as nisa
 import nki.language as nl
+from .mlp import pack_b1
 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"))
     else:
         out_hbm = nl.ndarray((R, K), dtype=_BF16, buffer=nl.shared_hbm)
+    # biases (tiny, resident). b1/bg are host-packed to [128, I/128] (pack_bias; b[p, ht] ==
+    # bias[ht*128 + p]) and each lands in ONE DMA on the Scalar HW-DGE queue -- replaces 2 * I/128
+    # [128, 1] DMAs (~54 per core at ViT-H+) that stalled the first silu/gate ops. b2 broadcast over
+    # partitions.
+    _kernel_assert(I % _P_MAX == 0, "I must be padded to a multiple of 128 (pack_bias)")
+    _kernel_assert(tuple(b1_hbm.shape) == (_P_MAX, n_h) and tuple(bg_hbm.shape) == (_P_MAX, n_h),
+                   f"b1/bg must be pack_bias(...): [{_P_MAX}, {n_h}]")
+    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)
+    bg_sb = nl.ndarray((_P_MAX, n_h), dtype=nl.float32, buffer=nl.sbuf)
+    nisa.dma_copy(dst=bg_sb, src=bg_hbm, dge_mode=nisa.dge_mode.hwdge, engine=nisa.engine.scalar)
     b2_bc = nl.ndarray((_P_MAX, K), dtype=nl.float32, buffer=nl.sbuf)
     nisa.dma_copy(dst=b2_bc[0:_P_MAX, 0:K], src=b2_hbm.ap(pattern=[[0, _P_MAX], [1, K]], offset=0))
                                        moving=xt[kt][0:pk, q0:q0 + rw])
                     sa = nl.ndarray((_P_MAX, _F_MOV_MAX), dtype=nl.float32, buffer=nl.sbuf)
                     nisa.activation(dst=sa[0:ph, 0:rw], op=nl.silu, data=y1[0:ph, 0:rw],
+                                    bias=b1_sb[0:ph, hg:hg + 1])
                     g = nl.ndarray((_P_MAX, _F_MOV_MAX), dtype=_BF16, buffer=nl.sbuf)
                     # g = (yg + bg) * silu(y1 + b1)   -- one fused op, PSUM operand
                     nisa.scalar_tensor_tensor(dst=g[0:ph, 0:rw], data=yg[0:ph, 0:rw],
+                                              op0=nl.add, operand0=bg_sb[0:ph, hg:hg + 1],
                                               op1=nl.multiply, operand1=sa[0:ph, 0:rw])
                     g_t.append(g)
                 # ---- down-projection partial for this chunk, accumulated into acc (fp32 SBUF) ----
     return _body(x_hbm, w1t_hbm, b1_hbm, wgt_hbm, bg_hbm, w2t_hbm, b2_hbm, a_hbm, eps, S, CI, 1)
+def pack_bias(b: torch.Tensor) -> torch.Tensor:
+    """Pack an fp32 bias ``[I]`` (I % 128 == 0) to ``[128, I/128]`` so ``b_p[p, t] == b[t*128 + p]``."""
+    if b.dim() == 2:
+        return b
+    assert b.shape[0] % _P_MAX == 0, f"pack_bias needs I % {_P_MAX} == 0, got {b.shape[0]}"
+    return b.reshape(-1, _P_MAX).t().contiguous()
 def _no_grad(*ts):
     for t in ts:
         if t is not None and t.requires_grad:
     """x [R,K] bf16; w1t/wgt [K,I] bf16; b1/bg [I] f32; w2t [I,K] bf16; b2 [K] f32 -> [R,K] bf16."""
     _no_grad(x, w1t, b1, wgt, bg, w2t, b2)
     S, CI = _plan_for(x, w1t)
+    return _swiglu_stream_kernel[LNC](x, w1t, pack_bias(b1), wgt, pack_bias(bg), w2t, b2, S, CI)
 @nki_op("vit_enc::nki_swiglu_stream_ln", mutates_args={})
     """SwiGLU(x) + residual a, then affine-free LN. Returns (x_new [R,K] fp32, z [R,K] bf16)."""
     _no_grad(x, w1t, b1, wgt, bg, w2t, b2, a)
     S, CI = _plan_for(x, w1t)
+    return _swiglu_stream_ln_kernel[LNC](x, w1t, pack_bias(b1), wgt, pack_bias(bg), w2t, b2, a, eps, S, CI)