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

FIX: tma_partition tensors have 4 modes, not 8. Mode 2 is GMEM tile dim.

Browse Source 
The 8-mode indexing (tBgK[None,None,None,None,kt,None,None,None]) fails at
JIT compilation with 'coord and shape are weakly congruent' error. The actual
MLIR tensor shape is (((64,128),1),?,?,?) — 4 modes, not 8.

The working fix from commit 845ad98 on the B200 used 4-mode indexing all along:
  tBgK[(None, None, kt, 0)] — mode 2 = GMEM tile dim
  tVgV[(None, 0, kt, 0)] — mode 2 = GMEM tile dim

Updated all files: example10, test_fmha_v3_stage_c, README, docstrings.
This commit is contained in:
biondizzle
2026-05-22 23:08:27 +00:00
parent 265c4583b7
commit 782de32644
3 changed files with 64 additions and 76 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
tests/unit/test_fmha_v3_stage_c.py
View File

@@ -179,9 +179,9 @@ class FmhaV3StageCMulti:
b_lay = cute.make_layout(cute.slice_(cl_vmnk,(0,None,0,0)).shape)
tBsK,tBgK = cpasync.tma_partition(tma_k,0,b_lay,cute.group_modes(sK,0,3),cute.group_modes(tCgK,0,3))
tVsV,tVgV = cpasync.tma_partition(tma_v,0,b_lay,cute.group_modes(sV,0,3),cute.group_modes(tCgV,0,3))
# After tma_partition, tBgK/tVgV have 8 modes.
# Mode 4 is the GMEM tile iteration axis (size = n_kv_tiles).
# Do NOT pre-slice — index all 8 modes explicitly in cute.copy.
# After tma_partition, tBgK/tVgV have 4 modes: (((64,128),1),?,?,?).
# Mode 2 is the GMEM tile iteration axis (size = n_kv_tiles).
# Do NOT pre-slice — index all 4 modes explicitly in cute.copy.
# tAgQ is fine with 4-mode slice (Q has only 1 tile).
tAgQ = tAgQ[(None,0,None,0)]
@@ -221,8 +221,8 @@ class FmhaV3StageCMulti:
kvp.reset(); pk = kvp.try_acquire()
for kt in cutlass.range(0, n_kv_tiles, 1, unroll=1):
kvh = kvp.acquire_and_advance(pk)
cute.copy(tma_k, tBgK[None, None, None, None, kt, None, None, None], tBsK[(None, kvh.index)], tma_bar_ptr=kvh.barrier)
cute.copy(tma_v, tVgV[None, None, None, None, kt, None, None, None], tVsV[(None, kvh.index)], tma_bar_ptr=kvh.barrier)
cute.copy(tma_k, tBgK[(None, None, kt, 0)], tBsK[(None, kvh.index)], tma_bar_ptr=kvh.barrier)
cute.copy(tma_v, tVgV[(None, 0, kt, 0)], tVsV[(None, kvh.index)], tma_bar_ptr=kvh.barrier)
pk = cutlass.Boolean(1)
kvp.tail()