@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / b8c48c5d72
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 14, 01:45 UTC

kernels/moe test pruning (#27053)

Signed-off-by: Fardin Hoque <kfhfar@amazon.com> Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>

Fardin Hoque committed Oct 29, 2025 at 21:10 UTC b8c48c5d722298656074c559d0e8d702a6c28da1
13 files changed +34 -56
tests/kernels/moe/test_batched_moe.py
+18 -7
@@ -24,23 +24,16 @@ from vllm.triton_utils import tl
24
25 MNK_FACTORS = [
26 (1, 128, 128),
27 - (1, 128, 2048),
27 (1, 512, 512),
29 - (1, 1024, 128),
28 (1, 1024, 2048),
29 (32, 128, 128),
30 (32, 512, 512),
31 (32, 1024, 2048),
34 - (45, 128, 128),
32 (45, 128, 2048),
36 - (45, 512, 512),
33 (45, 1024, 128),
38 - (45, 1024, 2048),
34 (64, 512, 512),
35 (64, 1024, 2048),
41 - (222, 128, 128),
36 (222, 128, 2048),
43 - (222, 1024, 128),
37 (222, 1024, 2048),
38 ]
39 NUM_EXPERTS = [8, 64]
@@ -117,10 +110,19 @@ def test_batched_mm(
110 block_shape: list[int] | None,
111 per_act_token_quant: bool,
112 ):
113 + """Note: float8_e4m3fn is not supported on CUDA architecture < 89,
114 + and those tests will be skipped on unsupported hardware."""
115 current_platform.seed_everything(7)
116
117 use_fp8_w8a8 = dtype == torch.float8_e4m3fn
118
119 + if (dtype == torch.float8_e4m3fn) and not current_platform.has_device_capability(
120 + 89
121 + ):
122 + pytest.skip(
123 + "Triton limitation: fp8e4nv data type is not supported on CUDA arch < 89"
124 + )
125 +
126 if (per_act_token_quant or block_shape is not None) and not use_fp8_w8a8:
127 pytest.skip("Don't test blocking for non-quantized types.")
128
@@ -244,10 +246,19 @@ def test_fused_moe_batched_experts(
246 block_shape: list[int] | None,
247 input_scales: bool,
248 ):
249 + """Note: float8_e4m3fn is not supported on CUDA architecture < 89,
250 + and those tests will be skipped on unsupported hardware."""
251 current_platform.seed_everything(7)
252
253 use_fp8_w8a8 = dtype == torch.float8_e4m3fn
254
255 + if (dtype == torch.float8_e4m3fn) and not current_platform.has_device_capability(
256 + 89
257 + ):
258 + pytest.skip(
259 + "Triton limitation: fp8e4nv data type is not supported on CUDA arch < 89"
260 + )
261 +
262 if topk > e:
263 pytest.skip("topk > e")
264
tests/kernels/moe/test_block_fp8.py
-14
@@ -42,57 +42,43 @@ DTYPES = [torch.bfloat16] # [torch.half, torch.bfloat16, torch.float32]
42 # and its hidden size is 7168.
43 MNK_FACTORS = [
44 (1, 128, 128),
45 - (1, 512, 512),
45 (1, 128, 7168),
46 (1, 1024, 7168),
47 (1, 4608, 128),
49 - (1, 4608, 512),
48 (1, 4608, 7168),
49 (83, 128, 128),
50 (83, 512, 512),
53 - (83, 1024, 7168),
51 (83, 4608, 512),
52 (83, 4608, 7168),
56 - (128, 128, 128),
53 (128, 512, 512),
54 (128, 1024, 7168),
59 - (128, 4608, 512),
55 (128, 4608, 7168),
56 (2048, 128, 128),
57 (2048, 1024, 7168),
58 (2048, 4608, 512),
59 (2048, 4608, 7168),
60 (8192, 128, 128),
66 - (8192, 512, 512),
61 (8192, 128, 7168),
62 (8192, 1024, 7168),
69 - (8192, 4608, 512),
63 (8192, 4608, 7168),
64 ]
65
66 MNK_FACTORS_DG = [
67 (128, 128, 128),
75 - (128, 512, 512),
68 (128, 128, 7168),
69 (128, 1024, 7168),
70 (128, 4608, 128),
79 - (128, 4608, 512),
71 (128, 4608, 7168),
81 - (192, 128, 128),
72 (192, 512, 512),
73 (192, 1024, 7168),
84 - (192, 4608, 512),
74 (192, 4608, 7168),
75 (1335, 128, 128),
76 (1335, 1024, 7168),
77 (1335, 4608, 512),
78 (1335, 4608, 7168),
79 (2048, 128, 128),
91 - (2048, 512, 512),
80 (2048, 128, 7168),
81 (2048, 1024, 7168),
94 - (2048, 4608, 128),
95 - (2048, 4608, 512),
82 (2048, 4608, 7168),
83 ]
84
tests/kernels/moe/test_block_int8.py
+1 -9
@@ -21,36 +21,28 @@ vllm_config = VllmConfig()
21 vllm_config.scheduler_config.max_num_seqs = 128
22 vllm_config.scheduler_config.max_model_len = 8192
23
24 -DTYPES = [torch.half, torch.bfloat16]
24 +DTYPES = [torch.bfloat16]
25
26 MNK_FACTORS = [
27 (1, 128, 128),
28 - (1, 512, 512),
28 (1, 128, 7168),
29 (1, 1024, 7168),
31 - (1, 4096, 128),
30 (1, 4096, 512),
31 (1, 4096, 7168),
34 - (33, 128, 128),
32 (33, 512, 512),
33 (33, 128, 7168),
34 (33, 1024, 7168),
35 (33, 4096, 128),
39 - (33, 4096, 512),
36 (33, 4096, 7168),
37 (128, 128, 128),
42 - (128, 512, 512),
38 (128, 1024, 7168),
39 (128, 4096, 512),
40 (128, 4096, 7168),
46 - (222, 128, 128),
41 (222, 512, 512),
42 (222, 1024, 7168),
49 - (222, 4096, 512),
43 (222, 4096, 7168),
44 (2048, 128, 128),
45 (2048, 1024, 7168),
53 - (2048, 4096, 512),
46 (2048, 4096, 4096),
47 ]
48
tests/kernels/moe/test_cutlass_moe.py
-3
@@ -26,16 +26,13 @@ TOP_KS = [6, 8]
26
27 MNK_FACTORS = [
28 (2, 1024, 1024),
29 - (2, 1024, 1536),
29 (2, 3072, 1024),
30 (2, 3072, 1536),
31 (7, 3072, 1536),
32 (64, 1024, 1024),
33 (64, 1024, 1536),
34 (64, 3072, 1024),
36 - (64, 3072, 1536),
35 (224, 1024, 1024),
38 - (224, 1024, 1536),
36 (224, 3072, 1024),
37 (224, 3072, 1536),
38 (32768, 1024, 1024),
tests/kernels/moe/test_deepep_deepgemm_moe.py
-1
@@ -393,7 +393,6 @@ def _test_deepep_deepgemm_moe(
393 MNKs = [
394 (8, 128, 128),
395 (8, 128, 512),
396 - (8, 512, 512),
396 (3, 1024, 2048),
397 (32, 128, 1024),
398 (45, 512, 2048),
tests/kernels/moe/test_deepgemm.py
-2
@@ -130,10 +130,8 @@ def run_single_case(m, n, k, topk, num_experts, block_size):
130 # Note: N <= 512 will disable the deepgemm path due to performance issues.
131 MNKs = [
132 (1024, 768, 128),
133 - (1024, 768, 512),
133 (2048, 768, 512),
134 (512, 1024, 1024),
136 - (512, 2048, 2048),
135 (4096, 4096, 1024),
136 ]
137
tests/kernels/moe/test_flashinfer.py
-2
@@ -34,8 +34,6 @@ TOP_KS = [1]
34
35 MNK_FACTORS = [
36 (256, 8192, 5120),
37 - (256, 4096, 5120),
38 - (127, 8192, 5120),
37 (127, 4096, 5120),
38 (10, 8192, 5120),
39 (10, 4096, 5120),
tests/kernels/moe/test_flashinfer_moe.py
+1 -3
@@ -34,10 +34,8 @@ if not has_flashinfer_cutlass_fused_moe() or not current_platform.has_device_cap
34
35 MNK_FACTORS = [
36 (2, 1024, 1024),
37 - (2, 1024, 1536),
37 (2, 3072, 1024),
38 (2, 3072, 1536),
40 - (64, 1024, 1024),
39 (64, 1024, 1536),
40 (64, 3072, 1024),
41 (64, 2048, 1536),
@@ -49,7 +47,7 @@ MNK_FACTORS = [
47 @pytest.mark.parametrize("m,n,k", MNK_FACTORS)
48 @pytest.mark.parametrize("e", [40, 64, 256])
49 @pytest.mark.parametrize("topk", [1, 6, 8])
52 -@pytest.mark.parametrize("dtype", [torch.half, torch.bfloat16])
50 +@pytest.mark.parametrize("dtype", [torch.bfloat16])
51 @torch.inference_mode()
52 def test_flashinfer_fp4_moe_no_graph(
53 m: int, n: int, k: int, e: int, topk: int, dtype: torch.dtype
tests/kernels/moe/test_grouped_topk.py
+1 -1
@@ -27,7 +27,7 @@ from vllm.platforms import current_platform
27 @pytest.mark.parametrize("topk_group", [2])
28 @pytest.mark.parametrize("scoring_func", ["softmax", "sigmoid"])
29 @pytest.mark.parametrize("routed_scaling_factor", [1.0, 2.5])
30 -@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
30 +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32])
31 def test_grouped_topk(
32 monkeypatch: pytest.MonkeyPatch,
33 n_token: int,
tests/kernels/moe/test_modular_kernel_combinations.py
+8
@@ -295,6 +295,8 @@ def test_modular_kernel_combinations_singlegpu(
295 world_size: int,
296 pytestconfig,
297 ):
298 + """Note: float8_e4m3fn is not supported on CUDA architecture < 89,
299 + and those tests will be skipped on unsupported hardware."""
300 config = Config(
301 Ms=Ms,
302 K=k,
@@ -309,6 +311,12 @@ def test_modular_kernel_combinations_singlegpu(
311 world_size=world_size,
312 )
313
314 + if (
315 + quant_config is not None and quant_config.quant_dtype == torch.float8_e4m3fn
316 + ) and not current_platform.has_device_capability(89):
317 + pytest.skip(
318 + "Triton limitation: fp8e4nv data type is not supported on CUDA arch < 89"
319 + )
320 verbosity = pytestconfig.getoption("verbose")
321 run(config, verbosity > 0)
322
tests/kernels/moe/test_moe.py
+4 -7
@@ -66,8 +66,6 @@ FUSED_MOE_MNK_FACTORS = [
66 (1, 128, 128),
67 (1, 2048, 128),
68 (33, 2048, 128),
69 - (222, 1024, 1024),
70 - (32768, 128, 128),
69 (32768, 2048, 511),
70 (40000, 1024, 1024),
71 ]
@@ -76,7 +74,6 @@ FUSED_MOE_WN16_MNK_FACTORS = [
74 (1, 128, 128),
75 (1, 1024, 1024),
76 (32, 2048, 128),
79 - (32, 1024, 1024),
77 (222, 2048, 1024),
78 ]
79
@@ -512,8 +509,8 @@ def marlin_moe_generate_valid_test_cases():
509 e_list = [4, 12]
510 topk_list = [2, 3]
511 ep_size_list = [1, 4]
515 - dtype_list = [torch.half, torch.bfloat16]
516 - group_size_list = [-1, 16, 32, 128]
512 + dtype_list = [torch.bfloat16]
513 + group_size_list = [-1, 32, 128]
514 act_order_list = [True, False]
515 quant_type_list = [
516 scalar_types.float4_e2m1f,
@@ -885,10 +882,10 @@ def test_batched_moe_align_block_size_opcheck():
882 )
883
884
888 -@pytest.mark.parametrize("m", [1, 33, 64, 222])
885 +@pytest.mark.parametrize("m", [1, 33, 222])
886 @pytest.mark.parametrize("topk", TOP_KS)
887 @pytest.mark.parametrize("k", [128, 511, 1024])
891 -@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
888 +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
889 @pytest.mark.skipif(current_platform.is_rocm(), reason="Skip for rocm")
890 def test_moe_sum(m: int, topk: int, k: int, dtype: torch.dtype):
891 input = torch.randn((m, topk, k), device="cuda", dtype=dtype)
tests/kernels/moe/test_nvfp4_moe.py
+1 -3
@@ -26,9 +26,7 @@ MNK_FACTORS = [
26 (2, 1024, 1024),
27 (2, 1024, 1536),
28 (2, 3072, 1024),
29 - (2, 3072, 1536),
29 (64, 1024, 1024),
31 - (64, 1024, 1536),
30 (64, 3072, 1024),
31 (64, 2048, 1536),
32 (224, 1024, 1024),
@@ -39,7 +37,7 @@ MNK_FACTORS = [
37 @pytest.mark.parametrize("m,n,k", MNK_FACTORS)
38 @pytest.mark.parametrize("e", [40, 64, 256])
39 @pytest.mark.parametrize("topk", [1, 6, 8])
42 -@pytest.mark.parametrize("dtype", [torch.half, torch.bfloat16])
40 +@pytest.mark.parametrize("dtype", [torch.bfloat16])
41 @torch.inference_mode()
42 def test_cutlass_fp4_moe_no_graph(
43 m: int, n: int, k: int, e: int, topk: int, dtype: torch.dtype
tests/kernels/moe/test_silu_mul_fp8_quant_deep_gemm.py
-4
@@ -19,20 +19,16 @@ CASES = [
19 (32, 64, 256, fp8_dtype),
20 (17, 31, 768, fp8_dtype),
21 (1, 1, 128 * 1, fp8_dtype),
22 - (1, 1, 128 * 2, fp8_dtype),
22 (1, 1, 128 * 3, fp8_dtype),
23 (1, 1, 128 * 4, fp8_dtype),
24 (8, 16, 128 * 1, fp8_dtype),
25 (8, 16, 128 * 2, fp8_dtype),
26 (8, 16, 128 * 3, fp8_dtype),
28 - (8, 16, 128 * 4, fp8_dtype),
27 (8, 64, 7168, fp8_dtype),
28 (8, 128, 7168, fp8_dtype),
31 - (8, 256, 7168, fp8_dtype),
29 (8, 512, 7168, fp8_dtype),
30 (8, 1024, 7168, fp8_dtype),
31 (256, 8, 7168, fp8_dtype),
35 - (256, 16, 7168, fp8_dtype),
32 (256, 32, 7168, fp8_dtype),
33 (256, 64, 7168, fp8_dtype),
34 # Only add a few fnuz tests to help with long CI times.