From 91e8dc371a8b631475cb620aae35d5628ed27155 Mon Sep 17 00:00:00 2001 From: Sam Date: Thu, 27 Nov 2025 08:53:45 +0800 Subject: [PATCH] [Feat][NVFP4] Enable NVFP4 MoE for Qwen series models (eg. Qwen3-Next) #13761 (#13761) Co-authored-by: Kaixi Hou --- .../srt/layers/moe/fused_moe_triton/layer.py | 16 ++++- python/sglang/srt/server_args.py | 4 +- test/nightly/test_qwen3_fp4_trtllm_gen_moe.py | 61 +++++++++++++++++++ test/run_suite_nightly.py | 1 + 4 files changed, 77 insertions(+), 5 deletions(-) create mode 100644 test/nightly/test_qwen3_fp4_trtllm_gen_moe.py diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py index 0847d62f2..224ab3c7f 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py +++ b/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] diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 4db37be09..5b13ac801 100644 --- a/python/sglang/srt/server_args.py +++ b/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" ): diff --git a/test/nightly/test_qwen3_fp4_trtllm_gen_moe.py b/test/nightly/test_qwen3_fp4_trtllm_gen_moe.py new file mode 100644 index 000000000..f6b562431 --- /dev/null +++ b/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() diff --git a/test/run_suite_nightly.py b/test/run_suite_nightly.py index d978953e7..24f9ac0c0 100644 --- a/test/run_suite_nightly.py +++ b/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),