gqa(adr): opt2 ex_composite decode needs tl.wait() before the reduction

The ex_composite #2 updates acc=(m,l,O) asynchronously on the scheduler and
the kernel never reads its output, so no auto-wait fires; acc is only final
after the last #2 in the serial chain. Add tl.wait() (drain all composites)
before hierarchical_reduce_and_store reads acc. opt1/opt3 don't need it
(their composite outputs are consumed in-iteration -> auto-wait). Docs only.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
2026-06-04 22:49:59 -07:00
parent 9c554e3f64
commit 8de13efb9d
@@ -110,6 +110,8 @@ def gqa_decode_v2(q_ptr, k_ptr, v_ptr, o_ptr, S_kv_local, d, C, P, scale, *, tl)
Sj = tl.composite("gemm", a=q_g, b=tl.ref(k_tile(j), (d, TILE)), epi=[scale]) # #1: reads K (priority) Sj = tl.composite("gemm", a=q_g, b=tl.ref(k_tile(j), (d, TILE)), epi=[scale]) # #1: reads K (priority)
tl.ex_composite("softmax_pv", s=Sj, v=tl.ref(v_tile(j), (TILE, d)), acc=acc, scale=scale) # #2: reads V tl.ex_composite("softmax_pv", s=Sj, v=tl.ref(v_tile(j), (TILE, d)), acc=acc, scale=scale) # #2: reads V
# ↑ both non-blocking; CPU never waits intra-tile → max run-ahead, fewest issues # ↑ both non-blocking; CPU never waits intra-tile → max run-ahead, fewest issues
tl.wait() # ← drain ALL composites: acc is updated async (no auto-wait,
# #2 is a serial chain through acc); only final after the last #2
hierarchical_reduce_and_store(*acc, cube_id, pe_id, C, P, o_base(kv)) hierarchical_reduce_and_store(*acc, cube_id, pe_id, C, P, o_base(kv))
``` ```