[Quantization] Enable per-layer fused shared expert selection

vllm/model_executor/layers/fused_moe/config.py

@@ -1338,7 +1338,7 @@ class FusedMoEConfig:
     # Use the allocated width as the checkpoint TP stride, including for scales.
     tp_shard_with_padding: bool = False
     rocm_aiter_fmoe_enabled: bool = False
-    aiter_fmoe_shared_expert_enabled: bool = False
+    is_fused_shared_expert_requested: bool = False
     # Whether to force MXFP4 (a4w4) MoE activations for DeepSeek V4.1 on
     # ROCm/AITER. Opt-in via VLLM_ROCM_USE_AITER_MOE_A4W4_DSV4=1; rejected
     # for any other model type. Resolved here, not in the forward path,
@@ -1371,7 +1371,7 @@ def __post_init__(self):
 
         if self.is_act_and_mul:
             self.rocm_aiter_fmoe_enabled = rocm_aiter_ops.is_fused_moe_enabled()
-            self.aiter_fmoe_shared_expert_enabled = (
+            self.is_fused_shared_expert_requested = bool(
                 rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()
             )
 
@@ -1402,7 +1402,7 @@ def __post_init__(self):
             assert self.rocm_aiter_fmoe_enabled, (
                 "Mori needs to be used with aiter fused_moe for now."
             )
-            assert not self.aiter_fmoe_shared_expert_enabled, (
+            assert not self.is_fused_shared_expert_requested, (
                 "Mori does not support fusion shared expert now. "
                 "Turn it off by setting VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS=0"
             )

vllm/model_executor/layers/fused_moe/routed_experts.py

@@ -725,6 +725,16 @@ def weight_loader(
 
         quant_method_name = base_quant_method.__class__.__name__
         global_expert_id = expert_id
+        num_global_experts = (
+            self.moe_config.num_experts
+            + self.expert_map_manager.num_fused_shared_experts
+        )
+        # `num_experts` counts logical and redundant experts; fused shared
+        # slots are appended after them.
+        assert global_expert_id < num_global_experts, (
+            f"Expert id {global_expert_id} is out of range for "
+            f"{self.layer_name}; expected < {num_global_experts}"
+        )
         expert_id = self._map_global_expert_id_to_local_expert_id(global_expert_id)
 
         use_global_sf = (

vllm/model_executor/layers/fused_moe/utils.py

@@ -1,13 +1,12 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import functools
-from collections.abc import Iterable
 from math import prod
 from typing import TYPE_CHECKING
 
+import regex as re
 import torch
 import torch.nn.functional as F
-from torch import nn
 
 import vllm.envs as envs
 from vllm import _custom_ops as ops
@@ -40,7 +39,6 @@
 from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
     per_tensor_dequantize,
 )
-from vllm.model_executor.models.utils import PPMissingLayer
 from vllm.platforms import current_platform
 from vllm.triton_utils import tl, triton
 from vllm.utils.math_utils import cdiv
@@ -52,11 +50,23 @@
 logger = init_logger(__name__)
 
 
+# Matches any layer-qualified module path (e.g.
+# ``model.layers.5.mlp.shared_expert`` or
+# ``model.language_model.model.layers.5.mlp.shared_expert``) so skip reasons
+# are reported once per cause instead of once per layer.
+_LAYER_MODULE_PATH_RE = re.compile(r"(?:[\w]+\.)*(?:layers|mtp)\.\d+\.[\w.]*")
+
+
+def _normalize_skip_reason(reason: str) -> str:
+    return _LAYER_MODULE_PATH_RE.sub("<shared_expert>", reason)
+
+
 def resolve_layer_fused_shared_expert(
     quant_config: "QuantizationConfig | None",
     prefix: str,
     shared_expert_name: str = "shared_experts",
     moe_backend: str | None = None,
+    disable_under_ep: bool = False,
 ) -> bool:
     """Resolve whether AITER fused shared-expert execution is enabled.
 
@@ -65,6 +75,9 @@ def resolve_layer_fused_shared_expert(
         prefix: MoE module prefix.
         shared_expert_name: Shared-expert module name under ``prefix``.
         moe_backend: MoE backend that runs the fused shared expert.
+        disable_under_ep: Whether this model family does not support FSE under
+            expert parallelism. Kept as an explicit call-site policy instead of
+            branching on model type here.
 
     Returns:
         Whether AITER fused shared experts are enabled.
@@ -76,16 +89,31 @@ def resolve_layer_fused_shared_expert(
     # NOTE: is_fusion_moe_shared_experts_enabled is decorated with @if_aiter_supported
     # that returns None if AITER is not available.
     fse_requested = bool(rocm_aiter_ops.is_fusion_moe_shared_experts_enabled())
-    fse_compatible, fse_reason = (
-        is_shared_expert_quant_fse_compatible(
+
+    # TODO: should it be disabled for all models?
+    skip_reason: str | None = None
+    if fse_requested and disable_under_ep:
+        from vllm.config import get_current_vllm_config_or_none
+
+        vllm_config = get_current_vllm_config_or_none()
+        parallel_config = (
+            vllm_config.parallel_config if vllm_config is not None else None
+        )
+        if parallel_config is not None and parallel_config.enable_expert_parallel:
+            fse_requested = False
+            skip_reason = "expert parallelism is enabled"
+
+    if fse_requested:
+        fse_compatible, fse_reason = is_shared_expert_quant_fse_compatible(
             quant_config,
             f"{prefix}.experts",
             f"{prefix}.{shared_expert_name}",
             moe_backend=moe_backend,
         )
-        if fse_requested
-        else (True, None)
-    )
+        if not fse_compatible:
+            skip_reason = fse_reason or "unknown reason"
+    else:
+        fse_compatible = True
     is_fused_shared_expert_enabled = fse_requested and fse_compatible
 
     # In online shared_expert quantization is used, register it to
@@ -121,52 +149,18 @@ def resolve_layer_fused_shared_expert(
                     quant_key_str,
                     target_pattern,
                 )
-    if fse_requested and not is_fused_shared_expert_enabled:
-        logger.warning(
+    if skip_reason is not None:
+        # Keep the full per-layer reason at debug level; the warning below is
+        # deduplicated by cause so a long model logs it once.
+        logger.debug("Skipping shared-expert fusion at %s: %s", prefix, skip_reason)
+        logger.warning_once(
             "VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS is enabled but "
-            "cannot be enabled - skipping for this layer: %s.",
-            fse_reason,
+            "cannot be enabled - skipping shared-expert fusion: %s.",
+            _normalize_skip_reason(skip_reason),
         )
     return is_fused_shared_expert_enabled
 
 
-def is_model_fused_shared_expert_compatible(
-    layers: nn.ModuleList | Iterable[nn.Module],
-    moe_cls: type[nn.Module],
-    moe_name: str,
-) -> bool:
-    """Resolve one fused-shared-expert state for a model's MoE layers."""
-
-    def get_moe_layer(layer: nn.Module) -> nn.Module | None:
-        for name in moe_name.split("."):
-            layer = getattr(layer, name, None)
-            if layer is None:
-                return None
-        return layer
-
-    moe_layers = (
-        moe_layer
-        for layer in layers
-        if not isinstance(layer, PPMissingLayer)
-        and (moe_layer := get_moe_layer(layer)) is not None
-        and isinstance(moe_layer, moe_cls)
-    )
-
-    enabled = [
-        getattr(layer, "is_fused_shared_expert_enabled", False) for layer in moe_layers
-    ]
-    enabled_count = sum(enabled)
-    disabled_count = len(enabled) - enabled_count
-    if enabled_count > 0 and disabled_count > 0:
-        raise NotImplementedError(
-            "Fused shared experts must be enabled for all MoE layers; found "
-            f"{enabled_count} enabled and {disabled_count} disabled layers. "
-            "Per-layer fused shared experts is not yet supported. Please open "
-            "an issue."
-        )
-    return enabled_count > 0 and disabled_count == 0
-
-
 @triton.jit
 def _count_expert_num_tokens(
     topk_ids_ptr,

vllm/model_executor/models/AXK1.py

@@ -49,7 +49,6 @@
     fused_moe_make_expert_params_mapping,
 )
 from vllm.model_executor.layers.fused_moe.utils import (
-    is_model_fused_shared_expert_compatible,
     resolve_layer_fused_shared_expert,
 )
 from vllm.model_executor.layers.layernorm import RMSNorm
@@ -84,6 +83,8 @@
 from .utils import (
     AutoWeightsLoader,
     PPMissingLayer,
+    extract_layer_index,
+    get_fused_shared_expert_layers,
     get_spec_layer_idx_from_weight_name,
     is_pp_missing_parameter,
     make_empty_intermediate_tensors_factory,
@@ -742,10 +743,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
             prefix=f"{prefix}.layers",
         )
 
-        self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible(
-            self.layers,
+        self.fused_shared_expert_layers = get_fused_shared_expert_layers(
+            self,
             AXK1MoE,
-            "mlp",
+            enable_eplb=vllm_config.parallel_config.enable_eplb,
         )
 
         if get_pp_group().is_last_rank:
@@ -839,10 +840,11 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
             ckpt_gate_proj_name="gate_proj",
             ckpt_down_proj_name="down_proj",
             ckpt_up_proj_name="up_proj",
+            # NOTE: global despite per-layer FSE; non-fused layers skip the extras.
             num_experts=self.config.n_routed_experts
             + (
                 (self.config.n_shared_experts or 0)
-                if self.is_fused_shared_expert_enabled
+                if self.fused_shared_expert_layers
                 else 0
             ),
             num_redundant_experts=self.num_redundant_experts,
@@ -858,8 +860,11 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
             if spec_layer is not None:
                 continue  # skip spec decode layers for main model
 
+            # Only parse the layer index when FSE is enabled.
             is_fusion_moe_shared_experts_layer = (
-                self.is_fused_shared_expert_enabled and ("mlp.shared_experts" in name)
+                "mlp.shared_experts" in name
+                and bool(rocm_aiter_ops.is_fusion_moe_shared_experts_enabled())
+                and extract_layer_index(name) in self.fused_shared_expert_layers
             )
 
             for param_name, weight_name, shard_id in stacked_params_mapping:

vllm/model_executor/models/alice_ai.py

@@ -20,7 +20,6 @@
 )
 from vllm.model_executor.layers.fused_moe import FusedMoEFactory, GateLinear
 from vllm.model_executor.layers.fused_moe.utils import (
-    is_model_fused_shared_expert_compatible,
     resolve_layer_fused_shared_expert,
 )
 from vllm.model_executor.layers.layernorm import GemmaRMSNorm
@@ -57,6 +56,7 @@
     AutoWeightsLoader,
     WeightsMapper,
     extract_layer_index,
+    get_fused_shared_expert_layers,
     make_layers,
     maybe_fuse_shared_experts,
     maybe_prefix,
@@ -500,8 +500,10 @@ def get_layer(prefix: str) -> AliceAIDecoderBlock:
             layer_fn=get_layer,
             prefix=f"{prefix}.layers",
         )
-        self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible(
-            layers=self.layers, moe_cls=AliceAISparseMoeBlock, moe_name="mlp"
+        self.fused_shared_expert_layers = get_fused_shared_expert_layers(
+            self,
+            AliceAISparseMoeBlock,
+            enable_eplb=vllm_config.parallel_config.enable_eplb,
         )
         self.norm = GemmaRMSNorm(
             hidden_size=config.hidden_size, eps=config.rms_norm_eps
@@ -567,7 +569,7 @@ def forward(
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         weights = maybe_fuse_shared_experts(
             weights=weights,
-            enabled=self.is_fused_shared_expert_enabled,
+            fused_shared_expert_layers=self.fused_shared_expert_layers,
             n_routed_experts=self.config.num_experts,
             n_shared_experts=1,
             ckpt_prefix="mlp.shared_expert",

vllm/model_executor/models/deepseek_mtp.py

@@ -7,15 +7,13 @@
 import torch.nn as nn
 from transformers import PreTrainedConfig
 
+from vllm._aiter_ops import rocm_aiter_ops
 from vllm.compilation.decorators import support_torch_compile
 from vllm.config import VllmConfig
 from vllm.distributed import tensor_model_parallel_all_gather
 from vllm.model_executor.layers.fused_moe import (
     fused_moe_make_expert_params_mapping,
 )
-from vllm.model_executor.layers.fused_moe.utils import (
-    is_model_fused_shared_expert_compatible,
-)
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
 from vllm.model_executor.layers.quantization import QuantizationConfig
@@ -41,6 +39,8 @@
 )
 from .interfaces import SupportsPP
 from .utils import (
+    extract_layer_index,
+    get_fused_shared_expert_layers,
     get_pp_missing_layer_names,
     get_spec_layer_idx_from_weight_name,
     make_empty_intermediate_tensors_factory,
@@ -242,6 +242,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         super().__init__()
         self.config = vllm_config.model_config.hf_config
         self.quant_config = vllm_config.quant_config
+        self.vllm_config = vllm_config
         self.model = DeepSeekMultiTokenPredictor(
             vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
         )
@@ -267,10 +268,10 @@ def set_moe_parameters(self):
                 self.moe_mlp_layers.append(layer.mlp)
                 self.moe_layers.append(layer.mlp.experts)
         self.extract_moe_parameters(example_moe)
-        self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible(
-            self.model.layers.values(),
+        self.fused_shared_expert_layers = get_fused_shared_expert_layers(
+            self,
             DeepseekV2MoE,
-            "mtp_block.mlp",
+            enable_eplb=self.vllm_config.parallel_config.enable_eplb,
         )
 
     def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
@@ -316,17 +317,15 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         ]
         stacked_params_mapping.extend(indexer_fused_mapping)
 
+        fused_shared_expert_layers = self.fused_shared_expert_layers
         expert_params_mapping = fused_moe_make_expert_params_mapping(
             self,
             ckpt_gate_proj_name="gate_proj",
             ckpt_down_proj_name="down_proj",
             ckpt_up_proj_name="up_proj",
+            # NOTE: global despite per-layer FSE; non-fused layers skip the extras.
             num_experts=self.config.n_routed_experts
-            + (
-                self.config.n_shared_experts
-                if self.is_fused_shared_expert_enabled
-                else 0
-            ),
+            + (self.config.n_shared_experts if fused_shared_expert_layers else 0),
         )
 
         pp_missing_layer_names = get_pp_missing_layer_names(self)
@@ -348,8 +347,11 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
                     weight_loader(param, loaded_weight)
                     loaded_params.add(name)
                 continue
+            # Only parse the layer index when FSE is enabled.
             is_fusion_moe_shared_experts_layer = (
-                self.is_fused_shared_expert_enabled and ("mlp.shared_experts" in name)
+                "mlp.shared_experts" in name
+                and bool(rocm_aiter_ops.is_fusion_moe_shared_experts_enabled())
+                and extract_layer_index(name) in fused_shared_expert_layers
             )
             name = self._rewrite_spec_layer_name(spec_layer, name)