[Fix] Fix FSE compatibility detection for Quark-produced models (#54573)
Signed-off-by: xuebwang-amd <xuebwang@amd.com> Signed-off-by: Felix Marty <Felix.Marty@amd.com> Co-authored-by: xuebwang-amd <xuebwang@amd.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>
fxmarty-amd committed
Sep 1, 2026 at 19:12 UTC
2fe5cef35eddcaf1c74104002f9abbfe32f9ef4d
3 files changed
+111
-7
tests/model_executor/layers/test_fused_shared_expert.py
+94
-2
@@ -458,7 +458,7 @@ def test_is_model_fused_shared_expert_compatible() -> None:
458
[
459
(False, []),
460
(True, []),
461
- (True, ["*.shared_experts.*"]),
461
+ (True, [r"re:.*\.shared_experts?\..*"]),
462
],
463
)
464
def test_models_fse_init(
@@ -583,7 +583,7 @@ def test_models_fse_init(
583
584
@pytest.mark.parametrize(
585
("exclude", "expected"),
586
- [([], True), (["*.shared_expert.*"], False)],
586
+ [([], True), ([r"re:.*\.shared_expert\..*"], False)],
587
)
588
def test_quark_shared_expert_fse_compatibility(
589
exclude: list[str], expected: bool
@@ -610,6 +610,98 @@ def test_quark_shared_expert_fse_compatibility(
610
)
611
612
613
+def test_quark_shared_expert_fse_exclude_is_scoped_to_the_layer() -> None:
614
+ """Excluding one layer's shared experts must not disable FSE elsewhere.
615
+
616
+ amd/GLM-5.2-MXFP4 excludes its MTP layer (78) wholesale, including
617
+ ``model.layers.78.mlp.shared_experts.{gate,up,down}_proj``, while the 75
618
+ target MoE layers keep MXFP4 shared experts.
619
+ """
620
+ quant_config = QuarkConfig(
621
+ {
622
+ "exclude": [
623
+ f"model.layers.78.mlp.shared_experts.{projection_name}"
624
+ for projection_name in ("gate_proj", "up_proj", "down_proj")
625
+ ],
626
+ "global_quant_config": {},
627
+ "layer_quant_config": {},
628
+ }
629
+ )
630
+ quant_config.packed_modules_mapping = {"gate_up_proj": ["gate_proj", "up_proj"]}
631
+
632
+ assert is_shared_expert_quant_fse_compatible(
633
+ quant_config,
634
+ "model.layers.3.mlp.experts",
635
+ "model.layers.3.mlp.shared_experts",
636
+ ) == (True, None)
637
+
638
+ compatible, reason = is_shared_expert_quant_fse_compatible(
639
+ quant_config,
640
+ "model.layers.78.mlp.experts",
641
+ "model.layers.78.mlp.shared_experts",
642
+ )
643
+ assert compatible is False
644
+ assert (
645
+ reason == "Quark excludes shared experts at model.layers.78.mlp.shared_experts"
646
+ )
647
+
648
+
649
+def test_quark_shared_expert_fse_rejects_partially_excluded_packed_projection() -> None:
650
+ """Packed shared-expert projections require consistent shard exclusions."""
651
+ shared_expert_prefix = "model.layers.0.mlp.shared_expert"
652
+ quant_config = QuarkConfig(
653
+ {
654
+ "exclude": [f"{shared_expert_prefix}.gate_proj"],
655
+ "global_quant_config": {},
656
+ "layer_quant_config": {},
657
+ }
658
+ )
659
+ quant_config.packed_modules_mapping = {"gate_up_proj": ["gate_proj", "up_proj"]}
660
+
661
+ with pytest.raises(ValueError, match="different quantization schemes"):
662
+ is_shared_expert_quant_fse_compatible(
663
+ quant_config, "model.layers.0.mlp.experts", shared_expert_prefix
664
+ )
665
+
666
+
667
+def test_quark_shared_expert_fse_ignores_sibling_gate_exclusions() -> None:
668
+ """Gates beside a shared expert must not disable FSE.
669
+
670
+ Qwen3.x MoE checkpoints (e.g. amd/Qwen3.5-397B-A17B-MoE-MXFP4) quantize the
671
+ shared expert to MXFP4 but leave the tiny routing gates next to it in BF16:
672
+ ``model.language_model.layers.N.mlp.shared_expert_gate`` and
673
+ ``model.language_model.layers.N.mlp.gate``. Neither is a shared-expert
674
+ projection, so fusion must stay enabled. Matching ``shared_expert`` against
675
+ raw exclude entries also matches ``shared_expert_gate``, which silently
676
+ disables FSE on every such checkpoint.
677
+ """
678
+ expert_prefix = "model.language_model.layers.0.mlp.experts"
679
+ shared_expert_prefix = "model.language_model.layers.0.mlp.shared_expert"
680
+ quant_config = QuarkConfig(
681
+ {
682
+ "exclude": [
683
+ "model.language_model.layers.0.mlp.gate",
684
+ f"{shared_expert_prefix}_gate",
685
+ ],
686
+ "global_quant_config": {},
687
+ "layer_quant_config": {},
688
+ }
689
+ )
690
+ quant_config.packed_modules_mapping = {"gate_up_proj": ["gate_proj", "up_proj"]}
691
+
692
+ assert is_shared_expert_quant_fse_compatible(
693
+ quant_config, expert_prefix, shared_expert_prefix
694
+ ) == (True, None)
695
+
696
+ # The gate exclusions must not mask a genuinely unquantized shared expert.
697
+ quant_config.quant_config["exclude"].append(f"{shared_expert_prefix}.down_proj")
698
+ compatible, reason = is_shared_expert_quant_fse_compatible(
699
+ quant_config, expert_prefix, shared_expert_prefix
700
+ )
701
+ assert compatible is False
702
+ assert reason == f"Quark excludes shared experts at {shared_expert_prefix}"
703
+
704
+
705
def test_quark_ocp_mx_group_size_for_layer() -> None:
706
quant_config = QuarkConfig({**_QUARK_FSE_CONFIG, "exclude": []})
707
assert (
vllm/model_executor/layers/fused_moe/utils.py
+1
-1
@@ -86,7 +86,7 @@ def resolve_layer_fused_shared_expert(
86
if fse_requested and not is_fused_shared_expert_enabled:
87
logger.warning(
88
"VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS is enabled but "
89
- "cannot be enabled: %s.",
89
+ "cannot be enabled - skipping for this layer: %s.",
90
fse_reason,
91
)
92
return is_fused_shared_expert_enabled
vllm/model_executor/layers/quantization/utils/config_utils.py
+16
-4
@@ -138,15 +138,27 @@ def is_shared_expert_quant_fse_compatible(
138
)
139
140
if isinstance(quant_config, QuarkConfig):
141
+ from vllm.model_executor.layers.quantization.quark.utils import (
142
+ should_ignore_layer,
143
+ )
144
+
145
# TODO: layer_type_quant_config is not taken into account here.
146
assert "exclude" in quant_config.quant_config
147
assert "global_quant_config" in quant_config.quant_config
148
145
- is_compatible = not any(
146
- "shared_expert" in str(entry)
147
- for entry in quant_config.quant_config["exclude"]
149
+ exclude = quant_config.quant_config["exclude"]
150
+
151
+ # should_ignore_layer raises a rightful ValueError in case different shards
152
+ # have a different ignore policy, as unsupported.
153
+ is_excluded = any(
154
+ should_ignore_layer(
155
+ f"{shared_expert_prefix}.{projection_name}",
156
+ ignore=exclude,
157
+ fused_mapping=quant_config.packed_modules_mapping,
158
+ )
159
+ for projection_name in projection_names
160
)
149
- if not is_compatible:
161
+ if is_excluded:
162
return False, f"Quark excludes shared experts at {shared_expert_prefix}"
163
164
global_quant_config = quant_config.quant_config["global_quant_config"]