From 27ffbfde8decd340fa0144aea06061374c53e456 Mon Sep 17 00:00:00 2001 From: Colin Z <59755453+ColinZ22@users.noreply.github.com> Date: Wed, 22 Jul 2026 17:34:17 -0700 Subject: [PATCH] Fused Shared Expert Support for AMD Quark DeepSeek-V4 Model Checkpoints (#48044) Signed-off-by: Colin Zeng Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- vllm/models/deepseek_v4/amd/model.py | 132 ++++++++++++++++++++++-- vllm/models/deepseek_v4/amd/mtp.py | 17 +++ vllm/models/deepseek_v4/quant_config.py | 39 ++++++- 3 files changed, 175 insertions(+), 13 deletions(-) diff --git a/vllm/models/deepseek_v4/amd/model.py b/vllm/models/deepseek_v4/amd/model.py index 9093bd41289..b382f7fd6f9 100644 --- a/vllm/models/deepseek_v4/amd/model.py +++ b/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"(? 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") diff --git a/vllm/models/deepseek_v4/amd/mtp.py b/vllm/models/deepseek_v4/amd/mtp.py index a12b401f0d4..eb94a86173d 100644 --- a/vllm/models/deepseek_v4/amd/mtp.py +++ b/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 diff --git a/vllm/models/deepseek_v4/quant_config.py b/vllm/models/deepseek_v4/quant_config.py index 721a9138914..89cf695baf0 100644 --- a/vllm/models/deepseek_v4/quant_config.py +++ b/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(