[Feat][NVFP4] Enable NVFP4 MoE for Qwen series models (eg. Qwen3-Next) #13761 (#13761)

Co-authored-by: Kaixi Hou <kaixih@nvidia.com>

python/sglang/srt/layers/moe/fused_moe_triton/layer.py

@@ -1118,7 +1118,16 @@ class FlashInferFP4MoE(FusedMoE):
 
         hs_fp4, hs_scale_linear = self._quantize_hidden_states_fp4(hidden_states)
 
-        router_logits = router_logits.to(torch.float32)
+        routing_method_type = self.routing_method_type
+        assert (
+            routing_method_type is not None
+        ), "flashinfer trtllm moe nvfp4 backend has not been adapted for the current moe layer, you can set routing_method_type (See definition of RoutingMethodType please) for the moe layer explicitly for a quick adaptation."
+
+        correction_bias = (
+            None
+            if topk_config.correction_bias is None
+            else topk_config.correction_bias.to(hidden_states.dtype)
+        )
 
         with use_symmetric_memory(
             get_tp_group(), disabled=not is_allocation_symmetric()
@@ -1132,9 +1141,10 @@ class FlashInferFP4MoE(FusedMoE):
             symm_output = torch.empty(
                 num_tokens, hidden_size, dtype=torch.bfloat16, device=hs_fp4.device
             )
+
         result = trtllm_fp4_block_scale_moe(
             routing_logits=router_logits,
-            routing_bias=topk_config.correction_bias.to(hidden_states.dtype),
+            routing_bias=correction_bias,
             hidden_states=hs_fp4,
             hidden_states_scale=hs_scale_linear.view(torch.float8_e4m3fn).flatten(),
             gemm1_weights=self.gemm1_weights_fp4_shuffled.data,
@@ -1162,7 +1172,7 @@ class FlashInferFP4MoE(FusedMoE):
             local_num_experts=self.num_local_experts,
             routed_scaling_factor=self.moe_runner_config.routed_scaling_factor,
             tile_tokens_dim=None,
-            routing_method_type=RoutingMethodType.DeepSeekV3,
+            routing_method_type=routing_method_type,
             do_finalize=True,
             output=symm_output,
         )[0]

python/sglang/srt/server_args.py

@@ -1202,7 +1202,7 @@ class ServerArgs:
                 if self.quantization is None and quant_method is not None:
                     self.quantization = quant_method
                 if (
-                    self.quantization == "fp8"
+                    self.quantization in ("fp8", "modelopt_fp4")
                     and self.moe_a2a_backend == "none"
                     and self.moe_runner_backend == "auto"
                 ):
@@ -1229,7 +1229,7 @@ class ServerArgs:
                 if self.quantization is None and quant_method is not None:
                     self.quantization = quant_method
                 if (
-                    self.quantization == "fp8"
+                    (self.quantization == "fp8" or self.quantization == "modelopt_fp4")
                     and self.moe_a2a_backend == "none"
                     and self.moe_runner_backend == "auto"
                 ):

test/nightly/test_qwen3_fp4_trtllm_gen_moe.py

@@ -0,0 +1,61 @@
+import unittest
+from types import SimpleNamespace
+
+from sglang.srt.utils import kill_process_tree
+from sglang.test.few_shot_gsm8k import run_eval
+from sglang.test.test_utils import (
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    popen_launch_server,
+)
+
+
+class TestFlashinferTrtllmGenMoeBackend(CustomTestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.model = "nvidia/Qwen3-30B-A3B-NVFP4"
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        cls.process = popen_launch_server(
+            cls.model,
+            cls.base_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=[
+                "--moe-runner-backend",
+                "flashinfer_trtllm",
+                "--quantization",
+                "modelopt_fp4",
+                "--trust-remote-code",
+                "--disable-radix-cache",
+                "--max-running-requests",
+                "1024",
+                "--chunked-prefill-size",
+                "16384",
+                "--mem-fraction-static",
+                "0.89",
+                "--max-prefill-tokens",
+                "16384",
+            ],
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.process.pid)
+
+    def test_gsm8k(self):
+        args = SimpleNamespace(
+            num_shots=8,
+            data_path=None,
+            num_questions=1319,
+            max_new_tokens=512,
+            parallel=1319,
+            host="http://127.0.0.1",
+            port=int(self.base_url.split(":")[-1]),
+        )
+        metrics = run_eval(args)
+        print(f"{metrics=}")
+        self.assertGreater(metrics["accuracy"], 0.88)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/run_suite_nightly.py

@@ -24,6 +24,7 @@ suites = {
         TestFile("test_flashinfer_trtllm_gen_attn_backend.py", 300),
         TestFile("test_deepseek_v3_fp4_cutlass_moe.py", 900),
         TestFile("test_fp4_moe.py", 300),
+        TestFile("test_qwen3_fp4_trtllm_gen_moe.py", 300),
     ],
     "nightly-8-gpu-b200": [
         TestFile("test_deepseek_r1_fp8_trtllm_backend.py", 3600),