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

FIX: 8-None no-op pre-slice opens full TMA coordinate space (8 dims)

Browse Source 
The tma_partition output has 8 TMA coordinate dimensions, not 4.
The Python-visible shape shows 4 modes, but the TMA descriptor uses
8 coordinates. Without the 8-None no-op pre-slice, modes 4-7 are
collapsed and the GMEM tile axis (mode 4) is pinned to 0.

Pattern that works (confirmed on B200 at n=256 in diag test):
  tBgK = tBgK[(None,None,None,None,None,None,None,None)]  # open 8D
  cute.copy(tma_k, tBgK[None,None,None,None,kt,None,None,None], ...)

The old 4-mode indexing tBgK[(None,None,kt,0)] fails with
'rank mismatch: got 2 and 1' because slicing a 4-mode tensor
produces wrong rank for the TMA coordinate space.

Matches working diag test test_fmha_v3_diag.py exactly.
This commit is contained in:
biondizzle
2026-05-22 23:18:40 +00:00
parent 9c5adcee46
commit 30eaba39aa
3 changed files with 73 additions and 67 deletions
Download Patch File Download Diff File Expand all files Collapse all files

62
tests/fmha_v3_stage_c_example10.py
View File

@@ -14,19 +14,18 @@ Three structural rules learned the hard way:
`utils.sm100.get_tmem_load_op` + `get_smem_store_op` works and is what
the CUTLASS Blackwell FMHA reference uses in `correction_rescale`.
(C) tma_partition produces a 4-mode tensor for tBgK/tVgV, not a 2-mode one. After
(C) tma_partition produces a tensor with 8 TMA coordinate dimensions, but only
4 are visible in the Python shape. After
tBsK, tBgK = cpasync.tma_partition(tma_k, 0, b_lay,
group_modes(sK,0,3),
group_modes(tCgK,0,3))
`tBgK` shape is (((64,128),1),?,?,?) with 4 modes. Mode 2 is the
GMEM-tile iteration axis. Pre-slicing with `tBgK[(None,None,0,0)]`
sets mode 2 to 0, so every TMA copy reads tile 0 regardless of
what's passed at mode 1. The bug pretends to be a JIT issue:
dynamic coords seem to be "constant-folded" because the axis they
vary along is pinned to 0 by the pre-slice.
Fix: do not pre-slice. Index all 4 modes explicitly in the producer's
`cute.copy`, putting `kt` at mode 2 and `0` or `None` everywhere else.
`tBgK` has 8 TMA coord modes: (1,1,1,1,n_kv_tiles,1,1,1).
Mode 4 is the GMEM-tile iteration axis.
Pre-slicing with `tBgK[(None,None,0,0)]` collapses the GMEM tile axis to
coordinate 0, so every TMA copy reads tile 0 regardless of the coord
value passed. The 8-None no-op pre-slice opens the full TMA coord space.
Fix: tBgK = tBgK[(None,None,None,None,None,None,None,None)], then
cute.copy(tma_k, tBgK[None,None,None,None,kt,None,None,None], ...)
Kernel structure:
@@ -186,13 +185,15 @@ 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))
# NOTE: after tma_partition, tBgK has 4 modes: (((64,128),1),?,?,?).
# Mode 2 is the GMEM tile-iteration axis (size = n_kv_tiles).
# We previously pre-sliced like tBgK[(None,None,0,0)] which set mode 2
# to 0, so TMA always read from tile 0. Fix: no pre-slice — index
# all 4 modes explicitly in cute.copy, putting kt at mode 2.
# tVgV similarly has 4 modes with mode 2 as the GMEM tile dim.
# tAgQ is fine with 4-mode slice (Q has only 1 tile).
# NOTE: after tma_partition, tBgK/tVgV have 8 TMA coordinate dimensions.
# Shape is (1,1,1,1,n_kv_tiles,1,1,1) in TMA coord space.
# Mode 4 is the GMEM tile-iteration axis.
# The no-op 8-None slice opens up the full TMA coordinate space.
# Without it, only 4 modes are visible and 8-mode indexing fails.
# The old (None,None,0,0) pre-slice collapsed the GMEM tile axis to 0.
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)]
tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK)
tCrV = pv_mma.make_fragment_B(sV)
@@ -216,14 +217,13 @@ class FmhaV3StageCMulti:
pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk)
# ===== TMA LOAD warp — fully unrolled =====
# The original pre-slice tBgK[(None,None,0,0)] pinned mode 2 (the
# GMEM tile dim) to 0, forcing all TMA reads to tile 0. With no
# pre-slice, we index all 4 modes explicitly in cute.copy, putting
# kt at mode 2. The pipeline's acquire/release machinery still
# tracks the kv_stage ring buffer dynamically at runtime, so the
# producer correctly blocks on consumer release when n_kv_tiles >
# kv_stage. The unroll only flattens the LOOP control flow, not
# the synchronization.
# The old pre-slice tBgK[(None,None,0,0)] collapsed the 8-mode TMA
# coordinate space to 4 modes, pinning mode 4 (GMEM tile dim) to 0.
# The 8-None no-op pre-slice opens the full TMA coord space so we
# can index mode 4 with kt in cute.copy. The pipeline's
# acquire/release machinery still tracks the kv_stage ring buffer
# dynamically at runtime, so the producer correctly blocks on
# consumer release when n_kv_tiles > kv_stage.
if warp_idx == self.tma_warp_id:
qp.reset(); qh = qp.acquire_and_advance()
cute.copy(tma_q, tAgQ[(None, 0, 0, 0)], tAsQ[(None, qh.index)], tma_bar_ptr=qh.barrier)
@@ -231,20 +231,18 @@ class FmhaV3StageCMulti:
kvp.reset()
for kt in cutlass.range_constexpr(self.n_kv_tiles):
kvh = kvp.acquire_and_advance()
# tBgK has 4 modes after tma_partition: (((64,128),1),?,?,?).
# Mode 2 is the GMEM tile iteration axis (size = n_kv_tiles).
# The old pre-slice (None,None,0,0) set mode 2 to 0, forcing
# all TMA reads to tile 0. With no pre-slice, we index all
# 4 modes explicitly, putting kt at mode 2.
# 8-mode TMA indexing: mode 4 = GMEM tile axis (size n_kv_tiles).
# The 8-None pre-slice above opened the full coord space.
# kt at mode 4 indexes the correct KV tile in GMEM.
cute.copy(
tma_k,
tBgK[(None, None, kt, 0)],
tBgK[None, None, None, None, kt, None, None, None],
tBsK[(None, kvh.index)],
tma_bar_ptr=kvh.barrier,
)
cute.copy(
tma_v,
tVgV[(None, 0, kt, 0)],
tVgV[None, None, None, None, kt, None, None, None],
tVsV[(None, kvh.index)],
tma_bar_ptr=kvh.barrier,
)

12
tests/unit/test_fmha_v3_stage_c.py
View File

@@ -179,11 +179,13 @@ 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 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.
# After tma_partition, tBgK/tVgV have 8 TMA coordinate dimensions.
# Mode 4 is the GMEM tile iteration axis (size = n_kv_tiles).
# 8-None no-op pre-slice opens the full TMA coord space.
# tAgQ is fine with 4-mode slice (Q has only 1 tile).
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)]
tCrQ = qk_mma.make_fragment_A(sQ); tCrK = qk_mma.make_fragment_B(sK)
tCrV = pv_mma.make_fragment_B(sV)
@@ -221,8 +223,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, 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)
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)
pk = cutlass.Boolean(1)
kvp.tail()