From a2ea5941d562faa5a9a20083ffc81dd0d9dd2d6a Mon Sep 17 00:00:00 2001 From: Zheng Duan <200704041+zhengd-nv@users.noreply.github.com> Date: Sat, 28 Feb 2026 10:28:47 +0800 Subject: [PATCH] [feat] Support nvfp4 quantized model of Qwen3-Next (#17627) --- python/sglang/srt/models/qwen3_next.py | 13 +++- .../models/test_qwen3_next_models_fp4.py | 71 +++++++++++++++++++ 2 files changed, 83 insertions(+), 1 deletion(-) create mode 100644 test/registered/models/test_qwen3_next_models_fp4.py diff --git a/python/sglang/srt/models/qwen3_next.py b/python/sglang/srt/models/qwen3_next.py index 6000e92db..e742dd543 100644 --- a/python/sglang/srt/models/qwen3_next.py +++ b/python/sglang/srt/models/qwen3_next.py @@ -625,13 +625,19 @@ class Qwen3HybridAttentionDecoderLayer(nn.Module): dtype=torch.get_default_dtype(), # see impl of get_rope ) + # qkv_proj is not quantized for fp4 self.qkv_proj = QKVParallelLinear( config.hidden_size, self.head_dim, self.total_num_heads * (1 + self.attn_output_gate), self.total_num_kv_heads, bias=False, - quant_config=quant_config, + quant_config=( + quant_config + if quant_config is not None + and quant_config.get_name() != "modelopt_fp4" + else None + ), tp_rank=self.attn_tp_rank, tp_size=self.attn_tp_size, prefix=add_prefix("qkv_proj", prefix), @@ -1123,6 +1129,11 @@ class Qwen3NextForCausalLM(nn.Module): # if is_pp_missing_parameter(name, self): # continue + if name.endswith("_scale") and name not in params_dict: + assert ( + abs(loaded_weight.item() - 1.0) < 1e-6 + ), f"Expected 1.0, got {loaded_weight.item()} in skipped {name}" + continue param = params_dict[name] weight_loader = getattr( param, "weight_loader", default_weight_loader diff --git a/test/registered/models/test_qwen3_next_models_fp4.py b/test/registered/models/test_qwen3_next_models_fp4.py new file mode 100644 index 000000000..14be5ffa0 --- /dev/null +++ b/test/registered/models/test_qwen3_next_models_fp4.py @@ -0,0 +1,71 @@ +import unittest +from types import SimpleNamespace + +from sglang.srt.utils import get_device_sm, kill_process_tree +from sglang.test.ci.ci_register import register_cuda_ci +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, +) + +register_cuda_ci(est_time=500, suite="nightly-4-gpu-b200", nightly=True) + +QWEN3_NEXT_MODEL_FP4 = "nvidia/Qwen3-Next-80B-A3B-Instruct-NVFP4" + +ACC_THRESHOLDS = { + QWEN3_NEXT_MODEL_FP4: {"kl_div": 0.0025, "gsm8k": 0.93}, +} + + +@unittest.skipIf( + get_device_sm() < 100, "Test requires CUDA SM 100 or higher (Blackwell)" +) +class TestQwen3NextFp4(CustomTestCase): + @classmethod + def setUpClass(cls): + cls.model = QWEN3_NEXT_MODEL_FP4 + 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=[ + "--tp-size", + "4", + "--chunked-prefill-size", + "2048", + "--quantization", + "modelopt_fp4", + "--mamba-scheduler-strategy", + "extra_buffer", + "--mamba-track-interval", + "128", + ], + ) + + @classmethod + def tearDownClass(cls): + kill_process_tree(cls.process.pid) + + def test_gsm8k(self): + args = SimpleNamespace( + num_shots=5, + data_path=None, + num_questions=200, + max_new_tokens=512, + parallel=128, + host="http://127.0.0.1", + port=int(self.base_url.split(":")[-1]), + ) + metrics = run_eval(args) + print(f"{metrics=}") + self.assertGreaterEqual( + metrics["accuracy"], ACC_THRESHOLDS[self.model]["gsm8k"] + ) + + +if __name__ == "__main__": + unittest.main()