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

D2: fix gC coordinate to match 5-mode flat_divide result

Browse Source
This commit is contained in:
biondizzle
2026-05-24 23:44:39 +00:00
parent a3559538cf
commit 34f5beb767
1 changed files with 2 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
dsv4/kernels/attention/fmha.py
View File

@@ -205,7 +205,8 @@ class FmhaKernel:
gQ = tC_gQ[None, None, head_cta_idx, None, None, None]
c_tiler = cute.slice_(self.pv_mma_tiler, (None, None, 0)) # (128, hd)
tC_gC = cute.flat_divide(mC, c_tiler)
gC = tC_gC[None, None, head_cta_idx, None, None, None]
print(f"D2 DEBUG: tC_gC shape={cute.shape(tC_gC)}")
gC = tC_gC[None, None, head_cta_idx, None, None]
else:
gQ = cute.local_tile(mQ,cute.slice_(self.qk_mma_tiler,(None,0,None)),(None,None,None))
gC = cute.local_tile(mC,cute.slice_(self.pv_mma_tiler,(None,None,0)),(None,None,None))