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