From 8b351bd8710b8e0d55d52c7ecf28f4d00ed9591e Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sat, 23 May 2026 06:36:39 +0000 Subject: [PATCH] fix: OperandMajorMode.MN not .M --- tests/unit/test_tmem_budget.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/unit/test_tmem_budget.py b/tests/unit/test_tmem_budget.py index e0ede7e2..4d7e8422 100644 --- a/tests/unit/test_tmem_budget.py +++ b/tests/unit/test_tmem_budget.py @@ -16,7 +16,7 @@ from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset # For PV TMEM-P: a_major = OperandMajorMode.K (P in TMEM, K-major) # For PV SMEM-P: a_major = OperandMajorMode.M (P from SMEM, M-major) -M = cute.nvgpu.OperandMajorMode.M +MN = cute.nvgpu.OperandMajorMode.MN K = cute.nvgpu.OperandMajorMode.K def probe_hd(hd): @@ -24,7 +24,7 @@ def probe_hd(hd): # QK MMA: always (128, 128), SMEM A + SMEM B qk_mma = utils.sm100.make_trivial_tiled_mma( - BFloat16, BFloat16, M, K, + BFloat16, BFloat16, MN, K, Float32, tcgen05.CtaGroup.ONE, (128, 128), tcgen05.OperandSource.SMEM, ) @@ -48,7 +48,7 @@ def probe_hd(hd): # PV MMA: (128, hd), SMEM-P (P from SMEM, M-major) pv_mma_smem = utils.sm100.make_trivial_tiled_mma( - BFloat16, BFloat16, M, K, + BFloat16, BFloat16, MN, K, Float32, tcgen05.CtaGroup.ONE, (128, hd), tcgen05.OperandSource.SMEM, ) @@ -86,7 +86,7 @@ def probe_hd(hd): if hd > 256: pv_n_tile = 256 pv_mma_split = utils.sm100.make_trivial_tiled_mma( - BFloat16, BFloat16, M, K, + BFloat16, BFloat16, MN, K, Float32, tcgen05.CtaGroup.ONE, (128, pv_n_tile), tcgen05.OperandSource.SMEM, )