From 444810c6e3fb0defe22c174acc0d8b3a806db5c3 Mon Sep 17 00:00:00 2001 From: Yangwook Date: Thu, 4 Jun 2026 22:01:25 -0700 Subject: [PATCH] =?UTF-8?q?gqa(adr):=20TL;DR=20=E2=80=94=20Q-replicated/M-?= =?UTF-8?q?fold=20wording=20+=20pointer=20to=20S5.6=20decode=20variants?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The prior commit updated S0/S10/S5.5 terminology and added S5.6 (3 decode variants) but left the TL;DR with the old 'all G heads replicated' wording and no pointer to S5.6. Sync the TL;DR: 'Q replicated (G heads stacked into the GEMM M-dim, M-fold)', '1 Q head per CUBE', and a one-line pointer to the 3 CPU-pipelining variants in S5.6. Docs only. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../ADR-0060-algo-gqa-fused-attention-ahbm.md | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 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 e8d6fdc..27af36a 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 @@ -42,11 +42,12 @@ thing*: - **Decode** (T_q=1): the output `O = [G, d]` is tiny, the KV cache is big. → keep KV **statically sharded & resident**, do the local sweep, and move - only the small `(m,ℓ,O)` via a **2-level reduce** (§4). `G` heads are - **replicated** (folded into the matmul M dim). No KV movement. + only the small `(m,ℓ,O)` via a **2-level reduce** (§4). **Q is replicated** + — all `G` query heads stacked into the GEMM M-dim (**M-fold**), so one + `Q·Kᵀ` does all heads sharing one `K`. No KV movement. - **Prefill** (T_q=S): the output `O = [S, d]` per head is big — reducing it - would move `[S,d]` per rank. → instead make **each CUBE own one query head** - (head-parallel) and **rotate the KV** so each head sees all KV + would move `[S,d]` per rank. → instead make **each CUBE own one Q head** + (head-parallel, `C=G`) and **rotate the KV** so each head sees all KV (**Ring KV**, §5.5). No `(m,ℓ,O)` reduce (each CUBE outputs a different head). Only KV blocks move. @@ -60,7 +61,7 @@ merge in kernel, lazy `tl.load`). def gqa_decode_sp(q_ptr, k_ptr, v_ptr, o_ptr, S_kv_local, d, C, P, scale, *, tl): cube_id = tl.program_id(axis=1); pe_id = tl.program_id(axis=0) kv = head_of_group(cube_id) - q_g = tl.load(q_base(kv), (G * 1, d)) # T_q=1; all G heads replicated (M-fold) + q_g = tl.load(q_base(kv), (G * 1, d)) # T_q=1; Q replicated: G heads stacked into M (M-fold) m, l, O = init_running(G, d) # persistent arena for j in range(ceil(S_kv_local / TILE)): # sweep ONLY my KV shard (resident, no move) with tl.scratch_scope(): @@ -97,7 +98,9 @@ def gqa_prefill_sp(q_ptr, k_ptr, v_ptr, o_ptr, T_q, S_kv_local, d, C, scale, q_b > transposed to sidestep the reshape-not-transpose caveat (§3, §B); no > bespoke "flash-composite" kind (§8 item 4). **Output head distribution > differs** — decode lands all `G` heads at the CUBE-Group root; prefill -> leaves one Q head per CUBE (§0.5.4). +> leaves one Q head per CUBE (§0.5.4). The decode kernel above is the +> reference shape; **§5.6 gives 3 CPU-pipelining variants** of it +> (current / software-pipelined / `ex_composite`). ---