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 54076b8..5dee82b 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 @@ -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 ```