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

fix: OperandMajorMode.MN not .M

Browse Source
This commit is contained in:
biondizzle
2026-05-23 06:36:39 +00:00
parent 5b3bd4876c
commit 8b351bd871
1 changed files with 4 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
tests/unit/test_tmem_budget.py
View File

@@ -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,
)