v1.3: attention-path schedule improvements (bitwise identical to v1.2) + SwiGLU fallback fix
Browse filesQ/K bias host-packed [2,128,Fp/128] and loaded in one Scalar HW-DGE DMA per Q/K (pack_qk_bias);
dead zero-fills and the chained exp-sum tensor_reduce removed in the single-section attention core;
row-max chained across key tiles. Attention kernels -1.7% to -11.1% device cycles on trn2.3xlarge,
output torch.equal to v1.2 at LNC=1 and LNC=2. Fixes a v1.2 regression where the SwiGLU torch
fallback (rows/core < 256) used packed biases and failed to compile.
- CHANGELOG.md +39 -0
- build/torch-neuron/encoder.py +7 -4
- build/torch-neuron/nki_kernels/__init__.py +2 -1
- build/torch-neuron/nki_kernels/attention_block.py +37 -23
- build/torch-neuron/nki_kernels/attention_block_l1.py +17 -18
- build/torch-neuron/nki_kernels/attention_block_l1_ln.py +17 -18
- build/torch-neuron/nki_kernels/attention_inner.py +57 -13
CHANGELOG.md
CHANGED
|
@@ -1,5 +1,44 @@
|
|
| 1 |
# Changelog
|
| 2 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 3 |
## v1.2 (2026-10)
|
| 4 |
|
| 5 |
Applied the v1.1 packed-bias one-DMA technique to the SwiGLU FFN kernels (`swiglu_stream`,
|
|
|
|
| 1 |
# Changelog
|
| 2 |
|
| 3 |
+
## v1.3 (2026-10)
|
| 4 |
+
|
| 5 |
+
Attention-path schedule improvements, plus a SwiGLU fallback fix. Output of every attention kernel is
|
| 6 |
+
**bitwise identical** to v1.2. Verified on trn2.3xlarge at LNC=1 and LNC=2, PyTorch Native Beta 6.
|
| 7 |
+
|
| 8 |
+
Device-cycle A/B, v1.2 vs v1.3 (median of 3 captures, same instance, container and inputs):
|
| 9 |
+
|
| 10 |
+
| Kernel / shape | LNC | change |
|
| 11 |
+
|---|---|---|
|
| 12 |
+
| `attention_block`, N=1024, B=2 | 2 | -1.9% |
|
| 13 |
+
| `attention_block`, N=2048, B=1 (V-JEPA2 clip) | 2 | **-5.2%** |
|
| 14 |
+
| `attention_block`, N=256, B=8 | 2 | -2.0% |
|
| 15 |
+
| `attention_block`, N=1024, B=2 | 1 | -3.3% |
|
| 16 |
+
| `attention_block_l1`, N=256, B=4 | 1 | **-11.1%** |
|
| 17 |
+
| `attention_block_l1_ln`, N=256, B=4 | 1 | **-10.2%** |
|
| 18 |
+
| `attention_block_l1_ln`, N=1024, B=2 | 1 | -4.5% |
|
| 19 |
+
|
| 20 |
+
End to end (real Hugging Face weights, single core, vs v1.0 on the same harness): DINOv3 ViT-L/16 +1.0% to
|
| 21 |
+
+12.7%, ViT-H+/16 +0.6% to +17.3%. Largest gains are at 224 px and batch 1, where attention's fixed costs
|
| 22 |
+
dominate. V-JEPA2 (16 frames, LNC=2): 24.78 -> 24.48 ms/clip vs v1.2, output identical to v1.2.
|
| 23 |
+
All 16 DINOv3 cases pass rel_l2 <= 5e-2 vs fp32 (0.56-0.98e-2).
|
| 24 |
+
|
| 25 |
+
Changes:
|
| 26 |
+
|
| 27 |
+
* **Q/K bias in one DMA.** `pack_weights` now returns the Q/K bias packed as `[2, 128, Fp/128]` fp32
|
| 28 |
+
(`pack_qk_bias`, exported). Each attention program loads its groups with one DMA per Q/K on the Scalar
|
| 29 |
+
HW-DGE queue, replacing one `[128, 1]` DMA plus one cast per 128-row group. The values are rounded to bf16
|
| 30 |
+
first, which keeps the output bit-identical. The kernels still accept the legacy `[2, Fp, 1]` layout.
|
| 31 |
+
* **Dead work removed in the attention core** (`attention_inner.py`, single-section path, N <= 2048):
|
| 32 |
+
the `mm2_sb` and `exp_partial_sum` zero-fills (never observed), and the `tensor_reduce` over the chained
|
| 33 |
+
exp-sum columns (the sum is already in one column; it is now read directly).
|
| 34 |
+
* **Chained row-max reduce**: the QK^T row-max accumulates across the 512-key tiles of a large tile and is
|
| 35 |
+
read out once, instead of once per tile followed by a `tensor_reduce`. Pattern from the upstream nkilib
|
| 36 |
+
`attention_cte` row-max chaining, re-implemented for this fork. It adds -3.5% at N=2048 and is neutral at
|
| 37 |
+
N <= 1024 (fewer key tiles to chain).
|
| 38 |
+
* **Fix (v1.2 regression):** the SwiGLU torch fallback (taken below 256 rows per core, e.g. ViT-H+/16 at
|
| 39 |
+
LNC=2, batch 1, 224 px) used the packed biases and failed to compile. It now uses unpacked copies. The
|
| 40 |
+
kernel path was unaffected.
|
| 41 |
+
|
| 42 |
## v1.2 (2026-10)
|
| 43 |
|
| 44 |
Applied the v1.1 packed-bias one-DMA technique to the SwiGLU FFN kernels (`swiglu_stream`,
|
build/torch-neuron/encoder.py
CHANGED
|
@@ -124,9 +124,11 @@ def pack_block(bw: BlockWeights, num_heads: int, rope: str, device) -> dict:
|
|
| 124 |
pi = lambda t: torch.cat([t, t.new_zeros(Ip - t.shape[0], *t.shape[1:])], 0)
|
| 125 |
p["swiglu"] = True
|
| 126 |
p["s_w1t"] = pi(gw * g2[None]).t().contiguous().to(torch.bfloat16) # gate (silu)
|
| 127 |
-
p["
|
|
|
|
| 128 |
p["s_wgt"] = pi(uw * g2[None]).t().contiguous().to(torch.bfloat16) # up
|
| 129 |
-
p["
|
|
|
|
| 130 |
p["s_w2t"] = pi((dw * ls2[:, None]).t().contiguous()).contiguous().to(torch.bfloat16)
|
| 131 |
p["s_b2"] = (_z(_f(bw.down_b), H) * ls2).contiguous()
|
| 132 |
else: # GELU MLP
|
|
@@ -197,8 +199,9 @@ class EncoderRunner:
|
|
| 197 |
if z2.shape[0] // LNC >= _SW_MIN_ROWS:
|
| 198 |
return nki_swiglu_stream(z2, p["s_w1t"], p["s_b1"], p["s_wgt"], p["s_bg"],
|
| 199 |
p["s_w2t"], p["s_b2"]).unsqueeze(0)
|
| 200 |
-
|
| 201 |
-
|
|
|
|
| 202 |
return (((g * u).to(torch.bfloat16) @ p["s_w2t"]).float() + p["s_b2"]).to(torch.bfloat16).unsqueeze(0)
|
| 203 |
return nki_mlp_k4_xp(z2, p["w1t"], p["b1"], p["w2t"], p["b2"]).unsqueeze(0)
|
| 204 |
|
|
|
|
| 124 |
pi = lambda t: torch.cat([t, t.new_zeros(Ip - t.shape[0], *t.shape[1:])], 0)
|
| 125 |
p["swiglu"] = True
|
| 126 |
p["s_w1t"] = pi(gw * g2[None]).t().contiguous().to(torch.bfloat16) # gate (silu)
|
| 127 |
+
p["s_b1u"] = pi(_z(_f(bw.gate_b), I) + gw @ t2).contiguous() # [Ip] for the torch fallback
|
| 128 |
+
p["s_b1"] = _mlp_pack_b1(p["s_b1u"]) # [128, Ip/128], one-DMA
|
| 129 |
p["s_wgt"] = pi(uw * g2[None]).t().contiguous().to(torch.bfloat16) # up
|
| 130 |
+
p["s_bgu"] = pi(_z(_f(bw.up_b), I) + uw @ t2).contiguous() # [Ip] for the torch fallback
|
| 131 |
+
p["s_bg"] = _mlp_pack_b1(p["s_bgu"]) # [128, Ip/128], one-DMA
|
| 132 |
p["s_w2t"] = pi((dw * ls2[:, None]).t().contiguous()).contiguous().to(torch.bfloat16)
|
| 133 |
p["s_b2"] = (_z(_f(bw.down_b), H) * ls2).contiguous()
|
| 134 |
else: # GELU MLP
|
|
|
|
| 199 |
if z2.shape[0] // LNC >= _SW_MIN_ROWS:
|
| 200 |
return nki_swiglu_stream(z2, p["s_w1t"], p["s_b1"], p["s_wgt"], p["s_bg"],
|
| 201 |
p["s_w2t"], p["s_b2"]).unsqueeze(0)
|
| 202 |
+
# torch fallback needs the UNPACKED [Ip] biases (v1.2 regression: it used the packed ones).
|
| 203 |
+
g = torch.nn.functional.silu((z2 @ p["s_w1t"]).float() + p["s_b1u"])
|
| 204 |
+
u = (z2 @ p["s_wgt"]).float() + p["s_bgu"]
|
| 205 |
return (((g * u).to(torch.bfloat16) @ p["s_w2t"]).float() + p["s_b2"]).to(torch.bfloat16).unsqueeze(0)
|
| 206 |
return nki_mlp_k4_xp(z2, p["w1t"], p["b1"], p["w2t"], p["b2"]).unsqueeze(0)
|
| 207 |
|
build/torch-neuron/nki_kernels/__init__.py
CHANGED
|
@@ -4,6 +4,7 @@ from .attention_block import (
|
|
| 4 |
attention_block,
|
| 5 |
lane_permutation,
|
| 6 |
pack_out_weights,
|
|
|
|
| 7 |
pack_tables,
|
| 8 |
pack_v_weights,
|
| 9 |
pack_weights,
|
|
@@ -20,5 +21,5 @@ from .swiglu_stream import nki_swiglu_stream_ln as swiglu_stream_ln
|
|
| 20 |
__all__ = [
|
| 21 |
"attention_block", "attention_block_l1", "attention_block_l1_ln", "res_layernorm_fold",
|
| 22 |
"mlp", "mlp_res_ln", "pack_b1", "swiglu_stream", "swiglu_stream_ln",
|
| 23 |
-
"pack_weights", "pack_tables", "pack_v_weights", "pack_out_weights", "lane_permutation",
|
| 24 |
]
|
|
|
|
| 4 |
attention_block,
|
| 5 |
lane_permutation,
|
| 6 |
pack_out_weights,
|
| 7 |
+
pack_qk_bias,
|
| 8 |
pack_tables,
|
| 9 |
pack_v_weights,
|
| 10 |
pack_weights,
|
|
|
|
| 21 |
__all__ = [
|
| 22 |
"attention_block", "attention_block_l1", "attention_block_l1_ln", "res_layernorm_fold",
|
| 23 |
"mlp", "mlp_res_ln", "pack_b1", "swiglu_stream", "swiglu_stream_ln",
|
| 24 |
+
"pack_weights", "pack_qk_bias", "pack_tables", "pack_v_weights", "pack_out_weights", "lane_permutation",
|
| 25 |
]
|
build/torch-neuron/nki_kernels/attention_block.py
CHANGED
|
@@ -170,13 +170,13 @@ def lane_permutation(heads: int, D: int, rope: str = "half") -> torch.Tensor:
|
|
| 170 |
def pack_weights(wq: torch.Tensor, wk: torch.Tensor,
|
| 171 |
bq: torch.Tensor | None, bk: torch.Tensor | None,
|
| 172 |
heads: int, D: int, rope: str = "half"):
|
| 173 |
-
"""``(w_packed [2,H,H] bf16, b_packed [2,
|
| 174 |
|
| 175 |
FOLD: in the affine-free-LN1 freeze this is called on gamma-folded inputs
|
| 176 |
``pack_weights(wq*g1, wk*g1, bq+wq@t1, bk+wk@t1, heads, D)``. ``wq``/``wk`` are torch
|
| 177 |
``Linear.weight`` ``[H_out, H_in]``; the kernel wants ``[H_in, H_out]`` with ``H_out``
|
| 178 |
-
permuted, i.e. ``w.t()[:, perm]``. The bias is
|
| 179 |
-
|
| 180 |
"""
|
| 181 |
H = wq.shape[0]
|
| 182 |
perm = lane_permutation(heads, D, rope).to(wq.device)
|
|
@@ -185,8 +185,23 @@ def pack_weights(wq: torch.Tensor, wk: torch.Tensor,
|
|
| 185 |
zero = torch.zeros(H, dtype=wq.dtype, device=wq.device)
|
| 186 |
b = torch.stack([(zero if bq is None else bq.detach())[perm].contiguous(),
|
| 187 |
(zero if bk is None else bk.detach())[perm].contiguous()], 0)
|
| 188 |
-
|
| 189 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 190 |
|
| 191 |
|
| 192 |
def pack_tables(cos: torch.Tensor, sin: torch.Tensor, heads: int, D: int, rope: str = "half"):
|
|
@@ -296,7 +311,7 @@ def _attention_block_kernel(h_hbm, wqk_hbm, bqk_hbm, tab_hbm, wvt_hbm, bv_hbm,
|
|
| 296 |
# exercises and which produced rel_l2 0.90 (wrong). See the V-JEPA2 kernels' verification notes.
|
| 297 |
|
| 298 |
_kernel_assert(wqk_hbm.shape == (2, H, Fp), f"wqk must be [2,{H},{Fp}]")
|
| 299 |
-
_kernel_assert(bqk_hbm.shape == (2,
|
| 300 |
_kernel_assert(tab_hbm.shape == (4, _P_MAX, N), f"tab must be [4,{_P_MAX},{N}]")
|
| 301 |
_kernel_assert(wvt_hbm.shape == (H, Fp) and bv_hbm.shape == (1, Fp), "bad V weights")
|
| 302 |
_kernel_assert(wot_hbm.shape == (Fp, H) and bo_hbm.shape == (1, H), "bad O weights")
|
|
@@ -323,6 +338,7 @@ def _attention_block_kernel(h_hbm, wqk_hbm, bqk_hbm, tab_hbm, wvt_hbm, bv_hbm,
|
|
| 323 |
head0 = pid * local_heads
|
| 324 |
col0 = pid * (local_heads * D) # first permuted qk weight column
|
| 325 |
vcol0 = pid * (local_heads * D) # first natural v weight column (same offset)
|
|
|
|
| 326 |
|
| 327 |
# ---- RESIDENT MODE (resident=1, N<=2048 only): stage the per-launch-invariant inputs ONCE,
|
| 328 |
# at the BOTTOM of SBUF, before the image loop. Everything per image is then allocated above
|
|
@@ -342,13 +358,13 @@ def _attention_block_kernel(h_hbm, wqk_hbm, bqk_hbm, tab_hbm, wvt_hbm, bv_hbm,
|
|
| 342 |
block_dim=[n_chunks], num_free_tiles=[n_chunks])
|
| 343 |
for c in range(n_chunks):
|
| 344 |
nisa.dma_copy(dst=r_wv_sb[c], src=wvt_hbm[c * _P_MAX:(c + 1) * _P_MAX, vcol0:vcol0 + local_heads * D])
|
| 345 |
-
|
| 346 |
-
|
| 347 |
-
|
|
|
|
| 348 |
for qk in range(2):
|
| 349 |
-
|
| 350 |
-
|
| 351 |
-
nisa.tensor_copy(dst=r_b_sb[qk][g], src=r_b_stage)
|
| 352 |
r_bv_bc = ra.alloc_sbuf_tensor(shape=(_V_TILE, local_heads * D), dtype=nl.float32, align_to=4)
|
| 353 |
nisa.dma_copy(dst=r_bv_bc, src=bv_hbm.ap(pattern=[[0, _V_TILE], [1, local_heads * D]], offset=vcol0))
|
| 354 |
r_tab = ra.alloc_sbuf_tensor(shape=(_P_MAX, _PSUM_FMAX), dtype=nl.float32,
|
|
@@ -467,17 +483,13 @@ def _attention_block_kernel(h_hbm, wqk_hbm, bqk_hbm, tab_hbm, wvt_hbm, bv_hbm,
|
|
| 467 |
vcol0:vcol0 + local_heads * D])
|
| 468 |
|
| 469 |
# biases: qk as (128,1) per group; v broadcast to (128, 8*64) once
|
| 470 |
-
|
| 471 |
-
|
| 472 |
-
|
| 473 |
-
|
| 474 |
for qk in range(2):
|
| 475 |
-
|
| 476 |
-
|
| 477 |
-
src=bqk_hbm[qk, col0 + g * _P_MAX:col0 + (g + 1) * _P_MAX, 0:1])
|
| 478 |
-
# SITE :224 -- bias CAST, left UNPINNED on purpose (0.09% of the window, already at
|
| 479 |
-
# the Vector fixed fee; nothing to move).
|
| 480 |
-
nisa.tensor_copy(dst=b_sb[qk][g], src=b_stage)
|
| 481 |
bv_bc = allocator.alloc_sbuf_tensor(
|
| 482 |
shape=(_V_TILE, local_heads * D), dtype=nl.float32, align_to=4)
|
| 483 |
nisa.dma_copy(dst=bv_bc,
|
|
@@ -592,7 +604,7 @@ def _attention_block_kernel(h_hbm, wqk_hbm, bqk_hbm, tab_hbm, wvt_hbm, bv_hbm,
|
|
| 592 |
# NKI 0.7: a Scalar-engine pin here fails ISA opcode_engine_check on newer compilers
|
| 593 |
# TOT compiler -> unpinned (compiler picks the engine).
|
| 594 |
nisa.tensor_scalar(dst=proj[par][:, 0:tw], data=ps[:, 0:tw],
|
| 595 |
-
op0=nl.add, operand0=b_sb[qk][g])
|
| 596 |
|
| 597 |
# rope: the six ops of nki_rope, same order, tables sliced per tile
|
| 598 |
nisa.tensor_tensor(dst=acc[:, 0:tw], data1=proj[0][:, 0:tw],
|
|
@@ -798,4 +810,6 @@ def attention_block(
|
|
| 798 |
f"the V-JEPA2 kernels' verification notes. Use N<=2048 (up to 16-frame clips) -- the bit-for-bit "
|
| 799 |
f"shipping path."
|
| 800 |
)
|
|
|
|
|
|
|
| 801 |
return _ATTN[grid](h, wqk, bqk, tab, wvt, bv, wot, bo, scale, head_dim, n_valid, hb, grid, nimg, resident)
|
|
|
|
| 170 |
def pack_weights(wq: torch.Tensor, wk: torch.Tensor,
|
| 171 |
bq: torch.Tensor | None, bk: torch.Tensor | None,
|
| 172 |
heads: int, D: int, rope: str = "half"):
|
| 173 |
+
"""``(w_packed [2,H,H] bf16, b_packed [2,128,H/128] fp32)`` -- transposed and column-permuted.
|
| 174 |
|
| 175 |
FOLD: in the affine-free-LN1 freeze this is called on gamma-folded inputs
|
| 176 |
``pack_weights(wq*g1, wk*g1, bq+wq@t1, bk+wk@t1, heads, D)``. ``wq``/``wk`` are torch
|
| 177 |
``Linear.weight`` ``[H_out, H_in]``; the kernel wants ``[H_in, H_out]`` with ``H_out``
|
| 178 |
+
permuted, i.e. ``w.t()[:, perm]``. The bias is packed by :func:`pack_qk_bias` so each program
|
| 179 |
+
loads all of its 128-row groups in one DMA per q/k.
|
| 180 |
"""
|
| 181 |
H = wq.shape[0]
|
| 182 |
perm = lane_permutation(heads, D, rope).to(wq.device)
|
|
|
|
| 185 |
zero = torch.zeros(H, dtype=wq.dtype, device=wq.device)
|
| 186 |
b = torch.stack([(zero if bq is None else bq.detach())[perm].contiguous(),
|
| 187 |
(zero if bk is None else bk.detach())[perm].contiguous()], 0)
|
| 188 |
+
return w.to(torch.bfloat16).contiguous(), pack_qk_bias(b.to(torch.bfloat16).float())
|
| 189 |
+
|
| 190 |
+
|
| 191 |
+
def pack_qk_bias(b: torch.Tensor) -> torch.Tensor:
|
| 192 |
+
"""``[2, Fp]`` (or legacy ``[2, Fp, 1]``) permuted Q/K bias -> ``[2, 128, Fp/128]`` fp32.
|
| 193 |
+
|
| 194 |
+
``out[qk, p, g] == b[qk, g*128 + p]``: column ``g`` is the per-partition bias of 128-row group
|
| 195 |
+
``g``, so a program loads all of its groups with one DMA per q/k. Values are first rounded to
|
| 196 |
+
bf16 (the precision the v1 kernel loaded and cast), which keeps the output bit-identical to v1.
|
| 197 |
+
"""
|
| 198 |
+
if b.dim() == 3 and b.shape[1] == _P_MAX:
|
| 199 |
+
return b.float().contiguous()
|
| 200 |
+
if b.dim() == 3:
|
| 201 |
+
b = b.squeeze(-1)
|
| 202 |
+
Fp = b.shape[1]
|
| 203 |
+
_kernel_assert(Fp % _P_MAX == 0, f"bias length {Fp} must be a multiple of 128")
|
| 204 |
+
return b.to(torch.bfloat16).float().reshape(2, Fp // _P_MAX, _P_MAX).transpose(1, 2).contiguous()
|
| 205 |
|
| 206 |
|
| 207 |
def pack_tables(cos: torch.Tensor, sin: torch.Tensor, heads: int, D: int, rope: str = "half"):
|
|
|
|
| 311 |
# exercises and which produced rel_l2 0.90 (wrong). See the V-JEPA2 kernels' verification notes.
|
| 312 |
|
| 313 |
_kernel_assert(wqk_hbm.shape == (2, H, Fp), f"wqk must be [2,{H},{Fp}]")
|
| 314 |
+
_kernel_assert(bqk_hbm.shape == (2, _P_MAX, Fp // _P_MAX), f"bqk must be pack_qk_bias(...): [2,128,{Fp // _P_MAX}]")
|
| 315 |
_kernel_assert(tab_hbm.shape == (4, _P_MAX, N), f"tab must be [4,{_P_MAX},{N}]")
|
| 316 |
_kernel_assert(wvt_hbm.shape == (H, Fp) and bv_hbm.shape == (1, Fp), "bad V weights")
|
| 317 |
_kernel_assert(wot_hbm.shape == (Fp, H) and bo_hbm.shape == (1, H), "bad O weights")
|
|
|
|
| 338 |
head0 = pid * local_heads
|
| 339 |
col0 = pid * (local_heads * D) # first permuted qk weight column
|
| 340 |
vcol0 = pid * (local_heads * D) # first natural v weight column (same offset)
|
| 341 |
+
gcol0 = pid * ((local_heads * D) // _P_MAX) # first 128-row bias group (multiply only: pid may be symbolic)
|
| 342 |
|
| 343 |
# ---- RESIDENT MODE (resident=1, N<=2048 only): stage the per-launch-invariant inputs ONCE,
|
| 344 |
# at the BOTTOM of SBUF, before the image loop. Everything per image is then allocated above
|
|
|
|
| 358 |
block_dim=[n_chunks], num_free_tiles=[n_chunks])
|
| 359 |
for c in range(n_chunks):
|
| 360 |
nisa.dma_copy(dst=r_wv_sb[c], src=wvt_hbm[c * _P_MAX:(c + 1) * _P_MAX, vcol0:vcol0 + local_heads * D])
|
| 361 |
+
# Q/K biases: host-packed fp32 [2, 128, Fp/128] (pack_qk_bias); this program's groups land in
|
| 362 |
+
# ONE DMA per q/k on the Scalar HW-DGE queue (was one [128,1] DMA + one cast per group).
|
| 363 |
+
r_b_sb = ra.alloc_sbuf_tensor(shape=(_P_MAX, local_quads * 2), dtype=nl.float32,
|
| 364 |
+
block_dim=[2], num_free_tiles=[2], align_to=4)
|
| 365 |
for qk in range(2):
|
| 366 |
+
nisa.dma_copy(dst=r_b_sb[qk], src=bqk_hbm[qk, 0:_P_MAX, gcol0:gcol0 + local_quads * 2],
|
| 367 |
+
dge_mode=nisa.dge_mode.hwdge, engine=nisa.engine.scalar)
|
|
|
|
| 368 |
r_bv_bc = ra.alloc_sbuf_tensor(shape=(_V_TILE, local_heads * D), dtype=nl.float32, align_to=4)
|
| 369 |
nisa.dma_copy(dst=r_bv_bc, src=bv_hbm.ap(pattern=[[0, _V_TILE], [1, local_heads * D]], offset=vcol0))
|
| 370 |
r_tab = ra.alloc_sbuf_tensor(shape=(_P_MAX, _PSUM_FMAX), dtype=nl.float32,
|
|
|
|
| 483 |
vcol0:vcol0 + local_heads * D])
|
| 484 |
|
| 485 |
# biases: qk as (128,1) per group; v broadcast to (128, 8*64) once
|
| 486 |
+
# Q/K biases: host-packed fp32 [2, 128, Fp/128] (pack_qk_bias); this program's groups land in
|
| 487 |
+
# ONE DMA per q/k on the Scalar HW-DGE queue (was one [128,1] DMA + one cast per group).
|
| 488 |
+
b_sb = allocator.alloc_sbuf_tensor(shape=(_P_MAX, local_quads * 2), dtype=nl.float32,
|
| 489 |
+
block_dim=[2], num_free_tiles=[2], align_to=4)
|
| 490 |
for qk in range(2):
|
| 491 |
+
nisa.dma_copy(dst=b_sb[qk], src=bqk_hbm[qk, 0:_P_MAX, gcol0:gcol0 + local_quads * 2],
|
| 492 |
+
dge_mode=nisa.dge_mode.hwdge, engine=nisa.engine.scalar)
|
|
|
|
|
|
|
|
|
|
|
|
|
| 493 |
bv_bc = allocator.alloc_sbuf_tensor(
|
| 494 |
shape=(_V_TILE, local_heads * D), dtype=nl.float32, align_to=4)
|
| 495 |
nisa.dma_copy(dst=bv_bc,
|
|
|
|
| 604 |
# NKI 0.7: a Scalar-engine pin here fails ISA opcode_engine_check on newer compilers
|
| 605 |
# TOT compiler -> unpinned (compiler picks the engine).
|
| 606 |
nisa.tensor_scalar(dst=proj[par][:, 0:tw], data=ps[:, 0:tw],
|
| 607 |
+
op0=nl.add, operand0=b_sb[qk][:, g:g + 1])
|
| 608 |
|
| 609 |
# rope: the six ops of nki_rope, same order, tables sliced per tile
|
| 610 |
nisa.tensor_tensor(dst=acc[:, 0:tw], data1=proj[0][:, 0:tw],
|
|
|
|
| 810 |
f"the V-JEPA2 kernels' verification notes. Use N<=2048 (up to 16-frame clips) -- the bit-for-bit "
|
| 811 |
f"shipping path."
|
| 812 |
)
|
| 813 |
+
if bqk.dim() == 3 and bqk.shape[1] != 128:
|
| 814 |
+
bqk = pack_qk_bias(bqk)
|
| 815 |
return _ATTN[grid](h, wqk, bqk, tab, wvt, bv, wot, bo, scale, head_dim, n_valid, hb, grid, nimg, resident)
|
build/torch-neuron/nki_kernels/attention_block_l1.py
CHANGED
|
@@ -36,6 +36,7 @@ from nki.isa import engine
|
|
| 36 |
from nkilib.core.utils.modular_allocator import ModularAllocator
|
| 37 |
|
| 38 |
from .attention_inner import _attention_cte
|
|
|
|
| 39 |
|
| 40 |
_P_MAX = 128
|
| 41 |
_PSUM_FMAX = 512
|
|
@@ -92,7 +93,7 @@ def _attention_block_l1_kernel(h_hbm, wqk_hbm, bqk_hbm, tab_hbm, wvt_hbm, bv_hbm
|
|
| 92 |
v_per_tt = _PSUM_FMAX // _V_TILE
|
| 93 |
|
| 94 |
_kernel_assert(wqk_hbm.shape == (2, H, Fp), "bad wqk")
|
| 95 |
-
_kernel_assert(bqk_hbm.shape == (2,
|
| 96 |
_kernel_assert(tab_hbm.shape == (4, _P_MAX, N), "bad tab")
|
| 97 |
_kernel_assert(wvt_hbm.shape == (H, Fp) and bv_hbm.shape == (1, Fp), "bad V weights")
|
| 98 |
_kernel_assert(wot_hbm.shape == (Fp, H) and bo_hbm.shape == (1, H), "bad O weights")
|
|
@@ -127,15 +128,13 @@ def _attention_block_l1_kernel(h_hbm, wqk_hbm, bqk_hbm, tab_hbm, wvt_hbm, bv_hbm
|
|
| 127 |
for c in range(n_chunks):
|
| 128 |
nisa.dma_copy(dst=wv_sb_[c],
|
| 129 |
src=wvt_hbm[c * _P_MAX:(c + 1) * _P_MAX, col0:col0 + local_heads * D])
|
| 130 |
-
|
| 131 |
-
|
| 132 |
-
|
| 133 |
-
|
| 134 |
for qk in range(2):
|
| 135 |
-
|
| 136 |
-
|
| 137 |
-
src=bqk_hbm[qk, col0 + gg * _P_MAX:col0 + (gg + 1) * _P_MAX, 0:1])
|
| 138 |
-
nisa.tensor_copy(dst=b_sb_[qk][gg], src=b_stage_)
|
| 139 |
bv_bc_ = ra.alloc_sbuf_tensor(shape=(_V_TILE, local_heads * D), dtype=nl.float32, align_to=4)
|
| 140 |
nisa.dma_copy(dst=bv_bc_,
|
| 141 |
src=bv_hbm.ap(pattern=[[0, _V_TILE], [1, local_heads * D]], offset=col0))
|
|
@@ -233,15 +232,13 @@ def _attention_block_l1_kernel(h_hbm, wqk_hbm, bqk_hbm, tab_hbm, wvt_hbm, bv_hbm
|
|
| 233 |
for c in range(n_chunks):
|
| 234 |
nisa.dma_copy(dst=wv_sb[c],
|
| 235 |
src=wvt_hbm[c * _P_MAX:(c + 1) * _P_MAX, col0:col0 + local_heads * D])
|
| 236 |
-
|
| 237 |
-
|
| 238 |
-
|
| 239 |
-
|
| 240 |
for qk in range(2):
|
| 241 |
-
|
| 242 |
-
|
| 243 |
-
src=bqk_hbm[qk, col0 + gg * _P_MAX:col0 + (gg + 1) * _P_MAX, 0:1])
|
| 244 |
-
nisa.tensor_copy(dst=b_sb[qk][gg], src=b_stage)
|
| 245 |
bv_bc = allocator.alloc_sbuf_tensor(shape=(_V_TILE, local_heads * D), dtype=nl.float32,
|
| 246 |
align_to=4)
|
| 247 |
nisa.dma_copy(dst=bv_bc,
|
|
@@ -286,7 +283,7 @@ def _attention_block_l1_kernel(h_hbm, wqk_hbm, bqk_hbm, tab_hbm, wvt_hbm, bv_hbm
|
|
| 286 |
stationary=w_sb[qk][c][:, cols:cols + _P_MAX],
|
| 287 |
moving=x_t[c][tt][:, 0:tw], accumulate=(c > 0))
|
| 288 |
nisa.tensor_scalar(dst=proj[par][:, 0:tw], data=ps[:, 0:tw],
|
| 289 |
-
op0=nl.add, operand0=b_sb[qk][gq])
|
| 290 |
nisa.tensor_tensor(dst=acc[:, 0:tw], data1=proj[0][:, 0:tw],
|
| 291 |
data2=tab[0][tt][:, 0:tw], op=nl.multiply)
|
| 292 |
nisa.tensor_tensor(dst=tmp[:, 0:tw], data1=proj[1][:, 0:tw],
|
|
@@ -399,4 +396,6 @@ def attention_block_l1(
|
|
| 399 |
raise RuntimeError("attention_block_l1 is inference-only; use torch.no_grad().")
|
| 400 |
if h.shape[0] // nimg > 2048:
|
| 401 |
raise NotImplementedError("attention_block_l1 supports N<=2048 only")
|
|
|
|
|
|
|
| 402 |
return _ATTN_L1(h, wqk, bqk, tab, wvt, bv, wot, bo, scale, head_dim, n_valid, nimg, resident)
|
|
|
|
| 36 |
from nkilib.core.utils.modular_allocator import ModularAllocator
|
| 37 |
|
| 38 |
from .attention_inner import _attention_cte
|
| 39 |
+
from .attention_block import pack_qk_bias as _pack_qk_bias
|
| 40 |
|
| 41 |
_P_MAX = 128
|
| 42 |
_PSUM_FMAX = 512
|
|
|
|
| 93 |
v_per_tt = _PSUM_FMAX // _V_TILE
|
| 94 |
|
| 95 |
_kernel_assert(wqk_hbm.shape == (2, H, Fp), "bad wqk")
|
| 96 |
+
_kernel_assert(bqk_hbm.shape == (2, _P_MAX, Fp // _P_MAX), "bqk must be pack_qk_bias(...)")
|
| 97 |
_kernel_assert(tab_hbm.shape == (4, _P_MAX, N), "bad tab")
|
| 98 |
_kernel_assert(wvt_hbm.shape == (H, Fp) and bv_hbm.shape == (1, Fp), "bad V weights")
|
| 99 |
_kernel_assert(wot_hbm.shape == (Fp, H) and bo_hbm.shape == (1, H), "bad O weights")
|
|
|
|
| 128 |
for c in range(n_chunks):
|
| 129 |
nisa.dma_copy(dst=wv_sb_[c],
|
| 130 |
src=wvt_hbm[c * _P_MAX:(c + 1) * _P_MAX, col0:col0 + local_heads * D])
|
| 131 |
+
# Q/K biases: host-packed fp32 [2, 128, Fp/128] (pack_qk_bias); this program's groups land in
|
| 132 |
+
# ONE DMA per q/k on the Scalar HW-DGE queue (was one [128,1] DMA + one cast per group).
|
| 133 |
+
b_sb_ = ra.alloc_sbuf_tensor(shape=(_P_MAX, local_quads * 2), dtype=nl.float32,
|
| 134 |
+
block_dim=[2], num_free_tiles=[2], align_to=4)
|
| 135 |
for qk in range(2):
|
| 136 |
+
nisa.dma_copy(dst=b_sb_[qk], src=bqk_hbm[qk, 0:_P_MAX, g * (_PSUM_FMAX // _P_MAX):g * (_PSUM_FMAX // _P_MAX) + local_quads * 2],
|
| 137 |
+
dge_mode=nisa.dge_mode.hwdge, engine=nisa.engine.scalar)
|
|
|
|
|
|
|
| 138 |
bv_bc_ = ra.alloc_sbuf_tensor(shape=(_V_TILE, local_heads * D), dtype=nl.float32, align_to=4)
|
| 139 |
nisa.dma_copy(dst=bv_bc_,
|
| 140 |
src=bv_hbm.ap(pattern=[[0, _V_TILE], [1, local_heads * D]], offset=col0))
|
|
|
|
| 232 |
for c in range(n_chunks):
|
| 233 |
nisa.dma_copy(dst=wv_sb[c],
|
| 234 |
src=wvt_hbm[c * _P_MAX:(c + 1) * _P_MAX, col0:col0 + local_heads * D])
|
| 235 |
+
# Q/K biases: host-packed fp32 [2, 128, Fp/128] (pack_qk_bias); this program's groups land in
|
| 236 |
+
# ONE DMA per q/k on the Scalar HW-DGE queue (was one [128,1] DMA + one cast per group).
|
| 237 |
+
b_sb = allocator.alloc_sbuf_tensor(shape=(_P_MAX, local_quads * 2), dtype=nl.float32,
|
| 238 |
+
block_dim=[2], num_free_tiles=[2], align_to=4)
|
| 239 |
for qk in range(2):
|
| 240 |
+
nisa.dma_copy(dst=b_sb[qk], src=bqk_hbm[qk, 0:_P_MAX, g * (_PSUM_FMAX // _P_MAX):g * (_PSUM_FMAX // _P_MAX) + local_quads * 2],
|
| 241 |
+
dge_mode=nisa.dge_mode.hwdge, engine=nisa.engine.scalar)
|
|
|
|
|
|
|
| 242 |
bv_bc = allocator.alloc_sbuf_tensor(shape=(_V_TILE, local_heads * D), dtype=nl.float32,
|
| 243 |
align_to=4)
|
| 244 |
nisa.dma_copy(dst=bv_bc,
|
|
|
|
| 283 |
stationary=w_sb[qk][c][:, cols:cols + _P_MAX],
|
| 284 |
moving=x_t[c][tt][:, 0:tw], accumulate=(c > 0))
|
| 285 |
nisa.tensor_scalar(dst=proj[par][:, 0:tw], data=ps[:, 0:tw],
|
| 286 |
+
op0=nl.add, operand0=b_sb[qk][:, gq:gq + 1])
|
| 287 |
nisa.tensor_tensor(dst=acc[:, 0:tw], data1=proj[0][:, 0:tw],
|
| 288 |
data2=tab[0][tt][:, 0:tw], op=nl.multiply)
|
| 289 |
nisa.tensor_tensor(dst=tmp[:, 0:tw], data1=proj[1][:, 0:tw],
|
|
|
|
| 396 |
raise RuntimeError("attention_block_l1 is inference-only; use torch.no_grad().")
|
| 397 |
if h.shape[0] // nimg > 2048:
|
| 398 |
raise NotImplementedError("attention_block_l1 supports N<=2048 only")
|
| 399 |
+
if bqk.dim() == 3 and bqk.shape[1] != 128:
|
| 400 |
+
bqk = _pack_qk_bias(bqk)
|
| 401 |
return _ATTN_L1(h, wqk, bqk, tab, wvt, bv, wot, bo, scale, head_dim, n_valid, nimg, resident)
|
build/torch-neuron/nki_kernels/attention_block_l1_ln.py
CHANGED
|
@@ -45,6 +45,7 @@ from nki.isa import engine
|
|
| 45 |
from nkilib.core.utils.modular_allocator import ModularAllocator
|
| 46 |
|
| 47 |
from .attention_inner import _attention_cte
|
|
|
|
| 48 |
|
| 49 |
_P_MAX = 128
|
| 50 |
_PSUM_FMAX = 512
|
|
@@ -101,7 +102,7 @@ def _attention_block_l1_ln_kernel(h_hbm, wqk_hbm, bqk_hbm, tab_hbm, wvt_hbm, bv_
|
|
| 101 |
v_per_tt = _PSUM_FMAX // _V_TILE
|
| 102 |
|
| 103 |
_kernel_assert(wqk_hbm.shape == (2, H, Fp), "bad wqk")
|
| 104 |
-
_kernel_assert(bqk_hbm.shape == (2,
|
| 105 |
_kernel_assert(tab_hbm.shape == (4, _P_MAX, N), "bad tab")
|
| 106 |
_kernel_assert(wvt_hbm.shape == (H, Fp) and bv_hbm.shape == (1, Fp), "bad V weights")
|
| 107 |
_kernel_assert(wot_hbm.shape == (Fp, H) and bo_hbm.shape == (1, H), "bad O weights")
|
|
@@ -140,15 +141,13 @@ def _attention_block_l1_ln_kernel(h_hbm, wqk_hbm, bqk_hbm, tab_hbm, wvt_hbm, bv_
|
|
| 140 |
for c in range(n_chunks):
|
| 141 |
nisa.dma_copy(dst=wv_sb_[c],
|
| 142 |
src=wvt_hbm[c * _P_MAX:(c + 1) * _P_MAX, col0:col0 + local_heads * D])
|
| 143 |
-
|
| 144 |
-
|
| 145 |
-
|
| 146 |
-
|
| 147 |
for qk in range(2):
|
| 148 |
-
|
| 149 |
-
|
| 150 |
-
src=bqk_hbm[qk, col0 + gg * _P_MAX:col0 + (gg + 1) * _P_MAX, 0:1])
|
| 151 |
-
nisa.tensor_copy(dst=b_sb_[qk][gg], src=b_stage_)
|
| 152 |
bv_bc_ = ra.alloc_sbuf_tensor(shape=(_V_TILE, local_heads * D), dtype=nl.float32, align_to=4)
|
| 153 |
nisa.dma_copy(dst=bv_bc_,
|
| 154 |
src=bv_hbm.ap(pattern=[[0, _V_TILE], [1, local_heads * D]], offset=col0))
|
|
@@ -246,15 +245,13 @@ def _attention_block_l1_ln_kernel(h_hbm, wqk_hbm, bqk_hbm, tab_hbm, wvt_hbm, bv_
|
|
| 246 |
for c in range(n_chunks):
|
| 247 |
nisa.dma_copy(dst=wv_sb[c],
|
| 248 |
src=wvt_hbm[c * _P_MAX:(c + 1) * _P_MAX, col0:col0 + local_heads * D])
|
| 249 |
-
|
| 250 |
-
|
| 251 |
-
|
| 252 |
-
|
| 253 |
for qk in range(2):
|
| 254 |
-
|
| 255 |
-
|
| 256 |
-
src=bqk_hbm[qk, col0 + gg * _P_MAX:col0 + (gg + 1) * _P_MAX, 0:1])
|
| 257 |
-
nisa.tensor_copy(dst=b_sb[qk][gg], src=b_stage)
|
| 258 |
bv_bc = allocator.alloc_sbuf_tensor(shape=(_V_TILE, local_heads * D), dtype=nl.float32,
|
| 259 |
align_to=4)
|
| 260 |
nisa.dma_copy(dst=bv_bc,
|
|
@@ -299,7 +296,7 @@ def _attention_block_l1_ln_kernel(h_hbm, wqk_hbm, bqk_hbm, tab_hbm, wvt_hbm, bv_
|
|
| 299 |
stationary=w_sb[qk][c][:, cols:cols + _P_MAX],
|
| 300 |
moving=x_t[c][tt][:, 0:tw], accumulate=(c > 0))
|
| 301 |
nisa.tensor_scalar(dst=proj[par][:, 0:tw], data=ps[:, 0:tw],
|
| 302 |
-
op0=nl.add, operand0=b_sb[qk][gq])
|
| 303 |
nisa.tensor_tensor(dst=acc[:, 0:tw], data1=proj[0][:, 0:tw],
|
| 304 |
data2=tab[0][tt][:, 0:tw], op=nl.multiply)
|
| 305 |
nisa.tensor_tensor(dst=tmp[:, 0:tw], data1=proj[1][:, 0:tw],
|
|
@@ -448,4 +445,6 @@ def attention_block_l1_ln(
|
|
| 448 |
raise RuntimeError("attention_block_l1_ln is inference-only; use torch.no_grad().")
|
| 449 |
if h.shape[0] // nimg > 2048:
|
| 450 |
raise NotImplementedError("attention_block_l1_ln supports N<=2048 only")
|
|
|
|
|
|
|
| 451 |
return _ATTN_L1_LN(h, wqk, bqk, tab, wvt, bv, wot, bo, x, scale, head_dim, n_valid, eps, nimg, resident)
|
|
|
|
| 45 |
from nkilib.core.utils.modular_allocator import ModularAllocator
|
| 46 |
|
| 47 |
from .attention_inner import _attention_cte
|
| 48 |
+
from .attention_block import pack_qk_bias as _pack_qk_bias
|
| 49 |
|
| 50 |
_P_MAX = 128
|
| 51 |
_PSUM_FMAX = 512
|
|
|
|
| 102 |
v_per_tt = _PSUM_FMAX // _V_TILE
|
| 103 |
|
| 104 |
_kernel_assert(wqk_hbm.shape == (2, H, Fp), "bad wqk")
|
| 105 |
+
_kernel_assert(bqk_hbm.shape == (2, _P_MAX, Fp // _P_MAX), "bqk must be pack_qk_bias(...)")
|
| 106 |
_kernel_assert(tab_hbm.shape == (4, _P_MAX, N), "bad tab")
|
| 107 |
_kernel_assert(wvt_hbm.shape == (H, Fp) and bv_hbm.shape == (1, Fp), "bad V weights")
|
| 108 |
_kernel_assert(wot_hbm.shape == (Fp, H) and bo_hbm.shape == (1, H), "bad O weights")
|
|
|
|
| 141 |
for c in range(n_chunks):
|
| 142 |
nisa.dma_copy(dst=wv_sb_[c],
|
| 143 |
src=wvt_hbm[c * _P_MAX:(c + 1) * _P_MAX, col0:col0 + local_heads * D])
|
| 144 |
+
# Q/K biases: host-packed fp32 [2, 128, Fp/128] (pack_qk_bias); this program's groups land in
|
| 145 |
+
# ONE DMA per q/k on the Scalar HW-DGE queue (was one [128,1] DMA + one cast per group).
|
| 146 |
+
b_sb_ = ra.alloc_sbuf_tensor(shape=(_P_MAX, local_quads * 2), dtype=nl.float32,
|
| 147 |
+
block_dim=[2], num_free_tiles=[2], align_to=4)
|
| 148 |
for qk in range(2):
|
| 149 |
+
nisa.dma_copy(dst=b_sb_[qk], src=bqk_hbm[qk, 0:_P_MAX, g * (_PSUM_FMAX // _P_MAX):g * (_PSUM_FMAX // _P_MAX) + local_quads * 2],
|
| 150 |
+
dge_mode=nisa.dge_mode.hwdge, engine=nisa.engine.scalar)
|
|
|
|
|
|
|
| 151 |
bv_bc_ = ra.alloc_sbuf_tensor(shape=(_V_TILE, local_heads * D), dtype=nl.float32, align_to=4)
|
| 152 |
nisa.dma_copy(dst=bv_bc_,
|
| 153 |
src=bv_hbm.ap(pattern=[[0, _V_TILE], [1, local_heads * D]], offset=col0))
|
|
|
|
| 245 |
for c in range(n_chunks):
|
| 246 |
nisa.dma_copy(dst=wv_sb[c],
|
| 247 |
src=wvt_hbm[c * _P_MAX:(c + 1) * _P_MAX, col0:col0 + local_heads * D])
|
| 248 |
+
# Q/K biases: host-packed fp32 [2, 128, Fp/128] (pack_qk_bias); this program's groups land in
|
| 249 |
+
# ONE DMA per q/k on the Scalar HW-DGE queue (was one [128,1] DMA + one cast per group).
|
| 250 |
+
b_sb = allocator.alloc_sbuf_tensor(shape=(_P_MAX, local_quads * 2), dtype=nl.float32,
|
| 251 |
+
block_dim=[2], num_free_tiles=[2], align_to=4)
|
| 252 |
for qk in range(2):
|
| 253 |
+
nisa.dma_copy(dst=b_sb[qk], src=bqk_hbm[qk, 0:_P_MAX, g * (_PSUM_FMAX // _P_MAX):g * (_PSUM_FMAX // _P_MAX) + local_quads * 2],
|
| 254 |
+
dge_mode=nisa.dge_mode.hwdge, engine=nisa.engine.scalar)
|
|
|
|
|
|
|
| 255 |
bv_bc = allocator.alloc_sbuf_tensor(shape=(_V_TILE, local_heads * D), dtype=nl.float32,
|
| 256 |
align_to=4)
|
| 257 |
nisa.dma_copy(dst=bv_bc,
|
|
|
|
| 296 |
stationary=w_sb[qk][c][:, cols:cols + _P_MAX],
|
| 297 |
moving=x_t[c][tt][:, 0:tw], accumulate=(c > 0))
|
| 298 |
nisa.tensor_scalar(dst=proj[par][:, 0:tw], data=ps[:, 0:tw],
|
| 299 |
+
op0=nl.add, operand0=b_sb[qk][:, gq:gq + 1])
|
| 300 |
nisa.tensor_tensor(dst=acc[:, 0:tw], data1=proj[0][:, 0:tw],
|
| 301 |
data2=tab[0][tt][:, 0:tw], op=nl.multiply)
|
| 302 |
nisa.tensor_tensor(dst=tmp[:, 0:tw], data1=proj[1][:, 0:tw],
|
|
|
|
| 445 |
raise RuntimeError("attention_block_l1_ln is inference-only; use torch.no_grad().")
|
| 446 |
if h.shape[0] // nimg > 2048:
|
| 447 |
raise NotImplementedError("attention_block_l1_ln supports N<=2048 only")
|
| 448 |
+
if bqk.dim() == 3 and bqk.shape[1] != 128:
|
| 449 |
+
bqk = _pack_qk_bias(bqk)
|
| 450 |
return _ATTN_L1_LN(h, wqk, bqk, tab, wvt, bv, wot, bo, x, scale, head_dim, n_valid, eps, nimg, resident)
|
build/torch-neuron/nki_kernels/attention_inner.py
CHANGED
|
@@ -857,7 +857,14 @@ def _update_max_impl(grp_i, ac: AttnConfig, atp: AttnTileParams, sp: SectionPara
|
|
| 857 |
# num_sections==1: the reduce writes straight into mm1_running_max (max-direct); the [128,1]
|
| 858 |
# copy the vendor once had here is gone (3,072 Vector instructions removed upstream).
|
| 859 |
_max_dst = bufs.mm1_running_max[:, grp_i]
|
| 860 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 861 |
# num_sections==1: no running-max copy / flash correction factor. sink block DEAD (sink=None).
|
| 862 |
|
| 863 |
|
|
@@ -873,7 +880,10 @@ def _exp_impl(grp_i, ac: AttnConfig, atp: AttnTileParams, sp: SectionParams,
|
|
| 873 |
the reduce_cmd discipline are preserved verbatim."""
|
| 874 |
# is_causal=False -> has_any_compute_pred is True.
|
| 875 |
q_seqlen_offset = grp_i * atp.sb_p
|
| 876 |
-
|
|
|
|
|
|
|
|
|
|
| 877 |
|
| 878 |
for large_tile_idx in range(atp.num_large_tiles_per_section):
|
| 879 |
kernel_assert(atp.exp_inst_elems == 512, "Internal validation failed.")
|
|
@@ -993,7 +1003,11 @@ def _exp_impl(grp_i, ac: AttnConfig, atp: AttnTileParams, sp: SectionParams,
|
|
| 993 |
def _pv_impl(grp_i, ac: AttnConfig, atp: AttnTileParams, sp: SectionParams, bufs: AttnInternalBuffers):
|
| 994 |
"""MM2 (P@V) for a Q group. LIVE: is_causal=False -> always compute."""
|
| 995 |
# is_causal=False -> has_any_compute_pred is True.
|
| 996 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
| 997 |
|
| 998 |
for large_tile_idx in range(atp.num_large_tiles_per_section):
|
| 999 |
_pv_large_tile_impl(grp_i, large_tile_idx, ac, atp, sp, bufs)
|
|
@@ -1022,9 +1036,18 @@ def _write_back_impl(grp_i, ac: AttnConfig, atp: AttnTileParams, sp: SectionPara
|
|
| 1022 |
The flash multi-section accumulation and the skip-normalization branches are DEAD."""
|
| 1023 |
# is_causal=False -> has_any_compute_pred is True.
|
| 1024 |
q_seqlen_offset = grp_i * atp.sb_p
|
| 1025 |
-
nisa.tensor_reduce(bufs.exp_section_sum[grp_i][...], nl.add, bufs.exp_partial_sum[grp_i], axis=1)
|
| 1026 |
# num_sections==1, not skip_output_normalization.
|
| 1027 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1028 |
|
| 1029 |
num_p = min(ac.seqlen_q - q_seqlen_offset, _Q_GRP_SZ)
|
| 1030 |
num_f = min(ac.seqlen_q - q_seqlen_offset, _Q_GRP_SZ)
|
|
@@ -1084,6 +1107,16 @@ def _qk_and_max_large_tile_impl(qkmax_grp, large_tile_idx, ac: AttnConfig, atp:
|
|
| 1084 |
q_seqlen_offset = qkmax_grp * atp.sb_p
|
| 1085 |
|
| 1086 |
num_k_tiles_in_large_tile = _LARGE_TILE_SZ // _K_TILE_SZ
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1087 |
for k_tile_idx in range(num_k_tiles_in_large_tile):
|
| 1088 |
mm1_psum_tile = bufs.mm1_psum[qkmax_grp][large_tile_idx][k_tile_idx]
|
| 1089 |
# dynamic_sel_mask=False, but the no-mask branch below reads neither mm1_copy_sb nor
|
|
@@ -1129,14 +1162,25 @@ def _qk_and_max_large_tile_impl(qkmax_grp, large_tile_idx, ac: AttnConfig, atp:
|
|
| 1129 |
num_p = min(ac.seqlen_q - q_seqlen_offset, _Q_GRP_SZ)
|
| 1130 |
num_f = min(seqlen_k - k_start_pos, _K_TILE_SZ)
|
| 1131 |
|
| 1132 |
-
|
| 1133 |
-
|
| 1134 |
-
|
| 1135 |
-
|
| 1136 |
-
|
| 1137 |
-
|
| 1138 |
-
|
| 1139 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1140 |
|
| 1141 |
|
| 1142 |
def _pv_large_tile_impl(pv_grp, large_tile_idx, ac: AttnConfig, atp: AttnTileParams,
|
|
|
|
| 857 |
# num_sections==1: the reduce writes straight into mm1_running_max (max-direct); the [128,1]
|
| 858 |
# copy the vendor once had here is gone (3,072 Vector instructions removed upstream).
|
| 859 |
_max_dst = bufs.mm1_running_max[:, grp_i]
|
| 860 |
+
_n_exec = min(div_ceil(atp.seqlen_k_active_updated, _K_TILE_SZ), _LARGE_TILE_SZ // _K_TILE_SZ,
|
| 861 |
+
atp.num_k_tiles_per_section)
|
| 862 |
+
if atp.num_large_tiles_per_section == 1 and _n_exec >= 2:
|
| 863 |
+
# v1.3: the chained row-max left the whole row max in the last executed k-tile's column.
|
| 864 |
+
nisa.tensor_scalar(dst=_max_dst, data=bufs.mm1_partial_max[grp_i][:, _n_exec - 1:_n_exec],
|
| 865 |
+
op0=nl.multiply, operand0=-1.0)
|
| 866 |
+
else:
|
| 867 |
+
nisa.tensor_reduce(_max_dst, nl.maximum, bufs.mm1_partial_max[grp_i], 1, negate=True)
|
| 868 |
# num_sections==1: no running-max copy / flash correction factor. sink block DEAD (sink=None).
|
| 869 |
|
| 870 |
|
|
|
|
| 880 |
the reduce_cmd discipline are preserved verbatim."""
|
| 881 |
# is_causal=False -> has_any_compute_pred is True.
|
| 882 |
q_seqlen_offset = grp_i * atp.sb_p
|
| 883 |
+
# v1.2: with one large tile only the last emitted column is ever read (see _write_back_impl),
|
| 884 |
+
# and it is written by the chained exp-sum, so the zero-fill is skipped.
|
| 885 |
+
if atp.num_large_tiles_per_section > 1:
|
| 886 |
+
nisa.memset(bufs.exp_partial_sum[grp_i][...], value=0.0) # UNPINNED on purpose (:2993).
|
| 887 |
|
| 888 |
for large_tile_idx in range(atp.num_large_tiles_per_section):
|
| 889 |
kernel_assert(atp.exp_inst_elems == 512, "Internal validation failed.")
|
|
|
|
| 1003 |
def _pv_impl(grp_i, ac: AttnConfig, atp: AttnTileParams, sp: SectionParams, bufs: AttnInternalBuffers):
|
| 1004 |
"""MM2 (P@V) for a Q group. LIVE: is_causal=False -> always compute."""
|
| 1005 |
# is_causal=False -> has_any_compute_pred is True.
|
| 1006 |
+
# v1.2: the mm2_sb zero-fill is skipped when there is a single large tile (N <= 2048, the only
|
| 1007 |
+
# configuration this fork serves): its first write is the PSUM->SBUF tensor_copy over exactly the
|
| 1008 |
+
# [:, :num_f] region that write-back reads, so the zeros were never observed.
|
| 1009 |
+
if atp.num_large_tiles_per_section > 1:
|
| 1010 |
+
nisa.memset(bufs.mm2_sb[grp_i][...], value=0.0)
|
| 1011 |
|
| 1012 |
for large_tile_idx in range(atp.num_large_tiles_per_section):
|
| 1013 |
_pv_large_tile_impl(grp_i, large_tile_idx, ac, atp, sp, bufs)
|
|
|
|
| 1036 |
The flash multi-section accumulation and the skip-normalization branches are DEAD."""
|
| 1037 |
# is_causal=False -> has_any_compute_pred is True.
|
| 1038 |
q_seqlen_offset = grp_i * atp.sb_p
|
|
|
|
| 1039 |
# num_sections==1, not skip_output_normalization.
|
| 1040 |
+
if atp.num_large_tiles_per_section == 1:
|
| 1041 |
+
# v1.2: with one large tile the chained exp-sum leaves the WHOLE row sum in one column
|
| 1042 |
+
# (the last emitted exp tile); every other column is the memset 0.0. Read that column
|
| 1043 |
+
# directly instead of reducing the zero-padded row (x + 0.0 is exact, so bit-identical).
|
| 1044 |
+
# last emitted exp tile = ceil(seqlen_k / exp_inst_elems) - 1 (clamped to the large tile)
|
| 1045 |
+
_last_col = min(div_ceil(atp.seqlen_k_active_updated, atp.exp_inst_elems),
|
| 1046 |
+
_LARGE_TILE_SZ // atp.exp_inst_elems) - 1
|
| 1047 |
+
nisa.reciprocal(bufs.exp_sum_reciprocal[:, grp_i], bufs.exp_partial_sum[grp_i][:, _last_col:_last_col + 1])
|
| 1048 |
+
else:
|
| 1049 |
+
nisa.tensor_reduce(bufs.exp_section_sum[grp_i][...], nl.add, bufs.exp_partial_sum[grp_i], axis=1)
|
| 1050 |
+
nisa.reciprocal(bufs.exp_sum_reciprocal[:, grp_i], bufs.exp_section_sum[grp_i])
|
| 1051 |
|
| 1052 |
num_p = min(ac.seqlen_q - q_seqlen_offset, _Q_GRP_SZ)
|
| 1053 |
num_f = min(ac.seqlen_q - q_seqlen_offset, _Q_GRP_SZ)
|
|
|
|
| 1107 |
q_seqlen_offset = qkmax_grp * atp.sb_p
|
| 1108 |
|
| 1109 |
num_k_tiles_in_large_tile = _LARGE_TILE_SZ // _K_TILE_SZ
|
| 1110 |
+
# v1.3: chain the row-max across the executed k-tiles of this large tile (reset on the first,
|
| 1111 |
+
# accumulate on the rest, write reduce_res only on the last) instead of one read-out per k-tile
|
| 1112 |
+
# followed by a tensor_reduce. Executed tiles are those with k_start < seqlen_k (non-causal:
|
| 1113 |
+
# every tile up to the active key length). max is exact under any grouping -> bit-identical.
|
| 1114 |
+
# single large tile (the only case chained): executed k-tiles are 0 .. ceil(seqlen_k/512)-1
|
| 1115 |
+
_n_exec = min(div_ceil(atp.seqlen_k_active_updated, _K_TILE_SZ), num_k_tiles_in_large_tile,
|
| 1116 |
+
atp.num_k_tiles_per_section)
|
| 1117 |
+
_chain = (atp.num_large_tiles_per_section == 1 and _n_exec >= 2)
|
| 1118 |
+
_first_kt = 0
|
| 1119 |
+
_last_kt = _n_exec - 1
|
| 1120 |
for k_tile_idx in range(num_k_tiles_in_large_tile):
|
| 1121 |
mm1_psum_tile = bufs.mm1_psum[qkmax_grp][large_tile_idx][k_tile_idx]
|
| 1122 |
# dynamic_sel_mask=False, but the no-mask branch below reads neither mm1_copy_sb nor
|
|
|
|
| 1162 |
num_p = min(ac.seqlen_q - q_seqlen_offset, _Q_GRP_SZ)
|
| 1163 |
num_f = min(seqlen_k - k_start_pos, _K_TILE_SZ)
|
| 1164 |
|
| 1165 |
+
if _chain:
|
| 1166 |
+
nisa.tensor_scalar_reduce(
|
| 1167 |
+
mm1_masked_tile[:num_p, nl.ds(k_tile_idx * _K_TILE_SZ, num_f)],
|
| 1168 |
+
data=mm1_psum_tile[:num_p, :num_f],
|
| 1169 |
+
op0=nl.multiply,
|
| 1170 |
+
operand0=ac.scale,
|
| 1171 |
+
reduce_op=nl.maximum,
|
| 1172 |
+
reduce_cmd=(nisa.reduce_cmd.reset_reduce if k_tile_idx == _first_kt else nisa.reduce_cmd.reduce),
|
| 1173 |
+
reduce_res=(mm1_partial_max_tile[:num_p, k_tile_idx_in_section] if k_tile_idx == _last_kt else None),
|
| 1174 |
+
)
|
| 1175 |
+
else:
|
| 1176 |
+
nisa.tensor_scalar_reduce(
|
| 1177 |
+
mm1_masked_tile[:num_p, nl.ds(k_tile_idx * _K_TILE_SZ, num_f)],
|
| 1178 |
+
data=mm1_psum_tile[:num_p, :num_f],
|
| 1179 |
+
op0=nl.multiply,
|
| 1180 |
+
operand0=ac.scale,
|
| 1181 |
+
reduce_op=nl.maximum,
|
| 1182 |
+
reduce_res=mm1_partial_max_tile[:num_p, k_tile_idx_in_section],
|
| 1183 |
+
)
|
| 1184 |
|
| 1185 |
|
| 1186 |
def _pv_large_tile_impl(pv_grp, large_tile_idx, ac: AttnConfig, atp: AttnTileParams,
|