v1.6: two images per producer pass in attention_block (LNC=2 path; LNC=1 single-group models)
Browse filesattention_block gains the pack2 producer packing v1.5 added to attention_block_l1_ln: for even batches up to
256 tokens per image the Q/K/V projections and RoPE run once over 2N columns, attention stays per image.
Output bit-identical to v1.5. Per instance at 224 px, batch 8: ViT-S/16 LNC=1 +9.0%; LNC=2 ViT-S +8.4%,
ViT-B +7.8%, ViT-L +4.9%, ViT-H+ +9.2%.
- CHANGELOG.md +19 -0
- build/torch-neuron/encoder.py +2 -1
- build/torch-neuron/nki_kernels/attention_block.py +76 -33
CHANGELOG.md
CHANGED
|
@@ -1,5 +1,24 @@
|
|
| 1 |
# Changelog
|
| 2 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 3 |
## v1.5 (2026-10)
|
| 4 |
|
| 5 |
* **Two images per producer pass** (LNC=1 all-group attention, `attention_block_l1_ln`, new `pack2` argument, on by
|
|
|
|
| 1 |
# Changelog
|
| 2 |
|
| 3 |
+
## v1.6 (2026-10)
|
| 4 |
+
|
| 5 |
+
* **Two images per producer pass in `attention_block` too** (new `pack2` argument; the encoder enables it for even
|
| 6 |
+
batches up to 256 tokens per image, e.g. 224 px). This kernel serves every model at LNC=2 and single-head-group models
|
| 7 |
+
such as ViT-S/16 at LNC=1. v1.5 had packing only in the LNC=1 all-group kernel. Q/K/V projections and RoPE run once
|
| 8 |
+
over 2N columns; attention stays per image. Output bit-identical to v1.5 (same rel_l2 vs fp32 on all 20 cases).
|
| 9 |
+
Kernel device cycles, ViT-L 224 px LNC=2: batch 4 -8.4% (resident weights) / -12.0% (restaged), batch 2 -8.4%.
|
| 10 |
+
Per instance at 224 px, batch 8:
|
| 11 |
+
|
| 12 |
+
| model | config | v1.5 img/s | v1.6 img/s | change |
|
| 13 |
+
|---|---|---|---|---|
|
| 14 |
+
| ViT-S/16 | LNC=1 DP=8 | 8,592.8 | **9,365.6** | +9.0% |
|
| 15 |
+
| ViT-S/16 | LNC=2 DP=4 | 4,667.6 | 5,061.6 | +8.4% |
|
| 16 |
+
| ViT-B/16 | LNC=2 DP=4 | 3,578.3 | 3,856.8 | +7.8% |
|
| 17 |
+
| ViT-L/16 | LNC=2 DP=4 | 1,514.8 | 1,589.7 | +4.9% |
|
| 18 |
+
| ViT-H+/16 | LNC=2 DP=4 | 507.3 | 554.0 | +9.2% |
|
| 19 |
+
|
| 20 |
+
Single core, batch 4: +5.9% to +10.8%. Batch 1 and 448 px unchanged (packing off). `VIT_ENC_PACK2=0` disables it.
|
| 21 |
+
|
| 22 |
## v1.5 (2026-10)
|
| 23 |
|
| 24 |
* **Two images per producer pass** (LNC=1 all-group attention, `attention_block_l1_ln`, new `pack2` argument, on by
|
build/torch-neuron/encoder.py
CHANGED
|
@@ -209,7 +209,8 @@ class EncoderRunner:
|
|
| 209 |
p["bo_full"], p["scale"], self.D, n_valid, B, self._res(_ab1, N))
|
| 210 |
outs = [_ab.attention_block(z1, p["wqk"], p["bqk"], tab, p["wvt"], p["bv"], p["wot"], p["bo"],
|
| 211 |
p["scale"], self.D, n_valid, hb, gr, B,
|
| 212 |
-
int(z1.shape[0] // B <= _RESIDENT_MAX_N)
|
|
|
|
| 213 |
for hb, gr in self.plan]
|
| 214 |
P = outs[0] if len(outs) == 1 else torch.cat(outs, 0)
|
| 215 |
if P.shape[0] > 4: # residual/LN kernel takes <= 4 branches
|
|
|
|
| 209 |
p["bo_full"], p["scale"], self.D, n_valid, B, self._res(_ab1, N))
|
| 210 |
outs = [_ab.attention_block(z1, p["wqk"], p["bqk"], tab, p["wvt"], p["bv"], p["wot"], p["bo"],
|
| 211 |
p["scale"], self.D, n_valid, hb, gr, B,
|
| 212 |
+
int(z1.shape[0] // B <= _RESIDENT_MAX_N),
|
| 213 |
+
int(_PACK2 and B % 2 == 0 and 2 * (z1.shape[0] // B) <= 512))
|
| 214 |
for hb, gr in self.plan]
|
| 215 |
P = outs[0] if len(outs) == 1 else torch.cat(outs, 0)
|
| 216 |
if P.shape[0] > 4: # residual/LN kernel takes <= 4 branches
|
build/torch-neuron/nki_kernels/attention_block.py
CHANGED
|
@@ -267,7 +267,7 @@ def pack_out_weights(wo, bo):
|
|
| 267 |
# ==========================================================================================
|
| 268 |
@nki.jit
|
| 269 |
def _attention_block_kernel(h_hbm, wqk_hbm, bqk_hbm, tab_hbm, wvt_hbm, bv_hbm,
|
| 270 |
-
wot_hbm, bo_hbm, scale, head_dim, n_valid, hb, grid, nimg=1, resident=0):
|
| 271 |
"""out[16, N, D] bf16 = attention(rope(h@Wq), rope(h@Wk), h@Wv), one kernel (per-core partials).
|
| 272 |
|
| 273 |
BATCHED (nimg>1): h_hbm is the STACKED [nimg*N, H] buffer (images back to back, N rows each).
|
|
@@ -384,8 +384,18 @@ def _attention_block_kernel(h_hbm, wqk_hbm, bqk_hbm, tab_hbm, wvt_hbm, bv_hbm,
|
|
| 384 |
nisa.dma_copy(dst=r_bo_bc, src=bo_hbm.ap(pattern=[[0, _P_MAX], [1, H]], offset=0))
|
| 385 |
res_top = ra.get_current_address()
|
| 386 |
|
| 387 |
-
|
| 388 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 389 |
allocator = ModularAllocator(initial_address=res_top)
|
| 390 |
|
| 391 |
# ---- persistent consumer tiles (written by producer, read by attention) ---------
|
|
@@ -416,11 +426,11 @@ def _attention_block_kernel(h_hbm, wqk_hbm, bqk_hbm, tab_hbm, wvt_hbm, bv_hbm,
|
|
| 416 |
for _lh in range(local_heads):
|
| 417 |
kt_ = allocator.alloc_sbuf_tensor(
|
| 418 |
shape=(D, _K_TILE), dtype=nl.bfloat16,
|
| 419 |
-
block_dim=[
|
| 420 |
k_tiles.append(kt_)
|
| 421 |
vt_ = allocator.alloc_sbuf_tensor(
|
| 422 |
shape=(_V_TILE, D), dtype=nl.bfloat16,
|
| 423 |
-
block_dim=[
|
| 424 |
v_tiles.append(vt_)
|
| 425 |
# HIGH WATER MARK for phase 2: k_tiles/v_tiles are the only phase-1 products attention
|
| 426 |
# needs, so every buffer above this address is dead once the producer loop finishes.
|
|
@@ -442,7 +452,7 @@ def _attention_block_kernel(h_hbm, wqk_hbm, bqk_hbm, tab_hbm, wvt_hbm, bv_hbm,
|
|
| 442 |
# combine); its large-tile loop iterates 4x over _LARGE_TILE_SZ-sized buffers.
|
| 443 |
# The N<=2048 path enters NEITHER branch and is bit-for-bit unchanged.
|
| 444 |
_big = N > _PSUM_FMAX * 4 # N>2048 -> the 64-frame geometry
|
| 445 |
-
XT_RING = min(n_tt, 4) if _big else
|
| 446 |
|
| 447 |
# ---- resident producer inputs ----------------------------------------------------
|
| 448 |
# x^T: hidden on partitions, tokens on the free axis, as a per-512-token-block ring.
|
|
@@ -451,17 +461,17 @@ def _attention_block_kernel(h_hbm, wqk_hbm, bqk_hbm, tab_hbm, wvt_hbm, bv_hbm,
|
|
| 451 |
# is a pure dma_transpose (no arithmetic), so any fill ORDER yields identical values.
|
| 452 |
x_t = allocator.alloc_sbuf_tensor(
|
| 453 |
shape=(_P_MAX, _PSUM_FMAX), dtype=nl.bfloat16,
|
| 454 |
-
block_dim=[n_chunks,
|
| 455 |
|
| 456 |
if not _big:
|
| 457 |
# N<=2048: fill the whole buffer up front, exactly as the shipping path did (pure
|
| 458 |
# dma_transpose; block-indexed AP lands at the same addresses as the old (128,N) buffer).
|
| 459 |
-
for s in range(
|
| 460 |
for c in range(n_chunks):
|
| 461 |
nisa.dma_transpose(
|
| 462 |
dst=x_t[c][s].ap([[_PSUM_FMAX, _P_MAX], [1, 1], [1, 1],
|
| 463 |
-
[1, min(_PSUM_FMAX,
|
| 464 |
-
src=h_hbm.ap([[H, min(_PSUM_FMAX,
|
| 465 |
offset=rofs + (s * _PSUM_FMAX) * H + c * _P_MAX),
|
| 466 |
)
|
| 467 |
|
|
@@ -503,8 +513,16 @@ def _attention_block_kernel(h_hbm, wqk_hbm, bqk_hbm, tab_hbm, wvt_hbm, bv_hbm,
|
|
| 503 |
# ring, same gating as x_t: at N<=2048 TAB_RING==n_tt holds every block (== old (128,N)
|
| 504 |
# buffer, bit-for-bit); at N=8192 a rolling 4-block window filled inside the tt loop. Pure
|
| 505 |
# dma_copy -> identical values regardless of fill order.
|
| 506 |
-
if _res:
|
| 507 |
tab = r_tab
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 508 |
else:
|
| 509 |
tab = allocator.alloc_sbuf_tensor(
|
| 510 |
shape=(_P_MAX, _PSUM_FMAX), dtype=_ROPE_DT,
|
|
@@ -535,12 +553,15 @@ def _attention_block_kernel(h_hbm, wqk_hbm, bqk_hbm, tab_hbm, wvt_hbm, bv_hbm,
|
|
| 535 |
# Q staging via private_hbm, one per LOCAL head (vendor segmented-kernel pattern)
|
| 536 |
q_priv = []
|
| 537 |
for _lh in range(local_heads):
|
| 538 |
-
|
|
|
|
|
|
|
|
|
|
| 539 |
|
| 540 |
# ---- producer: per 512-token tile, straight through the pipeline ----------------
|
| 541 |
-
for tt in range(
|
| 542 |
t0 = tt * _PSUM_FMAX
|
| 543 |
-
tw = min(_PSUM_FMAX,
|
| 544 |
|
| 545 |
# N>2048: fill this token-block's x_t / tab slice NOW (rolling ring window). At
|
| 546 |
# N<=2048 the whole buffer was already filled up front (bit-for-bit), so skip.
|
|
@@ -641,12 +662,18 @@ def _attention_block_kernel(h_hbm, wqk_hbm, bqk_hbm, tab_hbm, wvt_hbm, bv_hbm,
|
|
| 641 |
dst=k_tiles[lh][tt][par2 * half:(par2 + 1) * half,
|
| 642 |
0:tw],
|
| 643 |
src=roped[par2][r0:r0 + half, 0:tw])
|
| 644 |
-
|
| 645 |
for par2 in range(2):
|
| 646 |
nisa.dma_copy(
|
| 647 |
-
dst=q_priv[lh][0, par2 * half:(par2 + 1) * half,
|
| 648 |
-
|
| 649 |
src=roped[par2][r0:r0 + half, 0:tw])
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 650 |
|
| 651 |
# ---- PHASE 2: reclaim the dead producer buffers. NO ALIASING. -------------------
|
| 652 |
allocator.set_current_address(phase2_base)
|
|
@@ -671,9 +698,9 @@ def _attention_block_kernel(h_hbm, wqk_hbm, bqk_hbm, tab_hbm, wvt_hbm, bv_hbm,
|
|
| 671 |
# group (feature-major), which the projection contracts over, so O never transposes.
|
| 672 |
# The list handed to the vendor is SLICES into o_feat, so _write_back_o_impl copies
|
| 673 |
# straight into final position. ------------------------------------------------------
|
| 674 |
-
n_grp =
|
| 675 |
o_feat = allocator.alloc_sbuf_tensor(
|
| 676 |
-
shape=(_P_MAX,
|
| 677 |
block_dim=[n_ochunks], num_free_tiles=[n_ochunks], align_to=32)
|
| 678 |
o_sb = []
|
| 679 |
for _lh in range(local_heads):
|
|
@@ -707,7 +734,7 @@ def _attention_block_kernel(h_hbm, wqk_hbm, bqk_hbm, tab_hbm, wvt_hbm, bv_hbm,
|
|
| 707 |
for _i in range(n_vtiles):
|
| 708 |
v1d.append(v_stage[_i])
|
| 709 |
_attention_cte(
|
| 710 |
-
q_priv[lh], None, None,
|
| 711 |
scale=scale, causal_mask=False,
|
| 712 |
tp_q=False, tp_k=False, tp_out=True,
|
| 713 |
cache_softmax=False,
|
|
@@ -719,17 +746,32 @@ def _attention_block_kernel(h_hbm, wqk_hbm, bqk_hbm, tab_hbm, wvt_hbm, bv_hbm,
|
|
| 719 |
)
|
| 720 |
else:
|
| 721 |
# N<=2048: the pinned single-section flatten (h8/h10 engine pins intact, bit-for-bit).
|
| 722 |
-
|
| 723 |
-
|
| 724 |
-
|
| 725 |
-
|
| 726 |
-
|
| 727 |
-
|
| 728 |
-
|
| 729 |
-
|
| 730 |
-
|
| 731 |
-
|
| 732 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 733 |
allocator.set_current_address(attn_base)
|
| 734 |
|
| 735 |
# ---- PARTIAL out projection: tokens on partitions, output columns on the free axis
|
|
@@ -739,7 +781,7 @@ def _attention_block_kernel(h_hbm, wqk_hbm, bqk_hbm, tab_hbm, wvt_hbm, bv_hbm,
|
|
| 739 |
out_sb = allocator.alloc_sbuf_tensor(shape=(_P_MAX, _PSUM_FMAX),
|
| 740 |
dtype=nl.bfloat16, align_to=32)
|
| 741 |
n_ob = (H + _PSUM_FMAX - 1) // _PSUM_FMAX
|
| 742 |
-
for tb in range(
|
| 743 |
for ob in range(n_ob):
|
| 744 |
ow = min(_PSUM_FMAX, H - ob * _PSUM_FMAX) # ragged last column block (H%512!=0)
|
| 745 |
ps_out = nl.ndarray((_P_MAX, _PSUM_FMAX), dtype=nl.float32, buffer=nl.psum,
|
|
@@ -789,6 +831,7 @@ def attention_block(
|
|
| 789 |
grid: int = LNC,
|
| 790 |
nimg: int = 1,
|
| 791 |
resident: int = 0,
|
|
|
|
| 792 |
) -> torch.Tensor:
|
| 793 |
"""Fused norm1'd-hidden -> QKV+rope -> attention -> fused out-projection PARTIALS.
|
| 794 |
|
|
@@ -815,4 +858,4 @@ def attention_block(
|
|
| 815 |
)
|
| 816 |
if bqk.dim() == 3 and bqk.shape[1] != 128:
|
| 817 |
bqk = pack_qk_bias(bqk)
|
| 818 |
-
return _ATTN[grid](h, wqk, bqk, tab, wvt, bv, wot, bo, scale, head_dim, n_valid, hb, grid, nimg, resident)
|
|
|
|
| 267 |
# ==========================================================================================
|
| 268 |
@nki.jit
|
| 269 |
def _attention_block_kernel(h_hbm, wqk_hbm, bqk_hbm, tab_hbm, wvt_hbm, bv_hbm,
|
| 270 |
+
wot_hbm, bo_hbm, scale, head_dim, n_valid, hb, grid, nimg=1, resident=0, pack2=0):
|
| 271 |
"""out[16, N, D] bf16 = attention(rope(h@Wq), rope(h@Wk), h@Wv), one kernel (per-core partials).
|
| 272 |
|
| 273 |
BATCHED (nimg>1): h_hbm is the STACKED [nimg*N, H] buffer (images back to back, N rows each).
|
|
|
|
| 384 |
nisa.dma_copy(dst=r_bo_bc, src=bo_hbm.ap(pattern=[[0, _P_MAX], [1, H]], offset=0))
|
| 385 |
res_top = ra.get_current_address()
|
| 386 |
|
| 387 |
+
# v1.6 (pack2): two images share one producer pass when 2N <= 512 (e.g. 224 px): their rows are contiguous,
|
| 388 |
+
# so one 512-token x^T tile holds both and the Q/K/V projections + RoPE stream 2N columns per matmul.
|
| 389 |
+
# Attention runs per image (separate _attention_cte calls on per-image K/V/Q slices); the out-projection
|
| 390 |
+
# covers both images' rows. Same arithmetic per element -> bit-identical to the unpacked kernel.
|
| 391 |
+
PK = 2 if (pack2 and nimg % 2 == 0 and 2 * N <= _PSUM_FMAX) else 1
|
| 392 |
+
NP = N * PK
|
| 393 |
+
n_ttp = (NP + _PSUM_FMAX - 1) // _PSUM_FMAX
|
| 394 |
+
n_ktp = (NP + _K_TILE - 1) // _K_TILE
|
| 395 |
+
n_vtp = NP // _V_TILE
|
| 396 |
+
n_grp1 = N // _P_MAX
|
| 397 |
+
for img in range(nimg // PK): # unrolled at trace time (python int)
|
| 398 |
+
rofs = img * NP * H # element offset of this image (pair)'s rows in h / result
|
| 399 |
allocator = ModularAllocator(initial_address=res_top)
|
| 400 |
|
| 401 |
# ---- persistent consumer tiles (written by producer, read by attention) ---------
|
|
|
|
| 426 |
for _lh in range(local_heads):
|
| 427 |
kt_ = allocator.alloc_sbuf_tensor(
|
| 428 |
shape=(D, _K_TILE), dtype=nl.bfloat16,
|
| 429 |
+
block_dim=[n_ktp], num_free_tiles=[n_ktp], align_to=32)
|
| 430 |
k_tiles.append(kt_)
|
| 431 |
vt_ = allocator.alloc_sbuf_tensor(
|
| 432 |
shape=(_V_TILE, D), dtype=nl.bfloat16,
|
| 433 |
+
block_dim=[n_vtp], num_free_tiles=[n_vtp])
|
| 434 |
v_tiles.append(vt_)
|
| 435 |
# HIGH WATER MARK for phase 2: k_tiles/v_tiles are the only phase-1 products attention
|
| 436 |
# needs, so every buffer above this address is dead once the producer loop finishes.
|
|
|
|
| 452 |
# combine); its large-tile loop iterates 4x over _LARGE_TILE_SZ-sized buffers.
|
| 453 |
# The N<=2048 path enters NEITHER branch and is bit-for-bit unchanged.
|
| 454 |
_big = N > _PSUM_FMAX * 4 # N>2048 -> the 64-frame geometry
|
| 455 |
+
XT_RING = min(n_tt, 4) if _big else n_ttp # token-block ring depth for x_t / tab
|
| 456 |
|
| 457 |
# ---- resident producer inputs ----------------------------------------------------
|
| 458 |
# x^T: hidden on partitions, tokens on the free axis, as a per-512-token-block ring.
|
|
|
|
| 461 |
# is a pure dma_transpose (no arithmetic), so any fill ORDER yields identical values.
|
| 462 |
x_t = allocator.alloc_sbuf_tensor(
|
| 463 |
shape=(_P_MAX, _PSUM_FMAX), dtype=nl.bfloat16,
|
| 464 |
+
block_dim=[n_chunks, n_ttp], num_free_tiles=[n_chunks, XT_RING], align_to=32)
|
| 465 |
|
| 466 |
if not _big:
|
| 467 |
# N<=2048: fill the whole buffer up front, exactly as the shipping path did (pure
|
| 468 |
# dma_transpose; block-indexed AP lands at the same addresses as the old (128,N) buffer).
|
| 469 |
+
for s in range(n_ttp):
|
| 470 |
for c in range(n_chunks):
|
| 471 |
nisa.dma_transpose(
|
| 472 |
dst=x_t[c][s].ap([[_PSUM_FMAX, _P_MAX], [1, 1], [1, 1],
|
| 473 |
+
[1, min(_PSUM_FMAX, NP - s * _PSUM_FMAX)]], offset=0),
|
| 474 |
+
src=h_hbm.ap([[H, min(_PSUM_FMAX, NP - s * _PSUM_FMAX)], [1, 1], [1, 1], [1, _P_MAX]],
|
| 475 |
offset=rofs + (s * _PSUM_FMAX) * H + c * _P_MAX),
|
| 476 |
)
|
| 477 |
|
|
|
|
| 513 |
# ring, same gating as x_t: at N<=2048 TAB_RING==n_tt holds every block (== old (128,N)
|
| 514 |
# buffer, bit-for-bit); at N=8192 a rolling 4-block window filled inside the tt loop. Pure
|
| 515 |
# dma_copy -> identical values regardless of fill order.
|
| 516 |
+
if _res and PK == 1:
|
| 517 |
tab = r_tab
|
| 518 |
+
elif PK == 2:
|
| 519 |
+
# pack2: the same per-image tables side by side in one 512-wide tile
|
| 520 |
+
tab = allocator.alloc_sbuf_tensor(
|
| 521 |
+
shape=(_P_MAX, _PSUM_FMAX), dtype=_ROPE_DT,
|
| 522 |
+
block_dim=[4, 1], num_free_tiles=[4, 1], align_to=4)
|
| 523 |
+
for t_i in range(4):
|
| 524 |
+
for sub in range(2):
|
| 525 |
+
nisa.dma_copy(dst=tab[t_i][0][:, sub * N:(sub + 1) * N], src=tab_hbm[t_i, 0:_P_MAX, 0:N])
|
| 526 |
else:
|
| 527 |
tab = allocator.alloc_sbuf_tensor(
|
| 528 |
shape=(_P_MAX, _PSUM_FMAX), dtype=_ROPE_DT,
|
|
|
|
| 553 |
# Q staging via private_hbm, one per LOCAL head (vendor segmented-kernel pattern)
|
| 554 |
q_priv = []
|
| 555 |
for _lh in range(local_heads):
|
| 556 |
+
qs_ = []
|
| 557 |
+
for _sub in range(PK):
|
| 558 |
+
qs_.append(nl.ndarray((1, D, N), dtype=nl.bfloat16, buffer=nl.private_hbm))
|
| 559 |
+
q_priv.append(qs_)
|
| 560 |
|
| 561 |
# ---- producer: per 512-token tile, straight through the pipeline ----------------
|
| 562 |
+
for tt in range(n_ttp):
|
| 563 |
t0 = tt * _PSUM_FMAX
|
| 564 |
+
tw = min(_PSUM_FMAX, NP - t0) # ragged last token tile (N % 512 != 0)
|
| 565 |
|
| 566 |
# N>2048: fill this token-block's x_t / tab slice NOW (rolling ring window). At
|
| 567 |
# N<=2048 the whole buffer was already filled up front (bit-for-bit), so skip.
|
|
|
|
| 662 |
dst=k_tiles[lh][tt][par2 * half:(par2 + 1) * half,
|
| 663 |
0:tw],
|
| 664 |
src=roped[par2][r0:r0 + half, 0:tw])
|
| 665 |
+
elif PK == 1:
|
| 666 |
for par2 in range(2):
|
| 667 |
nisa.dma_copy(
|
| 668 |
+
dst=q_priv[lh][0][0, par2 * half:(par2 + 1) * half,
|
| 669 |
+
t0:t0 + tw],
|
| 670 |
src=roped[par2][r0:r0 + half, 0:tw])
|
| 671 |
+
else:
|
| 672 |
+
for par2 in range(2):
|
| 673 |
+
for sub in range(2):
|
| 674 |
+
nisa.dma_copy(
|
| 675 |
+
dst=q_priv[lh][sub][0, par2 * half:(par2 + 1) * half, 0:N],
|
| 676 |
+
src=roped[par2][r0:r0 + half, sub * N:(sub + 1) * N])
|
| 677 |
|
| 678 |
# ---- PHASE 2: reclaim the dead producer buffers. NO ALIASING. -------------------
|
| 679 |
allocator.set_current_address(phase2_base)
|
|
|
|
| 698 |
# group (feature-major), which the projection contracts over, so O never transposes.
|
| 699 |
# The list handed to the vendor is SLICES into o_feat, so _write_back_o_impl copies
|
| 700 |
# straight into final position. ------------------------------------------------------
|
| 701 |
+
n_grp = NP // _P_MAX
|
| 702 |
o_feat = allocator.alloc_sbuf_tensor(
|
| 703 |
+
shape=(_P_MAX, NP), dtype=nl.bfloat16,
|
| 704 |
block_dim=[n_ochunks], num_free_tiles=[n_ochunks], align_to=32)
|
| 705 |
o_sb = []
|
| 706 |
for _lh in range(local_heads):
|
|
|
|
| 734 |
for _i in range(n_vtiles):
|
| 735 |
v1d.append(v_stage[_i])
|
| 736 |
_attention_cte(
|
| 737 |
+
q_priv[lh][0], None, None,
|
| 738 |
scale=scale, causal_mask=False,
|
| 739 |
tp_q=False, tp_k=False, tp_out=True,
|
| 740 |
cache_softmax=False,
|
|
|
|
| 746 |
)
|
| 747 |
else:
|
| 748 |
# N<=2048: the pinned single-section flatten (h8/h10 engine pins intact, bit-for-bit).
|
| 749 |
+
for sub in range(PK):
|
| 750 |
+
if PK == 1:
|
| 751 |
+
k_l = k_tiles[lh]
|
| 752 |
+
v_l = v_tiles[lh]
|
| 753 |
+
o_l = o_sb[lh]
|
| 754 |
+
else:
|
| 755 |
+
# pack2: this image's K columns / V tiles / O slices of the shared buffers
|
| 756 |
+
k_l = [k_tiles[lh][0][:, sub * N:(sub + 1) * N]]
|
| 757 |
+
v_l = []
|
| 758 |
+
for _v in range(n_grp1):
|
| 759 |
+
v_l.append(v_tiles[lh][sub * n_grp1 + _v])
|
| 760 |
+
o_l = []
|
| 761 |
+
for _o in range(n_grp1):
|
| 762 |
+
o_l.append(o_sb[lh][sub * n_grp1 + _o])
|
| 763 |
+
_attention_cte(
|
| 764 |
+
q_priv[lh][sub], None, None,
|
| 765 |
+
scale=scale, causal_mask=False,
|
| 766 |
+
tp_q=False, tp_k=False, tp_out=True,
|
| 767 |
+
cache_softmax=False,
|
| 768 |
+
k_cache_sbuf=k_l,
|
| 769 |
+
v_cache_sbuf=v_l,
|
| 770 |
+
out_o_hbm=o_l,
|
| 771 |
+
init_sbuf_addr=attn_base,
|
| 772 |
+
seqlen_k_valid=n_valid,
|
| 773 |
+
)
|
| 774 |
+
allocator.set_current_address(attn_base)
|
| 775 |
allocator.set_current_address(attn_base)
|
| 776 |
|
| 777 |
# ---- PARTIAL out projection: tokens on partitions, output columns on the free axis
|
|
|
|
| 781 |
out_sb = allocator.alloc_sbuf_tensor(shape=(_P_MAX, _PSUM_FMAX),
|
| 782 |
dtype=nl.bfloat16, align_to=32)
|
| 783 |
n_ob = (H + _PSUM_FMAX - 1) // _PSUM_FMAX
|
| 784 |
+
for tb in range(NP // _P_MAX):
|
| 785 |
for ob in range(n_ob):
|
| 786 |
ow = min(_PSUM_FMAX, H - ob * _PSUM_FMAX) # ragged last column block (H%512!=0)
|
| 787 |
ps_out = nl.ndarray((_P_MAX, _PSUM_FMAX), dtype=nl.float32, buffer=nl.psum,
|
|
|
|
| 831 |
grid: int = LNC,
|
| 832 |
nimg: int = 1,
|
| 833 |
resident: int = 0,
|
| 834 |
+
pack2: int = 0,
|
| 835 |
) -> torch.Tensor:
|
| 836 |
"""Fused norm1'd-hidden -> QKV+rope -> attention -> fused out-projection PARTIALS.
|
| 837 |
|
|
|
|
| 858 |
)
|
| 859 |
if bqk.dim() == 3 and bqk.shape[1] != 128:
|
| 860 |
bqk = pack_qk_bias(bqk)
|
| 861 |
+
return _ATTN[grid](h, wqk, bqk, tab, wvt, bv, wot, bo, scale, head_dim, n_valid, hb, grid, nimg, resident, pack2)
|