From 99000cba8d56e4314c858ffa17c0e77f5d1f1acc Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sat, 23 May 2026 06:41:42 +0000 Subject: [PATCH] D1.2: fix probe for hd=512 (MMA max N=256, use pv_n_tile) --- tests/unit/test_tmem_budget.py | 30 +++++++++--------------------- 1 file changed, 9 insertions(+), 21 deletions(-) diff --git a/tests/unit/test_tmem_budget.py b/tests/unit/test_tmem_budget.py index b2321215..b380cbc1 100644 --- a/tests/unit/test_tmem_budget.py +++ b/tests/unit/test_tmem_budget.py @@ -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)