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

Switch gK/gV to flat_divide (CUTLASS FMHA reference pattern) for proper TMA strides

Browse Source
This commit is contained in:
biondizzle
2026-05-22 22:30:49 +00:00
parent daa4017505
commit f64af7b0e0
1 changed files with 9 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
tests/unit/test_fmha_v3_stage_c.py
View File

@@ -166,10 +166,16 @@ class FmhaV3StageCMulti:
sC = smem.allocate_tensor(element_type=self.o_dtype,layout=c_smem_s.outer,byte_alignment=128,swizzle=c_smem_s.inner)
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))
# Use flat_divide for gK/gV (matching CUTLASS FMHA reference pattern).
# local_tile produces dynamic Int32(?) modes for the kv_tiles dimension,
# resulting in stride-0 in the TMA partition tensor (broken multi-tile).
# flat_divide produces concrete mode sizes with proper TMA strides.
gK = cute.flat_divide(mK, cute.select(self.qk_mma_tiler, mode=[1, 2]))
gV = cute.flat_divide(mV, cute.select(self.pv_mma_tiler, mode=[1, 2]))
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])
# flat_divide layout: (tile_K, tile_L, n_kv_tiles, ...)
# n_kv_tiles is at mode 2 for K (0-indexed)
n_kv_tiles = cute.size(gK, mode=[2])
qk_thr = qk_mma.get_slice(0); pv_thr = pv_mma.get_slice(0)
tCgQ = qk_thr.partition_A(gQ); tCgK = qk_thr.partition_B(gK)
@@ -181,10 +187,6 @@ class FmhaV3StageCMulti:
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,0,None,0)]; tVgV = tVgV[(None,0,None,0)]
# DEBUG: print shapes and strides to verify kv_tiles dimension has nonzero stride
print(f"tBgK pre-sliced shape: {cute.shape(tBgK)} layout: {tBgK.layout}")
print(f"tVgV pre-sliced shape: {cute.shape(tVgV)} layout: {tVgV.layout}")
tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK)
tCrV = pv_mma.make_fragment_B(sV)