biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: all epilogue warps do TMA store, no dynamic if inside method

Browse Source
This commit is contained in:
biondizzle
2026-05-23 01:41:36 +00:00
parent 7123c214bf
commit f9bb2c18e3
1 changed files with 4 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
tests/unit/test_fmha_v3_stage_c.py
View File

@@ -190,11 +190,11 @@ class FmhaV3StageCMulti:
)
epilog_sync_bar.arrive_and_wait()
# All warps in the epilogue group do the TMA store together
c_buffer = 0
if warp_idx == self.epilogue_warp_id[0]:
cute.copy(tma_c, tCsC[(None, c_buffer)], tCgC_tma[(None, 0, 0, 0, 0, 0, 0)])
cute.arch.cp_async_bulk_commit_group()
cute.arch.cp_async_bulk_wait_group(0, read=True)
cute.copy(tma_c, tCsC[(None, c_buffer)], tCgC_tma[(None, 0, 0, 0, 0, 0, 0)])
cute.arch.cp_async_bulk_commit_group()
cute.arch.cp_async_bulk_wait_group(0, read=True)
epilog_sync_bar.arrive_and_wait()
@cute.kernel