test: fix NVFP4 mega_moe test dimensions for SMEM alignment

This commit is contained in:
2026-05-12 15:07:35 +00:00
parent 8737fd57c0
commit 787d427847

View File

@@ -9,12 +9,14 @@ import os
import sys
def test_nvfp4_mega_moe():
# Small but aligned dimensions
# Use dimensions that satisfy all alignment requirements:
# - hidden and intermediate_hidden must be multiples of 128 and 64
# - block_m will be at least 32 (SMEM alignment: 32 * 64 = 2048 >= 1024)
num_experts = 2
num_tokens = 8 # must be multiple of alignment (8 for block_m=8)
num_tokens = 32 # must be multiple of alignment
top_k = 2
hidden = 256 # must be multiple of 128 and 64
intermediate_hidden = 512 # must be multiple of 128 and 64
hidden = 512 # multiple of 128 and 64
intermediate_hidden = 1024 # multiple of 128 and 64
device = "cuda"
torch.cuda.set_device(0)