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, 05:30 UTC
[CI Perf] Prune tests in `tests/kernels/moe/` (#22939)
Signed-off-by: mgoin <mgoin64@gmail.com>
Michael Goin committed
Aug 14, 2025 at 23:33 UTC
d2b0e97ea6538c7f3f38880efc830afd7abfa5d4
6 files changed
+46
-31
tests/kernels/moe/test_batched_moe.py
+5
-8
@@ -89,14 +89,11 @@ class BatchedMMTensors:
89
return BatchedMMTensors(A, B, C, num_expert_tokens)
90
91
92
-@pytest.mark.parametrize("num_experts", [8, 16, 32])
93
-@pytest.mark.parametrize("max_tokens_per_expert",
94
- [32, 64, 128, 192, 224, 256, 512])
95
-@pytest.mark.parametrize("K", [128, 256, 1024])
96
-@pytest.mark.parametrize("N", [128, 256, 1024])
97
-@pytest.mark.parametrize(
98
- "dtype",
99
- [torch.float8_e4m3fn, torch.float32, torch.float16, torch.bfloat16])
92
+@pytest.mark.parametrize("num_experts", [8, 32])
93
+@pytest.mark.parametrize("max_tokens_per_expert", [32, 224, 512])
94
+@pytest.mark.parametrize("K", [128, 1024])
95
+@pytest.mark.parametrize("N", [128, 1024])
96
+@pytest.mark.parametrize("dtype", [torch.float8_e4m3fn, torch.bfloat16])
97
@pytest.mark.parametrize("block_shape", [None, [128, 128]])
98
@pytest.mark.parametrize("per_act_token_quant", [False, True])
99
def test_batched_mm(num_experts: int, max_tokens_per_expert: int, K: int,
tests/kernels/moe/test_count_expert_num_tokens.py
+2
-3
@@ -113,8 +113,7 @@ def do_test_compute_expert_num_tokens(num_tokens: int, num_topk: int,
113
rtol=0)
114
115
116
-@pytest.mark.parametrize(
117
- "num_tokens", [1, 4, 8, 11, 19, 128, 127, 405, 1024, 3333, 6666, 7317])
116
+@pytest.mark.parametrize("num_tokens", [1, 4, 8, 11, 127, 128, 3333, 7317])
117
@pytest.mark.parametrize("num_topk", [2, 6, 8])
118
@pytest.mark.parametrize("num_experts", [64])
119
@pytest.mark.parametrize("ep_size", [1, 2, 4])
@@ -126,7 +125,7 @@ def test_compute_expert_num_tokens(num_tokens: int, num_topk: int,
125
ep_size, topk_ids_dtype)
126
127
129
-@pytest.mark.parametrize("numel", list(range(1, 8192, 11)))
128
+@pytest.mark.parametrize("numel", list(range(1, 8192, 111)))
129
@pytest.mark.parametrize("num_experts", [32])
130
@pytest.mark.parametrize("ep_size", [2])
131
@pytest.mark.parametrize("topk_ids_dtype", [torch.int64])
tests/kernels/moe/test_moe.py
+23
-10
@@ -42,6 +42,24 @@ NUM_EXPERTS = [8, 64, 192]
42
EP_SIZE = [1, 4]
43
TOP_KS = [2, 6]
44
45
+FUSED_MOE_MNK_FACTORS = [
46
+ (1, 128, 128),
47
+ (1, 2048, 128),
48
+ (33, 2048, 128),
49
+ (222, 1024, 1024),
50
+ (32768, 128, 128),
51
+ (32768, 2048, 511),
52
+ (40000, 1024, 1024),
53
+]
54
+
55
+FUSED_MOE_WN16_MNK_FACTORS = [
56
+ (1, 128, 128),
57
+ (1, 1024, 1024),
58
+ (32, 2048, 128),
59
+ (32, 1024, 1024),
60
+ (222, 2048, 1024),
61
+]
62
+
63
vllm_config = VllmConfig()
64
vllm_config.scheduler_config.max_num_seqs = 128
65
vllm_config.scheduler_config.max_model_len = 8192
@@ -116,13 +134,11 @@ def run_moe_test(
134
return baseline_output
135
136
119
-@pytest.mark.parametrize("m", [1, 33, 64, 222, 32768, 40000])
120
-@pytest.mark.parametrize("n", [128, 1024, 2048])
121
-@pytest.mark.parametrize("k", [128, 511, 1024])
137
+@pytest.mark.parametrize("m,n,k", FUSED_MOE_MNK_FACTORS)
138
@pytest.mark.parametrize("e", NUM_EXPERTS)
139
@pytest.mark.parametrize("topk", TOP_KS)
140
@pytest.mark.parametrize("ep_size", EP_SIZE)
125
-@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
141
+@pytest.mark.parametrize("dtype", [torch.bfloat16])
142
@pytest.mark.parametrize("padding", [True, False])
143
@pytest.mark.parametrize("chunk_size", [8192])
144
def test_fused_moe(
@@ -235,13 +251,11 @@ def test_fused_moe(
251
use_cudagraph=use_cudagraph)
252
253
238
-@pytest.mark.parametrize("m", [1, 32, 222])
239
-@pytest.mark.parametrize("n", [128, 1024, 2048])
240
-@pytest.mark.parametrize("k", [128, 1024])
254
+@pytest.mark.parametrize("m,n,k", FUSED_MOE_WN16_MNK_FACTORS)
255
@pytest.mark.parametrize("e", NUM_EXPERTS)
256
@pytest.mark.parametrize("topk", TOP_KS)
257
@pytest.mark.parametrize("ep_size", EP_SIZE)
244
-@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
258
+@pytest.mark.parametrize("dtype", [torch.bfloat16])
259
@pytest.mark.parametrize("group_size", [64, 128])
260
@pytest.mark.parametrize("has_zp", [True, False])
261
@pytest.mark.parametrize("weight_bits", [4, 8])
@@ -352,8 +366,7 @@ def test_fused_moe_wn16(m: int, n: int, k: int, e: int, topk: int,
366
torch.testing.assert_close(triton_output, torch_output, atol=2e-2, rtol=0)
367
368
355
-@pytest.mark.parametrize("dtype",
356
- [torch.float32, torch.float16, torch.bfloat16])
369
+@pytest.mark.parametrize("dtype", [torch.bfloat16])
370
@pytest.mark.parametrize("padding", [True, False])
371
@pytest.mark.parametrize(
372
"use_rocm_aiter", [True, False] if current_platform.is_rocm() else [False])
tests/kernels/moe/test_moe_align_block_size.py
+3
-3
@@ -15,10 +15,10 @@ from vllm.model_executor.layers.fused_moe.moe_align_block_size import (
15
from vllm.platforms import current_platform
16
from vllm.utils import round_up
17
18
-NUM_TOKENS = [1, 3, 7, 16, 256, 2256, 4096]
19
-NUM_EXPERTS = [32, 160, 256, 257, 512]
18
+NUM_TOKENS = [1, 3, 256, 2256, 4096]
19
+NUM_EXPERTS = [32, 160, 256, 257]
20
TOP_KS = [1, 2, 16, 32]
21
-BLOCK_SIZES = [32, 64, 128, 256]
21
+BLOCK_SIZES = [32, 128]
22
current_platform.seed_everything(0)
23
24
tests/kernels/moe/test_moe_permute_unpermute.py
+4
-4
@@ -18,7 +18,7 @@ from vllm.model_executor.layers.fused_moe.moe_permute_unpermute import (
18
from vllm.platforms import current_platform
19
20
NUM_EXPERTS = [16, 64, 256]
21
-TOP_KS = [2, 4, 6, 8]
21
+TOP_KS = [2, 6, 8]
22
EP_SIZE = [1, 4, 16]
23
current_platform.seed_everything(0)
24
@@ -177,11 +177,11 @@ def torch_unpermute(permuted_hidden_states: torch.Tensor,
177
return output
178
179
180
-@pytest.mark.parametrize("n_token", [1, 33, 64, 222, 1024, 2048, 3000, 5000])
181
-@pytest.mark.parametrize("n_hidden", [2048, 4096, 7168])
180
+@pytest.mark.parametrize("n_token", [1, 33, 1024, 5000])
181
+@pytest.mark.parametrize("n_hidden", [2048, 7168])
182
@pytest.mark.parametrize("n_expert", NUM_EXPERTS)
183
@pytest.mark.parametrize("topk", TOP_KS)
184
-@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
184
+@pytest.mark.parametrize("dtype", [torch.bfloat16])
185
@pytest.mark.parametrize("ep_size", EP_SIZE)
186
@pytest.mark.parametrize("align_block_size", [None, 128])
187
def test_moe_permute_unpermute(n_token: int, n_hidden: int, topk: int,
tests/kernels/moe/test_pplx_moe.py
+9
-3
@@ -44,6 +44,14 @@ requires_pplx = pytest.mark.skipif(
44
reason="Requires PPLX kernels",
45
)
46
47
+BATCHED_MOE_MNK_FACTORS = [
48
+ (1, 128, 128),
49
+ (33, 2048, 128),
50
+ (64, 128, 2048),
51
+ (222, 128, 128),
52
+ (222, 2048, 1024),
53
+]
54
+
55
PPLX_COMBOS = [
56
# TODO: figure out why this fails, seems to be test problem
57
#(1, 128, 128),
@@ -152,9 +160,7 @@ def torch_batched_moe(
160
return torch_finalize(out, topk_weight, topk_ids)
161
162
155
-@pytest.mark.parametrize("m", [1, 33, 64, 222])
156
-@pytest.mark.parametrize("n", [128, 1024, 2048])
157
-@pytest.mark.parametrize("k", [128, 512, 1024])
163
+@pytest.mark.parametrize("m,n,k", BATCHED_MOE_MNK_FACTORS)
164
@pytest.mark.parametrize("e", NUM_EXPERTS)
165
@pytest.mark.parametrize("topk", TOP_KS)
166
@pytest.mark.parametrize("dtype", [torch.bfloat16])