[AMD] Fix Quark Shared Experts Fusion Gate after load-time-override Removal (#35200)

This commit is contained in:
Colin Z
2026-08-17 22:10:01 -07:00
committed by GitHub
parent 9401db3f29
commit ea27e3ddab
6 changed files with 70 additions and 23 deletions
@@ -336,29 +336,6 @@ class QuarkConfig(QuantizationConfig):
if isinstance(self.dequantization_config, Fp8Config): if isinstance(self.dequantization_config, Fp8Config):
self.weight_block_size = self.dequantization_config.weight_block_size 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 @property
def quantized_layers(self) -> tuple[list[str], int]: def quantized_layers(self) -> tuple[list[str], int]:
# Consumed by `report_online_quantization` in model_runner. Returns the # Consumed by `report_online_quantization` in model_runner. Returns the
@@ -152,6 +152,16 @@ def is_wint4afp8_or_wint4a16_config(
) or quant_config._is_wint4abf16(weight_quant, input_quant) ) 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: def yarn_get_mscale(scale: float = 1, mscale: float = 1) -> float:
if scale <= 1: if scale <= 1:
return 1.0 return 1.0
+9
View File
@@ -192,6 +192,7 @@ from sglang.srt.models.deepseek_common.utils import (
_use_aiter_bpreshuffle_gfx95, _use_aiter_bpreshuffle_gfx95,
_use_aiter_gfx95, _use_aiter_gfx95,
is_wint4afp8_or_wint4a16_config, is_wint4afp8_or_wint4a16_config,
quant_blocks_shared_experts_fusion,
) )
from sglang.srt.runtime_context import ( from sglang.srt.runtime_context import (
attention_backends, attention_backends,
@@ -3015,6 +3016,14 @@ class DeepseekV2ForCausalLM(nn.Module, DeepseekV2WeightLoaderMixin):
``install_shared_experts_fusion_decision``), so it takes the config and ``install_shared_experts_fusion_decision``), so it takes the config and
quantization it is asked about rather than reading an instance. 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: if get_exec().moe.enforce_shared_experts_fusion:
return None return None
if is_sbo_enabled() or is_tbo_enabled(): if is_sbo_enabled() or is_tbo_enabled():
+9
View File
@@ -143,6 +143,7 @@ from sglang.srt.models.deepseek_common.amd.deepseek_v4_fused_mhc import (
from sglang.srt.models.deepseek_common.utils import ( from sglang.srt.models.deepseek_common.utils import (
_use_aiter_bpreshuffle_gfx95, _use_aiter_bpreshuffle_gfx95,
is_wint4afp8_or_wint4a16_config, is_wint4afp8_or_wint4a16_config,
quant_blocks_shared_experts_fusion,
) )
from sglang.srt.models.deepseek_v2 import ( from sglang.srt.models.deepseek_v2 import (
ParallelLMHead, ParallelLMHead,
@@ -3015,6 +3016,14 @@ class DeepseekV4ForCausalLM(nn.Module):
"""V4 only fuses when explicitly asked to, and then the checkpoint must """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 carry exactly one shared expert. Asked by the loader before any layer is
built.""" 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: if not get_exec().moe.enforce_shared_experts_fusion:
return "Config does not support fused shared expert(s)." return "Config does not support fused shared expert(s)."
if hf_config.n_shared_experts != 1: if hf_config.n_shared_experts != 1:
@@ -86,6 +86,29 @@ class TestDeepseekV4SharedExpertFusionPolicy(CustomTestCase):
SimpleNamespace(n_shared_experts=2), None 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): def test_dspark_entry_class_uses_the_v4_gate(self):
"""A DSV4 DSpark draft must inherit the target's default fusion policy.""" """A DSV4 DSpark draft must inherit the target's default fusion policy."""
self._publish(enforce=False) self._publish(enforce=False)
@@ -117,6 +117,25 @@ class TestDeepseekV2Gate(_FusionGateCase):
self._reason(DeepseekV2ForCausalLM, self._config(), moe_ep_size=2) 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): class TestGlmMoeLiteGate(_FusionGateCase):
def _config(self, **kw): def _config(self, **kw):