Skip to content

Fix SM90 WY input buffer reuse before WGMMA completion #130

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
zheyang0825 merged 1 commit into from
Oct 19, 2026
Merged

Fix SM90 WY input buffer reuse before WGMMA completion #130

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 13 additions & 7 deletions cula/ops/kda/sm90/bwd_wy_dqkg.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2864,15 +2864,17 @@ def _kernel(
)
warpgroup.commit_group()

if v_iter > 0:
warpgroup.wait_group(1)
# This release returns the current stage to DMA. Wait for
# its WGMMA reads, not only the preceding group, before
# allowing the producer to overwrite the shared buffer.
warpgroup.wait_group(0)

pipeline_load_tv.consumer_release(load_tv_cs)
# h NOT released here — reused by dw loop below
load_tv_cs.advance()
load_h_wait_cs.advance()

# Wait for g/k TMA before WGMMA completes — ldmatrix overlaps with WGMMA.
# Wait for g/k TMA before loading the gating operands.
# q is only needed for dg_part1 and is loaded in a separate pass
# after dq is gated.
pipeline_load_g.consumer_wait(load_g_cs)
Expand Down Expand Up @@ -3017,8 +3019,10 @@ def _kernel(
)
warpgroup.commit_group()

if v_iter > 0:
warpgroup.wait_group(1)
# This release returns the current stage to DMA. Wait for
# its WGMMA reads, not only the preceding group, before
# allowing the producer to overwrite the shared buffer.
warpgroup.wait_group(0)

pipeline_load_dv.consumer_release(load_dv_cs)
pipeline_load_h.consumer_release(load_h_release_cs)
Expand Down Expand Up @@ -3089,8 +3093,10 @@ def _kernel(
)
warpgroup.commit_group()

if v_iter > 0:
warpgroup.wait_group(1)
# This release returns the current stage to DMA. Wait for
# its WGMMA reads, not only the preceding group, before
# allowing the producer to overwrite the shared buffer.
warpgroup.wait_group(0)

pipeline_load_tv.consumer_release(load_tv_cs)
pipeline_load_dh.consumer_release(load_dh_cs)
Expand Down
Loading