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

D1.2: fix probe for hd=512 (MMA max N=256, use pv_n_tile)

Browse Source
This commit is contained in:
biondizzle
2026-05-23 06:41:42 +00:00
parent 60824b62db
commit 99000cba8d
1 changed files with 9 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files

30
tests/unit/test_tmem_budget.py
View File

@@ -35,29 +35,31 @@ class BudgetProbe:
cute.printf("hd=%d s_cols=%d", Int32(self.hd), Int32(s_cols))
# PV MMA: (128, hd), TMEM-P (P from TMEM, K-major)
# tcgen05 MMA max N = 256. At hd > 256, skip TMEM-P and SMEM-P full-hd probes.
pv_n_tile = self.hd if self.hd <= 256 else 256
pv_a_major_tmem = cute.nvgpu.OperandMajorMode.K
pv_mma_tmem = utils.sm100.make_trivial_tiled_mma(
BFloat16, BFloat16, pv_a_major_tmem, b_major,
Float32, tcgen05.CtaGroup.ONE, (128, self.hd),
Float32, tcgen05.CtaGroup.ONE, (128, pv_n_tile),
tcgen05.OperandSource.TMEM,
)
pv_thr_tmem = pv_mma_tmem.get_slice(0)
pv_as_tmem = pv_thr_tmem.partition_shape_C((128, self.hd))
pv_as_tmem = pv_thr_tmem.partition_shape_C((128, pv_n_tile))
tOtO_tmem = pv_thr_tmem.make_fragment_C(pv_as_tmem)
o_cols_tmem = find_tmem_tensor_col_offset(tOtO_tmem)
cute.printf("hd=%d o_cols_tmem=%d", Int32(self.hd), Int32(o_cols_tmem))
cute.printf("hd=%d pv_n_tile=%d o_cols_tmem=%d", Int32(self.hd), Int32(pv_n_tile), Int32(o_cols_tmem))
# PV MMA: (128, hd), SMEM-P (P from SMEM, same a_major as Q)
# PV MMA: (128, pv_n_tile), SMEM-P (P from SMEM, same a_major as Q)
pv_mma_smem = utils.sm100.make_trivial_tiled_mma(
BFloat16, BFloat16, a_major, b_major,
Float32, tcgen05.CtaGroup.ONE, (128, self.hd),
Float32, tcgen05.CtaGroup.ONE, (128, pv_n_tile),
tcgen05.OperandSource.SMEM,
)
pv_thr_smem = pv_mma_smem.get_slice(0)
pv_as_smem = pv_thr_smem.partition_shape_C((128, self.hd))
pv_as_smem = pv_thr_smem.partition_shape_C((128, pv_n_tile))
tOtO_smem = pv_thr_smem.make_fragment_C(pv_as_smem)
o_cols_smem = find_tmem_tensor_col_offset(tOtO_smem)
cute.printf("hd=%d o_cols_smem=%d", Int32(self.hd), Int32(o_cols_smem))
cute.printf("hd=%d pv_n_tile=%d o_cols_smem=%d", Int32(self.hd), Int32(pv_n_tile), Int32(o_cols_smem))
# P columns in TMEM (TMEM-P path only)
p_cols = 128 * 16 // 32 # pv_mma_tiler[2] * bf16_width / fp32_width
@@ -74,20 +76,6 @@ class BudgetProbe:
# SMEM-P total (O reuses S space, starting at col 0)
cute.printf("hd=%d SMEM-P_total=%d", Int32(self.hd), Int32(o_cols_smem))
# Split-PV at hd > 256
if self.hd > 256:
pv_n_tile = 256
pv_mma_split = utils.sm100.make_trivial_tiled_mma(
BFloat16, BFloat16, a_major, b_major,
Float32, tcgen05.CtaGroup.ONE, (128, pv_n_tile),
tcgen05.OperandSource.SMEM,
)
pv_thr_split = pv_mma_split.get_slice(0)
pv_as_split = pv_thr_split.partition_shape_C((128, pv_n_tile))
tOtO_split = pv_thr_split.make_fragment_C(pv_as_split)
o_cols_split = find_tmem_tensor_col_offset(tOtO_split)
cute.printf("hd=%d split_pv_256_o_cols=%d", Int32(self.hd), Int32(o_cols_split))
def probe_hd(hd):
print(f"\n=== HEAD_DIM={hd} ===", flush=True)