From 9b59fc5db514e3748a0f6bb4ac90128602d09c87 Mon Sep 17 00:00:00 2001 From: YAMY <74099316+YAMY1234@users.noreply.github.com> Date: Mon, 21 Sep 2026 23:45:58 -0500 Subject: [PATCH] [ModelOpt][PP] Keep BF16 shared experts out of the NVFP4 fusion so TP1 pipeline stages can load (#40628) --- .../srt/layers/quantization/modelopt_quant.py | 9 ++++++++ .../quantization/test_modelopt_nvfp4.py | 22 +++++++++++++++++++ 2 files changed, 31 insertions(+) diff --git a/python/sglang/srt/layers/quantization/modelopt_quant.py b/python/sglang/srt/layers/quantization/modelopt_quant.py index 129b17b5a..7a2536933 100755 --- a/python/sglang/srt/layers/quantization/modelopt_quant.py +++ b/python/sglang/srt/layers/quantization/modelopt_quant.py @@ -1497,6 +1497,15 @@ class ModelOptFp4Config(ModelOptQuantConfig): def get_min_capability(cls) -> int: return 80 + def can_fuse_shared_expert(self) -> bool: + # A shared-expert body kept BF16 via exclude_modules cannot share the packed + # FP4 FusedMoE buffers. The shared_expert_gate is a separate linear (kept + # BF16 by e.g. Qwen3-Next NVFP4 checkpoints) and must not veto fusion. + return not any( + "shared_expert" in name and "shared_expert_gate" not in name + for name in self.exclude_modules + ) + @staticmethod def common_group_size(cfg: dict) -> int: """Return the unique group_size across the config; raise if missing/mismatched.""" diff --git a/test/registered/unit/layers/quantization/test_modelopt_nvfp4.py b/test/registered/unit/layers/quantization/test_modelopt_nvfp4.py index 923b886a1..8aeed7ae5 100644 --- a/test/registered/unit/layers/quantization/test_modelopt_nvfp4.py +++ b/test/registered/unit/layers/quantization/test_modelopt_nvfp4.py @@ -132,6 +132,28 @@ class TestModelOptNvfp4(CustomTestCase): use_per_token_activation=True, ) + def test_shared_expert_fusion_requires_matching_fp4_precision(self): + quantized_shared = ModelOptFp4Config( + is_checkpoint_nvfp4_serialized=True, + group_size=16, + ) + bf16_shared = ModelOptFp4Config( + is_checkpoint_nvfp4_serialized=True, + group_size=16, + exclude_modules=["model.layers.*.mlp.shared_experts*"], + ) + + gate_only_bf16 = ModelOptFp4Config( + is_checkpoint_nvfp4_serialized=True, + group_size=16, + exclude_modules=["model.layers.*.mlp.shared_expert_gate"], + ) + + self.assertTrue(quantized_shared.can_fuse_shared_expert()) + self.assertFalse(bf16_shared.can_fuse_shared_expert()) + # Only the gate is BF16 (Qwen3-Next NVFP4): the FP4 body still fuses. + self.assertTrue(gate_only_bf16.can_fuse_shared_expert()) + if __name__ == "__main__": unittest.main()