From ea27e3ddab1212f72e96d245627dfeefee52e096 Mon Sep 17 00:00:00 2001 From: Colin Z <59755453+ColinZ22@users.noreply.github.com> Date: Mon, 17 Aug 2026 22:10:01 -0700 Subject: [PATCH] [AMD] Fix Quark Shared Experts Fusion Gate after load-time-override Removal (#35200) --- .../srt/layers/quantization/quark/quark.py | 23 ------------------- .../srt/models/deepseek_common/utils.py | 10 ++++++++ python/sglang/srt/models/deepseek_v2.py | 9 ++++++++ python/sglang/srt/models/deepseek_v4.py | 9 ++++++++ .../test_deepseek_v4_shared_expert_fusion.py | 23 +++++++++++++++++++ .../test_shared_experts_fusion_gates.py | 19 +++++++++++++++ 6 files changed, 70 insertions(+), 23 deletions(-) diff --git a/python/sglang/srt/layers/quantization/quark/quark.py b/python/sglang/srt/layers/quantization/quark/quark.py index fed208fb5..a9a8acb68 100644 --- a/python/sglang/srt/layers/quantization/quark/quark.py +++ b/python/sglang/srt/layers/quantization/quark/quark.py @@ -336,29 +336,6 @@ class QuarkConfig(QuantizationConfig): if isinstance(self.dequantization_config, Fp8Config): self.weight_block_size = self.dequantization_config.weight_block_size - self._maybe_disable_shared_experts_fusion() - - def _maybe_disable_shared_experts_fusion(self) -> None: - """Turn off shared-expert fusion when the producer keeps shared experts - in a higher precision than the routed experts. - """ - if self.can_fuse_shared_expert(): - return - - from sglang.srt.arg_groups.overrides import declare_load_time_override - - declare_load_time_override( - "QuarkConfig._maybe_disable_shared_experts_fusion", - {"disable_shared_experts_fusion": True}, - ) - logger.info( - "Quark: shared experts are excluded from quantization (kept in " - "a higher precision) while routed experts are quantized; " - "disabling shared experts fusion to avoid loading " - "higher-precision shared experts through the quantized " - "routed-expert path." - ) - @property def quantized_layers(self) -> tuple[list[str], int]: # Consumed by `report_online_quantization` in model_runner. Returns the diff --git a/python/sglang/srt/models/deepseek_common/utils.py b/python/sglang/srt/models/deepseek_common/utils.py index 43a26b275..c4a267781 100644 --- a/python/sglang/srt/models/deepseek_common/utils.py +++ b/python/sglang/srt/models/deepseek_common/utils.py @@ -152,6 +152,16 @@ def is_wint4afp8_or_wint4a16_config( ) or quant_config._is_wint4abf16(weight_quant, input_quant) +def quant_blocks_shared_experts_fusion( + quant_config: Optional[QuantizationConfig], +) -> bool: + """Whether the quantization keeps shared experts at a higher precision than + the routed experts, which would require shared expert fusion to be disabled. + """ + can_fuse_fn = getattr(quant_config, "can_fuse_shared_expert", None) + return can_fuse_fn is not None and not can_fuse_fn() + + def yarn_get_mscale(scale: float = 1, mscale: float = 1) -> float: if scale <= 1: return 1.0 diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 11696d758..f2d76f735 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -192,6 +192,7 @@ from sglang.srt.models.deepseek_common.utils import ( _use_aiter_bpreshuffle_gfx95, _use_aiter_gfx95, is_wint4afp8_or_wint4a16_config, + quant_blocks_shared_experts_fusion, ) from sglang.srt.runtime_context import ( attention_backends, @@ -3015,6 +3016,14 @@ class DeepseekV2ForCausalLM(nn.Module, DeepseekV2WeightLoaderMixin): ``install_shared_experts_fusion_decision``), so it takes the config and quantization it is asked about rather than reading an instance. """ + # Need to disable if quant precision mismatch, even if + # --enforce-shared-experts-fusion is specified + if quant_blocks_shared_experts_fusion(quant_config): + return ( + "Quantization keeps shared experts at a higher precision than the " + "routed experts, so they cannot be fused into the quantized " + "routed-expert path." + ) if get_exec().moe.enforce_shared_experts_fusion: return None if is_sbo_enabled() or is_tbo_enabled(): diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index 20cb35665..d29aee963 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/python/sglang/srt/models/deepseek_v4.py @@ -143,6 +143,7 @@ from sglang.srt.models.deepseek_common.amd.deepseek_v4_fused_mhc import ( from sglang.srt.models.deepseek_common.utils import ( _use_aiter_bpreshuffle_gfx95, is_wint4afp8_or_wint4a16_config, + quant_blocks_shared_experts_fusion, ) from sglang.srt.models.deepseek_v2 import ( ParallelLMHead, @@ -3015,6 +3016,14 @@ class DeepseekV4ForCausalLM(nn.Module): """V4 only fuses when explicitly asked to, and then the checkpoint must carry exactly one shared expert. Asked by the loader before any layer is built.""" + # Need to disable if quant precision mismatch, even if + # --enforce-shared-experts-fusion is specified + if quant_blocks_shared_experts_fusion(quant_config): + return ( + "Quantization keeps shared experts at a higher precision than the " + "routed experts, so they cannot be fused into the quantized " + "routed-expert path." + ) if not get_exec().moe.enforce_shared_experts_fusion: return "Config does not support fused shared expert(s)." if hf_config.n_shared_experts != 1: diff --git a/test/registered/unit/models/test_deepseek_v4_shared_expert_fusion.py b/test/registered/unit/models/test_deepseek_v4_shared_expert_fusion.py index b7486b7e2..792691052 100644 --- a/test/registered/unit/models/test_deepseek_v4_shared_expert_fusion.py +++ b/test/registered/unit/models/test_deepseek_v4_shared_expert_fusion.py @@ -86,6 +86,29 @@ class TestDeepseekV4SharedExpertFusionPolicy(CustomTestCase): SimpleNamespace(n_shared_experts=2), None ) + def test_mixed_precision_quant_vetoes_even_when_enforced(self): + """A precision mismatch causes crash when shared expert fusion is enabled, + so --enforce-shared-experts-fusion must not override it. Guards the gap + where the enforce early-return skipped the quant check entirely.""" + self._publish(enforce=True) + mixed = SimpleNamespace( + get_name=lambda: "quark", can_fuse_shared_expert=lambda: False + ) + self.assertIn( + "higher precision", + DeepseekV4ForCausalLM.shared_experts_fusion_disable_reason( + SimpleNamespace(n_shared_experts=1), mixed + ), + ) + matched = SimpleNamespace( + get_name=lambda: "quark", can_fuse_shared_expert=lambda: True + ) + self.assertIsNone( + DeepseekV4ForCausalLM.shared_experts_fusion_disable_reason( + SimpleNamespace(n_shared_experts=1), matched + ) + ) + def test_dspark_entry_class_uses_the_v4_gate(self): """A DSV4 DSpark draft must inherit the target's default fusion policy.""" self._publish(enforce=False) diff --git a/test/registered/unit/models/test_shared_experts_fusion_gates.py b/test/registered/unit/models/test_shared_experts_fusion_gates.py index 73962523c..38fc544a2 100644 --- a/test/registered/unit/models/test_shared_experts_fusion_gates.py +++ b/test/registered/unit/models/test_shared_experts_fusion_gates.py @@ -117,6 +117,25 @@ class TestDeepseekV2Gate(_FusionGateCase): self._reason(DeepseekV2ForCausalLM, self._config(), moe_ep_size=2) ) + def test_mixed_precision_quant_vetoes_even_when_enforced(self): + """A precision mismatch causes crash when shared expert fusion is enabled, + so --enforce-shared-experts-fusion must not override it. Guards the gap + where the enforce early-return skipped the quant check entirely.""" + from sglang.srt.models.deepseek_v2 import DeepseekV2ForCausalLM + + self._seed(enforce_shared_experts_fusion=True) + mixed = SimpleNamespace( + get_name=lambda: "quark", can_fuse_shared_expert=lambda: False + ) + self.assertIn( + "higher precision", + self._reason(DeepseekV2ForCausalLM, self._config(), mixed), + ) + matched = SimpleNamespace( + get_name=lambda: "quark", can_fuse_shared_expert=lambda: True + ) + self.assertIsNone(self._reason(DeepseekV2ForCausalLM, self._config(), matched)) + class TestGlmMoeLiteGate(_FusionGateCase): def _config(self, **kw):