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

D2: fix local_tile coordinate for 4D Q (2 rest modes, not 3)

Browse Source
This commit is contained in:
biondizzle
2026-05-24 23:38:48 +00:00
parent b5cd1b88c9
commit 3e340a0eee
1 changed files with 4 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -191,18 +191,20 @@ class FmhaKernel:
# Print mQ shape at trace time to understand mode structure
print(f"D2 DEBUG: mQ shape={cute.shape(mQ)}, mK shape={cute.shape(mK)}")
# mQ has 4 modes: (n_h, T, hd, 1). Tiler covers (T, hd) = modes 1,2.
# Rest modes are (n_h, 1) = modes 0,3. Coordinate is (head_idx, 0).
# Q: if num_ctas > 1, mQ has a head dimension. local_tile indexes into it.
# K/V: shared (MQA), always coordinate 0.
# For single-CTA (num_ctas=1), head_cta_idx=0 and coordinates are the same as before.
if const_expr(self.num_ctas > 1):
gQ = cute.local_tile(mQ,cute.slice_(self.qk_mma_tiler,(None,0,None)),(head_cta_idx,None,None))
gQ = cute.local_tile(mQ,cute.slice_(self.qk_mma_tiler,(None,0,None)),(head_cta_idx,0))
else:
gQ = cute.local_tile(mQ,cute.slice_(self.qk_mma_tiler,(None,0,None)),(None,None,None))
gK = cute.local_tile(mK,cute.slice_(self.qk_mma_tiler,(0,None,None)),(None,None,None))
gV = cute.local_tile(mV,cute.slice_(self.pv_mma_tiler,(0,None,None)),(None,None,None))
if const_expr(self.num_ctas > 1):
gC = cute.local_tile(mC,cute.slice_(self.pv_mma_tiler,(None,None,0)),(head_cta_idx,None,None))
gC = cute.local_tile(mC,cute.slice_(self.pv_mma_tiler,(None,None,0)),(head_cta_idx,0))
else:
gC = cute.local_tile(mC,cute.slice_(self.pv_mma_tiler,(None,None,0)),(None,None,None))
n_kv_tiles = cute.size(gK, mode=[3])