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

Fix gO creation: use slice_(pv_mma_tiler) like fmha.py

Browse Source
This commit is contained in:
biondizzle
2026-05-27 05:30:50 +00:00
parent b39d7f1a14
commit 603f52de78
1 changed files with 1 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -625,7 +625,7 @@ class FmhaKernel:
# Step 2: TMA store sC_flat -> GMEM
# Use cute.copy with tma_c directly (AutoCopy-style)
gO = cute.local_tile(mC, cute.select(self.pv_mma_tiler, (None, None, 0)), (Int32(0), Int32(0), Int32(0)))
gO = cute.local_tile(mC, cute.slice_(self.pv_mma_tiler, (None, None, 0)), (None, None, None))
if warp_idx == self.epilogue_warp_id[0]:
cute.copy(tma_c, sC_flat, gO)
cute.arch.cp_async_bulk_commit_group()