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

DOCUMENT: TMA 8-mode indexing — the bug that cost us a full day. README + inline comments.

Browse Source
This commit is contained in:
biondizzle
2026-05-22 21:28:58 +00:00
parent 3a4524c318
commit dbd77f2bc4
2 changed files with 73 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
tests/unit/test_fmha_v3_diag.py
View File

@@ -138,7 +138,15 @@ class FmhaV3Diag:
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))
tAgQ = tAgQ[(None,0,None,0)]; tBgK = tBgK[(None,None,None,None,None,None,None,None)]; tVgV = tVgV[(None,None,None,None,None,None,None,None)]
# =====================================================================
# ⚠️⚠️⚠️ CRITICAL: TMA PARTITION TENSOR MODE ORDERING ⚠️⚠️⚠️
# After tma_partition, tBgK/tVgV have 8 modes: (1,1,1,1,n_kv_tiles,1,1,1)
# Mode 4 is the GMEM tile dimension. DO NOT pre-slice to fewer modes!
# See README.md for details.
# =====================================================================
tAgQ = tAgQ[(None,0,None,0)]
tBgK = tBgK[(None,None,None,None,None,None,None,None)] # 8 modes! No pre-slice!
tVgV = tVgV[(None,None,None,None,None,None,None,None)] # 8 modes! No pre-slice!
tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK)
tCrV = pv_mma.make_fragment_B(sV)
@@ -171,6 +179,8 @@ class FmhaV3Diag:
kv_coord = Int32(0 + 0)
for kt in cutlass.range(self.n_kv_tiles, unroll=1):
kvh = kvp.acquire_and_advance(pk)
# ⚠️ CRITICAL: kv_coord indexes MODE 4 of 8-mode tBgK/tVgV.
# Using (None, kv_coord) on a pre-sliced 4-mode tensor SILENTLY BREAKS multi-tile!
cute.copy(tma_k, tBgK[None, None, None, None, kv_coord, None, None, None], tBsK[(None, kvh.index)], tma_bar_ptr=kvh.barrier)
cute.copy(tma_v, tVgV[None, None, None, None, kv_coord, None, None, None], tVsV[(None, kvh.index)], tma_bar_ptr=kvh.barrier)
kv_coord += 1