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

Use tma_partition with group_modes on both sC_flat and gO

Browse Source
This commit is contained in:
biondizzle
2026-05-27 05:31:47 +00:00
parent 603f52de78
commit 2438826eee
1 changed files with 7 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
dsv4/kernels/attention/fmha_smem_acc.py
View File

@@ -624,10 +624,15 @@ class FmhaKernel:
cute.arch.fence_proxy("async.shared", space="cta")
# Step 2: TMA store sC_flat -> GMEM
# Use cute.copy with tma_c directly (AutoCopy-style)
gO = cute.local_tile(mC, cute.slice_(self.pv_mma_tiler, (None, None, 0)), (None, None, None))
# Group modes to match: sC_flat is 2D, gO needs to be grouped to 2D
tOsC, tOgO = cpasync.tma_partition(
tma_c, 0, cute.make_layout(1),
cute.group_modes(sC_flat, 0, 2),
cute.group_modes(gO, 0, len(cute.shape(gO))),
)
if warp_idx == self.epilogue_warp_id[0]:
cute.copy(tma_c, sC_flat, gO)
cute.copy(tma_c, tOsC, tOgO)
cute.arch.cp_async_bulk_commit_group()
cute.arch.cp_async_bulk_wait_group(0, read=True)