From 9c554e3f644b174cb32da98033c71cbda01cd2bb Mon Sep 17 00:00:00 2001 From: Yangwook Date: Thu, 4 Jun 2026 22:40:12 -0700 Subject: [PATCH] =?UTF-8?q?gqa(adr):=20fix=20prefill=20TL;DR=20pseudocode?= =?UTF-8?q?=20=E2=80=94=20K/V=20are=20TCM=20handles=20for=20the=20ring?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The ring rotates K/V over IPCQ, so they must be kernel-held TCM handles (tl.load), not HBM tl.ref streamed inside the composite. Correct the prefill kernel: load K pre-transposed [d, S/C] and pass the TCM-resident Kc/Vc directly to the composite (drop the bogus tl.ref(Kc,(d,TILE_S)) and TILE_S); recv shapes match ([d,S/C] for K, [S/C,d] for V). Comments note why K/V live in TCM. Docs only. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../ADR-0060-algo-gqa-fused-attention-ahbm.md | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/docs/adr-proposed/ADR-0060-algo-gqa-fused-attention-ahbm.md b/docs/adr-proposed/ADR-0060-algo-gqa-fused-attention-ahbm.md index 54076b8..5dee82b 100644 --- a/docs/adr-proposed/ADR-0060-algo-gqa-fused-attention-ahbm.md +++ b/docs/adr-proposed/ADR-0060-algo-gqa-fused-attention-ahbm.md @@ -128,19 +128,21 @@ max/sum/exp — the new part is the flash accumulator, not the ops.) ```python def gqa_prefill_sp(q_ptr, k_ptr, v_ptr, o_ptr, T_q, S_kv_local, d, C, scale, q_block, cube_start, *, tl): i = tl.program_id(axis=1) - cube_start # this CUBE's Q head = its start KV slice - q = tl.load(q_ptr, (T_q, d)) # MY Q head's query rows (resident) - Kc = tl.load(k_ptr, (S_kv_local, d)); Vc = tl.load(v_ptr, (S_kv_local, d)); src = i + q = tl.load(q_ptr, (T_q, d)) # MY Q head's query rows (resident, TCM) + Kc = tl.load(k_ptr, (d, S_kv_local)) # my K slice, pre-stored transposed [d, S/C] → TCM + Vc = tl.load(v_ptr, (S_kv_local, d)); src = i # my V slice [S/C, d] → TCM (K/V in TCM so the ring can send them) m, l, O = init_running(T_q, d) - for step in range(C): # ── Ring KV: rotate blocks around C CUBEs ── + for step in range(C): # ── Ring KV: rotate blocks around C CUBEs over IPCQ ── f = None - if step < C - 1: - tl.send("ring+", Kc); tl.send("ring+", Vc) # rotate out - f = (tl.recv_async("ring-", (S_kv_local, d)), tl.recv_async("ring-", (S_kv_local, d))) + if step < C - 1: # send current TCM block, recv next (overlap) + tl.send("ring+", Kc); tl.send("ring+", Vc) + f = (tl.recv_async("ring-", (d, S_kv_local)), tl.recv_async("ring-", (S_kv_local, d))) if not block_all_future(q_block, slice_pos(src)): # causal skip whole-future blocks - S = tl.composite("gemm", a=q, b=tl.ref(Kc, (d, TILE_S)), epi=[scale]) + S = tl.composite("gemm", a=q, b=Kc, epi=[scale]) # [T_q,d]·[d,S/C] → [T_q,S/C]; Kc is TCM-resident if block_partial(q_block, slice_pos(src)): S = S + causal_mask(q_block, slice_pos(src)) m2 = tl.maximum(m, tl.max(S, -1)); P = tl.exp(S - m2); corr = tl.exp(m - m2) - l = l * corr + tl.sum(P, -1); O = O * corr + tl.composite("gemm", a=P, b=tl.ref(Vc)); m = m2 + l = l * corr + tl.sum(P, -1) + O = O * corr + tl.composite("gemm", a=P, b=Vc); m = m2 # [T_q,S/C]·[S/C,d] → [T_q,d] if f: Kc, Vc = tl.wait(f[0]), tl.wait(f[1]); src = (src - 1) % C tl.store(o_ptr, O / l) # MY Q head's rows — NO reduce ```