Fused Shared Expert Support for AMD Quark DeepSeek-V4 Model Checkpoints (#48044)

Signed-off-by: Colin Zeng <[email protected]>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

vllm/models/deepseek_v4/amd/model.py

@@ -8,7 +8,8 @@ import regex as re
 import torch
 import torch.nn as nn
 
-from vllm.config import VllmConfig
+import vllm.envs as envs
+from vllm.config import VllmConfig, get_current_vllm_config
 from vllm.distributed import (
     get_pp_group,
     get_tensor_model_parallel_rank,
@@ -110,6 +111,51 @@ class DeepseekV4MLP(nn.Module):
         return x
 
 
+def _shared_experts_are_fp4(config, layer_idx: int | None = None) -> bool:
+    """Whether the shared experts are MXFP4 and thus fusable.
+
+    ``layer_idx=None`` resolves the model-wide default (global scheme), used by
+    the main-model weight loader / mapper callers that operate per-model.
+    """
+    quant_cfg = getattr(config, "quantization_config", None)
+    if quant_cfg is None:
+        return False
+    if layer_idx is None:
+        base = None
+    elif layer_idx >= config.num_hidden_layers:
+        base = f"mtp.{layer_idx - config.num_hidden_layers}.ffn.shared_experts"
+    else:
+        base = f"layers.{layer_idx}.ffn.shared_experts"
+    if base and any(e.startswith(base) for e in (quant_cfg.get("exclude") or [])):
+        return False
+    entry = (
+        (quant_cfg.get("layer_quant_config") or {}).get(f"{base}.w1") if base else None
+    )
+    if entry is None:
+        entry = quant_cfg.get("global_quant_config")
+    return ((entry or {}).get("weight") or {}).get("dtype") == "fp4"
+
+
+def _fuse_shared_experts_enabled(config, prefix: str = "") -> bool:
+    """Whether to fuse the shared expert into the routed MXFP4 grouped GEMM.
+
+    Fusion fuses the shared expert into the routed experts' MXFP4 grouped GEMM,
+    so it only applies where the shared expert is the same precision as the
+    routed experts. Some layers may carry a shared expert in a different quantization
+    than the routed experts; when so, it runs as its own linear and must not be fused.
+    """
+    if not (
+        current_platform.is_rocm()
+        and getattr(config, "n_shared_experts", None)
+        and envs.VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS
+        and not get_current_vllm_config().parallel_config.enable_expert_parallel
+    ):
+        return False
+    return _shared_experts_are_fp4(
+        config, extract_layer_index(prefix) if prefix else None
+    )
+
+
 class DeepseekV4MoE(nn.Module):
     def __init__(
         self,
@@ -164,7 +210,11 @@ class DeepseekV4MoE(nn.Module):
                 requires_grad=False,
             )
 
-        if config.n_shared_experts is None:
+        self.n_shared_experts = config.n_shared_experts
+
+        self.fuse_shared_experts = _fuse_shared_experts_enabled(config, prefix)
+
+        if config.n_shared_experts is None or self.fuse_shared_experts:
             self.shared_experts = None
         else:
             intermediate_size = config.moe_intermediate_size * config.n_shared_experts
@@ -188,6 +238,9 @@ class DeepseekV4MoE(nn.Module):
 
         self.experts = FusedMoE(
             shared_experts=self.shared_experts,
+            n_shared_experts=(
+                config.n_shared_experts if self.fuse_shared_experts else None
+            ),
             gate=self.gate,
             num_experts=config.n_routed_experts,
             top_k=config.num_experts_per_tok,
@@ -667,7 +720,38 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
         # Pre-compute expert mapping ONCE.
         expert_mapping = self.get_expert_mapping()
 
+        # Use each MoE's own per-layer fusion decision (computed with its prefix
+        # at init) as the single source of truth, so the redirect below cannot
+        # diverge from how the module was built if per-layer quantization ever
+        # mixes fused and non-fused layers.
+        fuse_by_layer = {
+            extract_layer_index(mod_name): mod.fuse_shared_experts
+            for mod_name, mod in self.named_modules()
+            if isinstance(mod, DeepseekV4MoE)
+        }
+        n_routed = self.config.n_routed_experts
+        # The redirect below maps the single shared-expert tensor group to one
+        # appended slot; multiple shared experts would need per-expert slicing
+        # (see deepseek_v2.py). DeepSeek-V4 has n_shared_experts == 1.
+        if any(fuse_by_layer.values()) and self.config.n_shared_experts != 1:
+            raise NotImplementedError(
+                "deepseek-v4 fused shared-expert loading supports only "
+                f"n_shared_experts == 1, got {self.config.n_shared_experts}"
+            )
+
         for name, loaded_weight in weights:
+            # Shared-expert fusion: redirect ``.ffn.shared_experts.w{1,2,3}``
+            # into appended routed-expert slot ``.ffn.experts.{n_routed}``
+            # so the MXFP4-quantized shared expert loads through the routed
+            # expert loader (grouped GEMM). Single shared expert only.
+            if ".ffn.shared_experts.w" in name and fuse_by_layer.get(
+                extract_layer_index(name), False
+            ):
+                name = name.replace(
+                    ".ffn.shared_experts.w",
+                    f".ffn.experts.{n_routed}.w",
+                )
+
             for param_name, weight_name, shard_id in stacked_params_mapping:
                 # Skip non-stacked layers and experts (experts handled below).
                 if ".experts." in name:
@@ -745,24 +829,41 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
     def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
         # Params for weights, fp8 weight scales, fp8 activation scales
         # (param_name, weight_name, expert_id, shard_id)
+        # When fusing shared experts, include the appended slots
+        # (ids n_routed_experts .. n_routed_experts + n_shared - 1) so the
+        # redirected shared-expert weights route through the expert loader.
+        n_shared = getattr(self.config, "n_shared_experts", 0) or 0
+        num_experts = self.config.n_routed_experts + (
+            n_shared if _fuse_shared_experts_enabled(self.config) else 0
+        )
         return fused_moe_make_expert_params_mapping(
             self,
             ckpt_gate_proj_name="w1",
             ckpt_down_proj_name="w2",
             ckpt_up_proj_name="w3",
-            num_experts=self.config.n_routed_experts,
+            num_experts=num_experts,
         )
 
 
-def _make_deepseek_v4_weights_mapper(expert_dtype: str) -> WeightsMapper:
+def _make_deepseek_v4_weights_mapper(
+    expert_dtype: str, fuse_shared_experts: bool = False
+) -> WeightsMapper:
     if expert_dtype == "fp4":
         # MXFP4 experts use Mxfp4MoEMethod, which registers scales as
         # ``w{1,2,3}_weight_scale`` (no _inv suffix). FP8 linear and
-        # shared experts use Fp8LinearMethod's block scales, which
-        # register as ``weight_scale_inv``.
+        # (non-fused) shared experts use Fp8LinearMethod's block scales,
+        # which register as ``weight_scale_inv``.
+        #
+        #  - DeepSeek native ``.scale``: expert scales -> ``.weight_scale``,
+        #    everything else -> ``.weight_scale_inv``.
+        #  - AMD-Quark ``.weight_scale``: linear/attn scales ->
+        #    ``.weight_scale_inv``. Expert and shared-expert
+        #    ``w{1,2,3}.weight_scale`` are left untouched (consumed as-is by
+        #    the MXFP4 expert loader, which produces ``w{13,2}_weight_scale``);
         scale_regex = {
             re.compile(r"(\.experts\.\d+\.w[123])\.scale$"): r"\1.weight_scale",
             re.compile(r"\.scale$"): ".weight_scale_inv",
+            re.compile(r"(?<!\.w[123])\.weight_scale$"): ".weight_scale_inv",
         }
     else:
         # FP8 experts use Fp8MoEMethod (block_quant=True), which registers
@@ -771,6 +872,14 @@ def _make_deepseek_v4_weights_mapper(expert_dtype: str) -> WeightsMapper:
         scale_regex = {
             re.compile(r"\.scale$"): ".weight_scale_inv",
         }
+    # When shared experts are fused into the routed MXFP4 grouped GEMM, the
+    # shared_experts tensors are redirected to routed expert slots ; leave
+    # their names untouched here.
+    substr_map = (
+        {}
+        if fuse_shared_experts
+        else {".shared_experts.w2": ".shared_experts.down_proj"}
+    )
     return WeightsMapper(
         orig_to_new_prefix={
             "layers.": "model.layers.",
@@ -785,9 +894,7 @@ def _make_deepseek_v4_weights_mapper(expert_dtype: str) -> WeightsMapper:
             "embed.weight": "embed_tokens.weight",
             ".ffn.gate.bias": ".ffn.gate.e_score_correction_bias",
         },
-        orig_to_new_substr={
-            ".shared_experts.w2": ".shared_experts.down_proj",
-        },
+        orig_to_new_substr=substr_map,
     )
 
 
@@ -804,8 +911,11 @@ class DeepseekV4ForCausalLM(nn.Module, SupportsPP, SupportsEagle3):
         config = vllm_config.model_config.hf_config
         self.config = config
         expert_dtype = getattr(config, "expert_dtype", "fp4")
-        if expert_dtype != "fp4":
-            self.hf_to_vllm_mapper = _make_deepseek_v4_weights_mapper(expert_dtype)
+        fuse_shared_experts = _fuse_shared_experts_enabled(config)
+        if expert_dtype != "fp4" or fuse_shared_experts:
+            self.hf_to_vllm_mapper = _make_deepseek_v4_weights_mapper(
+                expert_dtype, fuse_shared_experts=fuse_shared_experts
+            )
 
         self.model = self.model_cls(
             vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")

vllm/models/deepseek_v4/amd/mtp.py

@@ -334,6 +334,21 @@ class DeepSeekV4MTP(nn.Module):
         params_dict = dict(self.named_parameters())
         loaded_params: set[str] = set()
 
+        def _resolve_scale_name(name: str) -> str:
+            # Quark checkpoints name FP8 block scales ``.weight_scale``,
+            # but block-FP8 layers register them as ``.weight_scale_inv``
+            # while MXFP4 experts register ``.weight_scale``. Auto-detect:
+            # rename to ``_inv`` only when that variant exists and the plain
+            # one does not.
+            if name.endswith(".weight_scale") and name not in params_dict:
+                inv = name.removesuffix(".weight_scale") + ".weight_scale_inv"
+                if inv in params_dict:
+                    return inv
+            # Otherwise leave the name unchanged: either it already matches a
+            # param, or it is genuinely unknown and should surface the normal
+            # KeyError downstream rather than be silently rewritten.
+            return name
+
         # TP for attention
         tp_size = get_tensor_model_parallel_world_size()
         tp_rank = get_tensor_model_parallel_rank()
@@ -393,6 +408,7 @@ class DeepSeekV4MTP(nn.Module):
                 if weight_name not in name:
                     continue
                 name = name.replace(weight_name, param_name)
+                name = _resolve_scale_name(name)
 
                 param = params_dict[name]
                 weight_loader = param.weight_loader
@@ -447,6 +463,7 @@ class DeepSeekV4MTP(nn.Module):
                         )
                     if name.endswith(".ffn.gate.bias"):
                         name = name.replace(".bias", ".e_score_correction_bias")
+                    name = _resolve_scale_name(name)
                     param = params_dict[name]
                     weight_loader = getattr(
                         param, "weight_loader", default_weight_loader

vllm/models/deepseek_v4/quant_config.py

@@ -4,7 +4,7 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, cast
 
 from vllm.config import get_current_vllm_config
 from vllm.model_executor.layers.fused_moe import (
@@ -117,13 +117,29 @@ class DeepseekV4FP8Config(Fp8Config):
     def get_name(cls) -> QuantizationMethods:
         return "deepseek_v4_fp8"
 
+    @staticmethod
+    def _is_quark_mxfp4_ocp(hf_quant_cfg: dict) -> bool:
+        """True for AMD-Quark exports whose global scheme is MXFP4."""
+        weight = (hf_quant_cfg.get("global_quant_config") or {}).get("weight") or {}
+        return (
+            weight.get("dtype") == "fp4"
+            and weight.get("qscheme") == "per_group"
+            and weight.get("group_size") == 32
+        )
+
     @classmethod
     def override_quantization_method(
         cls, hf_quant_cfg, user_quant, hf_config=None
     ) -> QuantizationMethods | None:
         if not (
             isinstance(hf_quant_cfg, dict)
-            and hf_quant_cfg.get("quant_method") in ("fp8", "deepseek_v4_fp8")
+            and (
+                hf_quant_cfg.get("quant_method") in ("fp8", "deepseek_v4_fp8")
+                or (
+                    hf_quant_cfg.get("quant_method") == "quark"
+                    and cls._is_quark_mxfp4_ocp(hf_quant_cfg)
+                )
+            )
         ):
             return None
         model_type = getattr(hf_config, "model_type", None)
@@ -131,6 +147,25 @@ class DeepseekV4FP8Config(Fp8Config):
             return "deepseek_v4_fp8"
         return None
 
+    @classmethod
+    def from_config(cls, config: dict) -> DeepseekV4FP8Config:
+        # Reroute AMD-Quark fused shared expert MXFP4 checkpoints onto the fp8
+        # path: the runtime layout matches the DeepSeek-native fp8 checkpoint,
+        # so translate the schema into format Fp8Config.from_config expects.
+        if config.get("quant_method") == "quark":
+            quark_exclude = config.get("exclude") or []
+            config = {
+                "quant_method": "fp8",
+                "activation_scheme": "dynamic",
+                "fmt": "e4m3",
+                "scale_fmt": "ue8m0",
+                "weight_block_size": [128, 128],
+                "ignored_layers": [
+                    name for name in quark_exclude if isinstance(name, str)
+                ],
+            }
+        return cast("DeepseekV4FP8Config", super().from_config(config))
+
     def get_quant_method(self, layer, prefix):
         if isinstance(layer, RoutedExperts):
             if is_layer_skipped(