[Bugfix] Fix /moe/test_flashinfer.py

Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>

tests/kernels/moe/test_flashinfer.py

@@ -412,7 +412,7 @@ def test_flashinfer_blockscale_fp8_none_expert_group(monkeypatch):
         pytest.skip("Test requires SM >= 100 (Blackwell)")
 
     import vllm.model_executor.layers.fused_moe.flashinfer_trtllm_moe  # noqa: E501, F401
-    from tests.kernels.quant_utils import native_per_token_group_quant_fp8
+    from vllm.utils.deep_gemm import per_block_cast_to_fp8
 
     set_random_seed(7)
     monkeypatch.setenv("VLLM_FUSED_MOE_CHUNK_SIZE", "8192")
@@ -421,7 +421,6 @@ def test_flashinfer_blockscale_fp8_none_expert_group(monkeypatch):
     topk = 6  # top_k > 1 triggers DeepSeekV3 routing with sigmoid
     m, n, k = 10, 4096, 5120
     block_shape = [128, 128]
-    block_k = block_shape[1]
 
     with set_current_vllm_config(vllm_config):
         # Create BF16 hidden states
@@ -431,15 +430,15 @@ def test_flashinfer_blockscale_fp8_none_expert_group(monkeypatch):
         w13_bf16 = torch.randn((e, 2 * n, k), device="cuda", dtype=torch.bfloat16) / 10
         w2_bf16 = torch.randn((e, k, n), device="cuda", dtype=torch.bfloat16) / 10
 
-        # Quantize weights per-block to FP8
+        # Quantize weights per-block to FP8 (2D block quantization)
         w13_fp8_list, w13_scale_list = [], []
         w2_fp8_list, w2_scale_list = [], []
         for i in range(e):
-            wq, ws = native_per_token_group_quant_fp8(w13_bf16[i], block_k)
+            wq, ws = per_block_cast_to_fp8(w13_bf16[i], block_size=block_shape)
             w13_fp8_list.append(wq)
             w13_scale_list.append(ws)
 
-            wq, ws = native_per_token_group_quant_fp8(w2_bf16[i], block_k)
+            wq, ws = per_block_cast_to_fp8(w2_bf16[i], block_size=block_shape)
             w2_fp8_list.append(wq)
             w2_scale_list.append(ws)
 

vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py

@@ -210,9 +210,20 @@ def flashinfer_fused_moe_blockscale_fp8(
     # Routing kernel expects #experts <= #threads 512
     assert global_num_experts <= 512
 
+    # DeepSeekV3 routing requires grouped expert selection (n_group > 0).
+    # When num_expert_group is not provided (e.g. MiniMax-M2.1 which uses
+    # sigmoid + bias without grouped top-k), fall back to RenormalizeNaive
+    # which performs Sigmoid -> TopK -> Renormalize without grouping.
+    if (routing_method_type == RoutingMethodType.DeepSeekV3
+            and num_expert_group == 0):
+        routing_method_type = int(RoutingMethodType.RenormalizeNaive)
+
     # The DeepSeekV3 routing method requires float32 router logits.
+    # Other routing methods require bfloat16.
     if routing_method_type == RoutingMethodType.DeepSeekV3:
         routing_logits = routing_logits.to(torch.float32)
+    else:
+        routing_logits = routing_logits.to(torch.bfloat16)
 
     if routing_bias is not None:
         routing_bias = routing_bias.to(x.dtype)
@@ -310,8 +321,11 @@ def fi_trtllm_fp8_per_tensor_moe(
     from vllm.utils.flashinfer import flashinfer_trtllm_fp8_per_tensor_scale_moe
 
     # The DeepSeekV3 routing method requires float32 router logits.
+    # Other routing methods require bfloat16.
     if routing_method_type == RoutingMethodType.DeepSeekV3:
         routing_logits = routing_logits.to(torch.float32)
+    else:
+        routing_logits = routing_logits.to(torch.bfloat16)
 
     return flashinfer_trtllm_fp8_per_tensor_scale_moe(
         routing_logits=routing_logits,