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, 06:15 UTC
[Bugfix] Handle num_expert_group=None in flashinfer block-scale FP8 MoE (#34494)
Signed-off-by: haosdent <haosdent@gmail.com>
haosdent committed
Feb 15, 2026 at 15:25 UTC
79f3fab05a2d88a5db73591cd3b8afdf956ee723
2 files changed
+78
tests/kernels/moe/test_flashinfer.py
+77
@@ -398,3 +398,80 @@ 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/flashinfer_trtllm_moe.py
+1
@@ -201,6 +201,7 @@ def flashinfer_fused_moe_blockscale_fp8(
201
) -> torch.Tensor:
202
from vllm.utils.flashinfer import flashinfer_trtllm_fp8_block_scale_moe
203
204
+ num_expert_group = num_expert_group if num_expert_group is not None else 0
205
topk_group = topk_group if topk_group is not None else 0
206
assert top_k <= global_num_experts
207
assert top_k <= 10