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 1, 22:30 UTC
[Bugfix] Fix test_cutlass_moe.py (#44380)
Signed-off-by: Bill Nell <bnell@redhat.com>
bnellnm committed
Jun 4, 2026 at 14:18 UTC
439203d32c642db4e7db04847a0e231b728628e0
3 files changed
+20
-9
tests/kernels/moe/test_cutlass_moe.py
+13
-5
@@ -205,7 +205,10 @@ def run_with_expert_maps(
205
w2 = kwargs["w2"]
206
a = kwargs["hidden_states"]
207
moe_config = make_dummy_moe_config(
208
- num_experts=w2.shape[0],
208
+ max_num_tokens=kwargs.get("hidden_states").shape[0],
209
+ experts_per_token=kwargs.get("topk_ids").shape[1],
210
+ num_experts=num_experts,
211
+ num_local_experts=num_local_experts,
212
hidden_dim=w2.shape[1],
213
intermediate_size_per_partition=w2.shape[2],
214
in_dtype=a.dtype,
@@ -258,23 +261,27 @@ def run_8_bit(
261
a1_scale=None,
262
)
263
264
+ num_experts = moe_tensors.w1.size(0) # type: ignore[attr-defined]
265
+ with_ep = num_local_experts is not None or num_local_experts == num_experts
266
+
267
kwargs = {
268
"hidden_states": moe_tensors.a,
269
"w1": moe_tensors.w1_q, # type: ignore[union-attr]
270
"w2": moe_tensors.w2_q, # type: ignore[union-attr]
271
"topk_weights": topk_weights,
272
"topk_ids": topk_ids,
267
- "global_num_experts": moe_tensors.w1_q.shape[0], # type: ignore[union-attr]
273
+ "global_num_experts": num_experts,
274
"activation": MoEActivation.SILU,
275
"expert_map": None,
276
"apply_router_weight_on_input": False,
277
}
278
273
- num_experts = moe_tensors.w1.size(0) # type: ignore[attr-defined]
274
- with_ep = num_local_experts is not None or num_local_experts == num_experts
279
if not with_ep:
280
moe_config = make_dummy_moe_config(
277
- num_experts=moe_tensors.w2_q.shape[0], # type: ignore[union-attr]
281
+ max_num_tokens=moe_tensors.a.shape[0],
282
+ experts_per_token=topk_ids.shape[1],
283
+ num_experts=num_experts,
284
+ num_local_experts=num_local_experts,
285
hidden_dim=moe_tensors.w2_q.shape[1], # type: ignore[union-attr]
286
intermediate_size_per_partition=moe_tensors.w2_q.shape[2], # type: ignore[union-attr]
287
in_dtype=moe_tensors.a.dtype,
@@ -581,6 +588,7 @@ def test_run_cutlass_moe_fp8(
588
per_out_channel,
589
False,
590
topk_weights,
591
+ None,
592
)
593
594
workspace13.random_()
tests/kernels/moe/utils.py
+6
-2
@@ -49,10 +49,12 @@ def shuffle_weight(w: torch.Tensor) -> torch.Tensor:
49
50
def make_dummy_moe_config(
51
num_experts: int = 1,
52
+ num_local_experts: int | None = None,
53
experts_per_token: int = 1,
54
hidden_dim: int = 1,
55
intermediate_size_per_partition: int = 1,
56
in_dtype: torch.dtype = torch.bfloat16,
57
+ max_num_tokens: int = 512,
58
) -> FusedMoEConfig:
59
"""
60
This is a dummy config for the mk constructor interface
@@ -66,14 +68,16 @@ def make_dummy_moe_config(
68
experts_per_token=experts_per_token,
69
hidden_dim=hidden_dim,
70
intermediate_size_per_partition=intermediate_size_per_partition,
69
- num_local_experts=num_experts,
71
+ num_local_experts=num_local_experts
72
+ if num_local_experts is not None
73
+ else num_experts,
74
num_logical_experts=num_experts,
75
moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),
76
activation=MoEActivation.SILU,
77
in_dtype=in_dtype,
78
device="cuda",
79
routing_method=RoutingMethodType.TopK,
76
- max_num_tokens=512,
80
+ max_num_tokens=max_num_tokens,
81
)
82
83
vllm/model_executor/layers/fused_moe/experts/cutlass_moe.py
+1
-2
@@ -379,8 +379,7 @@ class CutlassExpertsFp8Base(mk.FusedMoEExpertsModular):
379
topk_ids,
380
activation,
381
global_num_experts,
382
- # the fp8 cutlass experts use their own expert map.
383
- None,
382
+ expert_map,
383
self.w1_scale,
384
self.w2_scale,
385
a1q_scale,