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 8, 02:30 UTC
[Bugfix][MoE Kernel] Fix incorrect routing selection for models without expert groups (e.g., MiniMax-M2.1) (#34673)
Signed-off-by: wwl2755 <wangwenlong2755@gmail.com> Signed-off-by: Robert Shaw <robshaw@redhat.com> Co-authored-by: Robert Shaw <robshaw@redhat.com> Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com>
Wenlong Wang committed
Feb 18, 2026 at 13:03 UTC
847a57cd1217c6fd51f8e1813af5fed31bae1a2f
5 files changed
+32
-96
tests/kernels/moe/test_flashinfer.py
-77
@@ -398,80 +398,3 @@ def test_convert_moe_weights_to_flashinfer_trtllm_block_layout(
398
399
assert w13_converted.shape[0] == num_experts
400
assert w2_converted.shape[0] == num_experts
401
-
402
-
403
-def test_flashinfer_blockscale_fp8_none_expert_group(monkeypatch):
404
- """Test that flashinfer_fused_moe_blockscale_fp8 handles num_expert_group=None.
405
-
406
- Regression test for https://github.com/vllm-project/vllm/issues/34477
407
- MiniMax-M2.1 uses sigmoid scoring with e_score_correction_bias but no
408
- grouped top-k, resulting in num_expert_group=None. This triggered a crash
409
- in the flashinfer kernel when DeepSeekV3 routing was selected.
410
- """
411
- if not current_platform.has_device_capability(100):
412
- pytest.skip("Test requires SM >= 100 (Blackwell)")
413
-
414
- import vllm.model_executor.layers.fused_moe.flashinfer_trtllm_moe # noqa: E501, F401
415
- from tests.kernels.quant_utils import native_per_token_group_quant_fp8
416
-
417
- set_random_seed(7)
418
- monkeypatch.setenv("VLLM_FUSED_MOE_CHUNK_SIZE", "8192")
419
-
420
- e = 16 # num_experts (must be divisible by 4)
421
- topk = 6 # top_k > 1 triggers DeepSeekV3 routing with sigmoid
422
- m, n, k = 10, 4096, 5120
423
- block_shape = [128, 128]
424
- block_k = block_shape[1]
425
-
426
- with set_current_vllm_config(vllm_config):
427
- # Create BF16 hidden states
428
- x = torch.randn((m, k), device="cuda", dtype=torch.bfloat16) / 10
429
-
430
- # Create FP8 block-scale quantized weights
431
- w13_bf16 = torch.randn((e, 2 * n, k), device="cuda", dtype=torch.bfloat16) / 10
432
- w2_bf16 = torch.randn((e, k, n), device="cuda", dtype=torch.bfloat16) / 10
433
-
434
- # Quantize weights per-block to FP8
435
- w13_fp8_list, w13_scale_list = [], []
436
- w2_fp8_list, w2_scale_list = [], []
437
- for i in range(e):
438
- wq, ws = native_per_token_group_quant_fp8(w13_bf16[i], block_k)
439
- w13_fp8_list.append(wq)
440
- w13_scale_list.append(ws)
441
-
442
- wq, ws = native_per_token_group_quant_fp8(w2_bf16[i], block_k)
443
- w2_fp8_list.append(wq)
444
- w2_scale_list.append(ws)
445
-
446
- w13_fp8 = torch.stack(w13_fp8_list)
447
- w13_scale = torch.stack(w13_scale_list)
448
- w2_fp8 = torch.stack(w2_fp8_list)
449
- w2_scale = torch.stack(w2_scale_list)
450
-
451
- # DeepSeekV3 routing uses float32 logits + optional bias
452
- routing_logits = torch.randn((m, e), device="cuda", dtype=torch.float32)
453
- routing_bias = torch.randn(e, device="cuda", dtype=torch.float32)
454
-
455
- # This should NOT crash with num_expert_group=None
456
- output = torch.ops.vllm.flashinfer_fused_moe_blockscale_fp8(
457
- routing_logits=routing_logits,
458
- routing_bias=routing_bias,
459
- x=x,
460
- w13_weight=w13_fp8,
461
- w13_weight_scale_inv=w13_scale,
462
- w2_weight=w2_fp8,
463
- w2_weight_scale_inv=w2_scale,
464
- global_num_experts=e,
465
- top_k=topk,
466
- num_expert_group=None,
467
- topk_group=None,
468
- intermediate_size=n,
469
- expert_offset=0,
470
- local_num_experts=e,
471
- block_shape=block_shape,
472
- routing_method_type=RoutingMethodType.DeepSeekV3,
473
- routed_scaling=1.0,
474
- )
475
-
476
- assert output is not None
477
- assert output.shape == (m, k)
vllm/model_executor/layers/fused_moe/config.py
+17
-10
@@ -8,11 +8,7 @@ import torch
8
9
import vllm.envs as envs
10
from vllm.config import ParallelConfig
11
-from vllm.distributed import (
12
- get_dp_group,
13
- get_pcp_group,
14
- get_tensor_model_parallel_rank,
15
-)
11
+from vllm.distributed import get_dp_group, get_pcp_group, get_tensor_model_parallel_rank
12
from vllm.logger import init_logger
13
from vllm.model_executor.layers.fused_moe.activation import MoEActivation
14
from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import (
@@ -126,20 +122,31 @@ class RoutingMethodType(IntEnum):
122
123
124
def get_routing_method_type(
129
- scoring_func: str, top_k: int, renormalize: bool
125
+ scoring_func: str,
126
+ top_k: int,
127
+ renormalize: bool,
128
+ num_expert_group: int | None,
129
+ has_e_score_bias: bool,
130
) -> RoutingMethodType:
131
+ if has_e_score_bias:
132
+ if (num_expert_group or 0) > 0 and scoring_func == "sigmoid":
133
+ return RoutingMethodType.DeepSeekV3
134
+ else:
135
+ return RoutingMethodType.Unspecified
136
+
137
if scoring_func == "sigmoid":
138
if top_k == 1:
139
return RoutingMethodType.Llama4
140
else:
135
- return RoutingMethodType.DeepSeekV3
136
- elif scoring_func == "softmax":
141
+ return RoutingMethodType.Unspecified
142
+
143
+ if scoring_func == "softmax":
144
if renormalize:
145
return RoutingMethodType.Renormalize
146
else:
147
return RoutingMethodType.Default
141
- else:
142
- return RoutingMethodType.Unspecified
148
+
149
+ return RoutingMethodType.Unspecified
150
151
152
@dataclass
vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py
+2
@@ -165,6 +165,8 @@ class FusedTopKBiasRouter(BaseRouter):
165
scoring_func=self.scoring_func,
166
top_k=self.top_k,
167
renormalize=self.renormalize,
168
+ num_expert_group=None,
169
+ has_e_score_bias=True,
170
)
171
172
def _compute_routing(
vllm/model_executor/layers/fused_moe/router/fused_topk_router.py
+2
@@ -142,6 +142,8 @@ class FusedTopKRouter(BaseRouter):
142
scoring_func=self.scoring_func,
143
top_k=self.top_k,
144
renormalize=self.renormalize,
145
+ num_expert_group=None,
146
+ has_e_score_bias=False,
147
)
148
149
def _compute_routing(
vllm/model_executor/layers/fused_moe/router/grouped_topk_router.py
+11
-9
@@ -13,7 +13,10 @@ from vllm.model_executor.custom_op import CustomOp
13
from vllm.model_executor.layers.batch_invariant import (
14
vllm_is_batch_invariant,
15
)
16
-from vllm.model_executor.layers.fused_moe.config import RoutingMethodType
16
+from vllm.model_executor.layers.fused_moe.config import (
17
+ RoutingMethodType,
18
+ get_routing_method_type,
19
+)
20
from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
21
rocm_aiter_grouped_topk,
22
)
@@ -277,16 +280,15 @@ class GroupedTopKRouter(BaseRouter):
280
self.e_score_correction_bias = e_score_correction_bias
281
self.num_fused_shared_experts = num_fused_shared_experts
282
280
- if scoring_func == "sigmoid":
281
- self._routing_method_type = RoutingMethodType.DeepSeekV3
282
- else:
283
- # NOTE: this prohibits the FLASHINFER_TRTLLM kernels from
284
- # being selected, since they only support DeepSeek-style.
285
- self._routing_method_type = RoutingMethodType.Unspecified
286
-
283
@property
284
def routing_method_type(self) -> RoutingMethodType:
289
- return self._routing_method_type
285
+ return get_routing_method_type(
286
+ scoring_func=self.scoring_func,
287
+ top_k=self.top_k,
288
+ renormalize=self.renormalize,
289
+ num_expert_group=self.num_expert_group,
290
+ has_e_score_bias=self.e_score_correction_bias is not None,
291
+ )
292
293
def _compute_routing(
294
self,