gqa(adr): fix prefill TL;DR pseudocode — K/V are TCM handles for the ring
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) <noreply@anthropic.com>
This commit is contained in:
@@ -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
|
||||
```
|
||||
|
||||
Reference in New Issue
Block a user