gqa(adr): TL;DR — Q-replicated/M-fold wording + pointer to S5.6 decode variants

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) <noreply@anthropic.com>
This commit is contained in:
2026-06-04 22:01:25 -07:00
parent 943626d758
commit 444810c6e3
@@ -42,11 +42,12 @@ thing*:
- **Decode** (T_q=1): the output `O = [G, d]` is tiny, the KV cache is big. - **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 → 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 only the small `(m,,O)` via a **2-level reduce** (§4). **Q is replicated**
**replicated** (folded into the matmul M dim). No KV movement. — 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 - **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** would move `[S,d]` per rank. → instead make **each CUBE own one Q head**
(head-parallel) and **rotate the KV** so each head sees all KV (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 (**Ring KV**, §5.5). No `(m,,O)` reduce (each CUBE outputs a different
head). Only KV blocks move. 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): 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) cube_id = tl.program_id(axis=1); pe_id = tl.program_id(axis=0)
kv = head_of_group(cube_id) 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 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) for j in range(ceil(S_kv_local / TILE)): # sweep ONLY my KV shard (resident, no move)
with tl.scratch_scope(): 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 > transposed to sidestep the reshape-not-transpose caveat (§3, §B); no
> bespoke "flash-composite" kind (§8 item 4). **Output head distribution > bespoke "flash-composite" kind (§8 item 4). **Output head distribution
> differs** — decode lands all `G` heads at the CUBE-Group root; prefill > 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`).
--- ---