[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>

tests/model_executor/layers/test_fused_shared_expert.py

@@ -458,7 +458,7 @@ def test_is_model_fused_shared_expert_compatible() -> None:
     [
         (False, []),
         (True, []),
-        (True, ["*.shared_experts.*"]),
+        (True, [r"re:.*\.shared_experts?\..*"]),
     ],
 )
 def test_models_fse_init(
@@ -583,7 +583,7 @@ def test_models_fse_init(
 
 @pytest.mark.parametrize(
     ("exclude", "expected"),
-    [([], True), (["*.shared_expert.*"], False)],
+    [([], True), ([r"re:.*\.shared_expert\..*"], False)],
 )
 def test_quark_shared_expert_fse_compatibility(
     exclude: list[str], expected: bool
@@ -610,6 +610,98 @@ def test_quark_shared_expert_fse_compatibility(
         )
 
 
+def test_quark_shared_expert_fse_exclude_is_scoped_to_the_layer() -> None:
+    """Excluding one layer's shared experts must not disable FSE elsewhere.
+
+    amd/GLM-5.2-MXFP4 excludes its MTP layer (78) wholesale, including
+    ``model.layers.78.mlp.shared_experts.{gate,up,down}_proj``, while the 75
+    target MoE layers keep MXFP4 shared experts.
+    """
+    quant_config = QuarkConfig(
+        {
+            "exclude": [
+                f"model.layers.78.mlp.shared_experts.{projection_name}"
+                for projection_name in ("gate_proj", "up_proj", "down_proj")
+            ],
+            "global_quant_config": {},
+            "layer_quant_config": {},
+        }
+    )
+    quant_config.packed_modules_mapping = {"gate_up_proj": ["gate_proj", "up_proj"]}
+
+    assert is_shared_expert_quant_fse_compatible(
+        quant_config,
+        "model.layers.3.mlp.experts",
+        "model.layers.3.mlp.shared_experts",
+    ) == (True, None)
+
+    compatible, reason = is_shared_expert_quant_fse_compatible(
+        quant_config,
+        "model.layers.78.mlp.experts",
+        "model.layers.78.mlp.shared_experts",
+    )
+    assert compatible is False
+    assert (
+        reason == "Quark excludes shared experts at model.layers.78.mlp.shared_experts"
+    )
+
+
+def test_quark_shared_expert_fse_rejects_partially_excluded_packed_projection() -> None:
+    """Packed shared-expert projections require consistent shard exclusions."""
+    shared_expert_prefix = "model.layers.0.mlp.shared_expert"
+    quant_config = QuarkConfig(
+        {
+            "exclude": [f"{shared_expert_prefix}.gate_proj"],
+            "global_quant_config": {},
+            "layer_quant_config": {},
+        }
+    )
+    quant_config.packed_modules_mapping = {"gate_up_proj": ["gate_proj", "up_proj"]}
+
+    with pytest.raises(ValueError, match="different quantization schemes"):
+        is_shared_expert_quant_fse_compatible(
+            quant_config, "model.layers.0.mlp.experts", shared_expert_prefix
+        )
+
+
+def test_quark_shared_expert_fse_ignores_sibling_gate_exclusions() -> None:
+    """Gates beside a shared expert must not disable FSE.
+
+    Qwen3.x MoE checkpoints (e.g. amd/Qwen3.5-397B-A17B-MoE-MXFP4) quantize the
+    shared expert to MXFP4 but leave the tiny routing gates next to it in BF16:
+    ``model.language_model.layers.N.mlp.shared_expert_gate`` and
+    ``model.language_model.layers.N.mlp.gate``. Neither is a shared-expert
+    projection, so fusion must stay enabled. Matching ``shared_expert`` against
+    raw exclude entries also matches ``shared_expert_gate``, which silently
+    disables FSE on every such checkpoint.
+    """
+    expert_prefix = "model.language_model.layers.0.mlp.experts"
+    shared_expert_prefix = "model.language_model.layers.0.mlp.shared_expert"
+    quant_config = QuarkConfig(
+        {
+            "exclude": [
+                "model.language_model.layers.0.mlp.gate",
+                f"{shared_expert_prefix}_gate",
+            ],
+            "global_quant_config": {},
+            "layer_quant_config": {},
+        }
+    )
+    quant_config.packed_modules_mapping = {"gate_up_proj": ["gate_proj", "up_proj"]}
+
+    assert is_shared_expert_quant_fse_compatible(
+        quant_config, expert_prefix, shared_expert_prefix
+    ) == (True, None)
+
+    # The gate exclusions must not mask a genuinely unquantized shared expert.
+    quant_config.quant_config["exclude"].append(f"{shared_expert_prefix}.down_proj")
+    compatible, reason = is_shared_expert_quant_fse_compatible(
+        quant_config, expert_prefix, shared_expert_prefix
+    )
+    assert compatible is False
+    assert reason == f"Quark excludes shared experts at {shared_expert_prefix}"
+
+
 def test_quark_ocp_mx_group_size_for_layer() -> None:
     quant_config = QuarkConfig({**_QUARK_FSE_CONFIG, "exclude": []})
     assert (

vllm/model_executor/layers/fused_moe/utils.py

@@ -86,7 +86,7 @@ def resolve_layer_fused_shared_expert(
     if fse_requested and not is_fused_shared_expert_enabled:
         logger.warning(
             "VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS is enabled but "
-            "cannot be enabled: %s.",
+            "cannot be enabled - skipping for this layer: %s.",
             fse_reason,
         )
     return is_fused_shared_expert_enabled

vllm/model_executor/layers/quantization/utils/config_utils.py

@@ -138,15 +138,27 @@ def is_shared_expert_quant_fse_compatible(
         )
 
     if isinstance(quant_config, QuarkConfig):
+        from vllm.model_executor.layers.quantization.quark.utils import (
+            should_ignore_layer,
+        )
+
         # TODO: layer_type_quant_config is not taken into account here.
         assert "exclude" in quant_config.quant_config
         assert "global_quant_config" in quant_config.quant_config
 
-        is_compatible = not any(
-            "shared_expert" in str(entry)
-            for entry in quant_config.quant_config["exclude"]
+        exclude = quant_config.quant_config["exclude"]
+
+        # should_ignore_layer raises a rightful ValueError in case different shards
+        # have a different ignore policy, as unsupported.
+        is_excluded = any(
+            should_ignore_layer(
+                f"{shared_expert_prefix}.{projection_name}",
+                ignore=exclude,
+                fused_mapping=quant_config.packed_modules_mapping,
+            )
+            for projection_name in projection_names
         )
-        if not is_compatible:
+        if is_excluded:
             return False, f"Quark excludes shared experts at {shared_expert_prefix}"
 
         global_quant_config = quant_config.quant_config["global_quant_config"]