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`). ---