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:
@@ -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`).
|
||||
|
||||
---
|
||||
|
||||
|
||||
Reference in New Issue
Block a user