@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 27ffbfde8d

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

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

Colin Z committed Jul 22, 2026 at 17:34 UTC 27ffbfde8decd340fa0144aea06061374c53e456

3 files changed +175 -13

vllm/models/deepseek_v4/amd/model.py

+121 -11

@@ -8,7 +8,8 @@ import regex as re

8 import torch

9 import torch.nn as nn

11 -from vllm.config import VllmConfig

11 +import vllm.envs as envs

12 +from vllm.config import VllmConfig, get_current_vllm_config

13 from vllm.distributed import (

14 get_pp_group,

15 get_tensor_model_parallel_rank,

@@ -110,6 +111,51 @@ class DeepseekV4MLP(nn.Module):

111 return x

112

113

114 +def _shared_experts_are_fp4(config, layer_idx: int | None = None) -> bool:

115 + """Whether the shared experts are MXFP4 and thus fusable.

116 +

117 + ``layer_idx=None`` resolves the model-wide default (global scheme), used by

118 + the main-model weight loader / mapper callers that operate per-model.

119 + """

120 + quant_cfg = getattr(config, "quantization_config", None)

121 + if quant_cfg is None:

122 + return False

123 + if layer_idx is None:

124 + base = None

125 + elif layer_idx >= config.num_hidden_layers:

126 + base = f"mtp.{layer_idx - config.num_hidden_layers}.ffn.shared_experts"

127 + else:

128 + base = f"layers.{layer_idx}.ffn.shared_experts"

129 + if base and any(e.startswith(base) for e in (quant_cfg.get("exclude") or [])):

130 + return False

131 + entry = (

132 + (quant_cfg.get("layer_quant_config") or {}).get(f"{base}.w1") if base else None

133 + )

134 + if entry is None:

135 + entry = quant_cfg.get("global_quant_config")

136 + return ((entry or {}).get("weight") or {}).get("dtype") == "fp4"

137 +

138 +

139 +def _fuse_shared_experts_enabled(config, prefix: str = "") -> bool:

140 + """Whether to fuse the shared expert into the routed MXFP4 grouped GEMM.

141 +

142 + Fusion fuses the shared expert into the routed experts' MXFP4 grouped GEMM,

143 + so it only applies where the shared expert is the same precision as the

144 + routed experts. Some layers may carry a shared expert in a different quantization

145 + than the routed experts; when so, it runs as its own linear and must not be fused.

146 + """

147 + if not (

148 + current_platform.is_rocm()

149 + and getattr(config, "n_shared_experts", None)

150 + and envs.VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS

151 + and not get_current_vllm_config().parallel_config.enable_expert_parallel

152 + ):

153 + return False

154 + return _shared_experts_are_fp4(

155 + config, extract_layer_index(prefix) if prefix else None

156 + )

157 +

158 +

159 class DeepseekV4MoE(nn.Module):

160 def __init__(

161 self,

@@ -164,7 +210,11 @@ class DeepseekV4MoE(nn.Module):

210 requires_grad=False,

211 )

212

167 - if config.n_shared_experts is None:

213 + self.n_shared_experts = config.n_shared_experts

214 +

215 + self.fuse_shared_experts = _fuse_shared_experts_enabled(config, prefix)

216 +

217 + if config.n_shared_experts is None or self.fuse_shared_experts:

218 self.shared_experts = None

219 else:

220 intermediate_size = config.moe_intermediate_size * config.n_shared_experts

@@ -188,6 +238,9 @@ class DeepseekV4MoE(nn.Module):

238

239 self.experts = FusedMoE(

240 shared_experts=self.shared_experts,

241 + n_shared_experts=(

242 + config.n_shared_experts if self.fuse_shared_experts else None

243 + ),

244 gate=self.gate,

245 num_experts=config.n_routed_experts,

246 top_k=config.num_experts_per_tok,

@@ -667,7 +720,38 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):

720 # Pre-compute expert mapping ONCE.

721 expert_mapping = self.get_expert_mapping()

722

723 + # Use each MoE's own per-layer fusion decision (computed with its prefix

724 + # at init) as the single source of truth, so the redirect below cannot

725 + # diverge from how the module was built if per-layer quantization ever

726 + # mixes fused and non-fused layers.

727 + fuse_by_layer = {

728 + extract_layer_index(mod_name): mod.fuse_shared_experts

729 + for mod_name, mod in self.named_modules()

730 + if isinstance(mod, DeepseekV4MoE)

731 + }

732 + n_routed = self.config.n_routed_experts

733 + # The redirect below maps the single shared-expert tensor group to one

734 + # appended slot; multiple shared experts would need per-expert slicing

735 + # (see deepseek_v2.py). DeepSeek-V4 has n_shared_experts == 1.

736 + if any(fuse_by_layer.values()) and self.config.n_shared_experts != 1:

737 + raise NotImplementedError(

738 + "deepseek-v4 fused shared-expert loading supports only "

739 + f"n_shared_experts == 1, got {self.config.n_shared_experts}"

740 + )

741 +

742 for name, loaded_weight in weights:

743 + # Shared-expert fusion: redirect ``.ffn.shared_experts.w{1,2,3}``

744 + # into appended routed-expert slot ``.ffn.experts.{n_routed}``

745 + # so the MXFP4-quantized shared expert loads through the routed

746 + # expert loader (grouped GEMM). Single shared expert only.

747 + if ".ffn.shared_experts.w" in name and fuse_by_layer.get(

748 + extract_layer_index(name), False

749 + ):

750 + name = name.replace(

751 + ".ffn.shared_experts.w",

752 + f".ffn.experts.{n_routed}.w",

753 + )

754 +

755 for param_name, weight_name, shard_id in stacked_params_mapping:

756 # Skip non-stacked layers and experts (experts handled below).

757 if ".experts." in name:

@@ -745,24 +829,41 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):

829 def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:

830 # Params for weights, fp8 weight scales, fp8 activation scales

831 # (param_name, weight_name, expert_id, shard_id)

832 + # When fusing shared experts, include the appended slots

833 + # (ids n_routed_experts .. n_routed_experts + n_shared - 1) so the

834 + # redirected shared-expert weights route through the expert loader.

835 + n_shared = getattr(self.config, "n_shared_experts", 0) or 0

836 + num_experts = self.config.n_routed_experts + (

837 + n_shared if _fuse_shared_experts_enabled(self.config) else 0

838 + )

839 return fused_moe_make_expert_params_mapping(

840 self,

841 ckpt_gate_proj_name="w1",

842 ckpt_down_proj_name="w2",

843 ckpt_up_proj_name="w3",

753 - num_experts=self.config.n_routed_experts,

844 + num_experts=num_experts,

845 )

846

847

757 -def _make_deepseek_v4_weights_mapper(expert_dtype: str) -> WeightsMapper:

848 +def _make_deepseek_v4_weights_mapper(

849 + expert_dtype: str, fuse_shared_experts: bool = False

850 +) -> WeightsMapper:

851 if expert_dtype == "fp4":

852 # MXFP4 experts use Mxfp4MoEMethod, which registers scales as

853 # ``w{1,2,3}_weight_scale`` (no _inv suffix). FP8 linear and

761 - # shared experts use Fp8LinearMethod's block scales, which

762 - # register as ``weight_scale_inv``.

854 + # (non-fused) shared experts use Fp8LinearMethod's block scales,

855 + # which register as ``weight_scale_inv``.

856 + #

857 + # - DeepSeek native ``.scale``: expert scales -> ``.weight_scale``,

858 + # everything else -> ``.weight_scale_inv``.

859 + # - AMD-Quark ``.weight_scale``: linear/attn scales ->

860 + # ``.weight_scale_inv``. Expert and shared-expert

861 + # ``w{1,2,3}.weight_scale`` are left untouched (consumed as-is by

862 + # the MXFP4 expert loader, which produces ``w{13,2}_weight_scale``);

863 scale_regex = {

864 re.compile(r"(\.experts\.\d+\.w[123])\.scale$"): r"\1.weight_scale",

865 re.compile(r"\.scale$"): ".weight_scale_inv",

866 + re.compile(r"(?<!\.w[123])\.weight_scale$"): ".weight_scale_inv",

867 }

868 else:

869 # FP8 experts use Fp8MoEMethod (block_quant=True), which registers

@@ -771,6 +872,14 @@ def _make_deepseek_v4_weights_mapper(expert_dtype: str) -> WeightsMapper:

872 scale_regex = {

873 re.compile(r"\.scale$"): ".weight_scale_inv",

874 }

875 + # When shared experts are fused into the routed MXFP4 grouped GEMM, the

876 + # shared_experts tensors are redirected to routed expert slots ; leave

877 + # their names untouched here.

878 + substr_map = (

879 + {}

880 + if fuse_shared_experts

881 + else {".shared_experts.w2": ".shared_experts.down_proj"}

882 + )

883 return WeightsMapper(

884 orig_to_new_prefix={

885 "layers.": "model.layers.",

@@ -785,9 +894,7 @@ def _make_deepseek_v4_weights_mapper(expert_dtype: str) -> WeightsMapper:

894 "embed.weight": "embed_tokens.weight",

895 ".ffn.gate.bias": ".ffn.gate.e_score_correction_bias",

896 },

788 - orig_to_new_substr={

789 - ".shared_experts.w2": ".shared_experts.down_proj",

790 - },

897 + orig_to_new_substr=substr_map,

898 )

899

900

@@ -804,8 +911,11 @@ class DeepseekV4ForCausalLM(nn.Module, SupportsPP, SupportsEagle3):

911 config = vllm_config.model_config.hf_config

912 self.config = config

913 expert_dtype = getattr(config, "expert_dtype", "fp4")

807 - if expert_dtype != "fp4":

808 - self.hf_to_vllm_mapper = _make_deepseek_v4_weights_mapper(expert_dtype)

914 + fuse_shared_experts = _fuse_shared_experts_enabled(config)

915 + if expert_dtype != "fp4" or fuse_shared_experts:

916 + self.hf_to_vllm_mapper = _make_deepseek_v4_weights_mapper(

917 + expert_dtype, fuse_shared_experts=fuse_shared_experts

918 + )

919

920 self.model = self.model_cls(

921 vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")

vllm/models/deepseek_v4/amd/mtp.py

+17

@@ -334,6 +334,21 @@ class DeepSeekV4MTP(nn.Module):

334 params_dict = dict(self.named_parameters())

335 loaded_params: set[str] = set()

336

337 + def _resolve_scale_name(name: str) -> str:

338 + # Quark checkpoints name FP8 block scales ``.weight_scale``,

339 + # but block-FP8 layers register them as ``.weight_scale_inv``

340 + # while MXFP4 experts register ``.weight_scale``. Auto-detect:

341 + # rename to ``_inv`` only when that variant exists and the plain

342 + # one does not.

343 + if name.endswith(".weight_scale") and name not in params_dict:

344 + inv = name.removesuffix(".weight_scale") + ".weight_scale_inv"

345 + if inv in params_dict:

346 + return inv

347 + # Otherwise leave the name unchanged: either it already matches a

348 + # param, or it is genuinely unknown and should surface the normal

349 + # KeyError downstream rather than be silently rewritten.

350 + return name

351 +

352 # TP for attention

353 tp_size = get_tensor_model_parallel_world_size()

354 tp_rank = get_tensor_model_parallel_rank()

@@ -393,6 +408,7 @@ class DeepSeekV4MTP(nn.Module):

408 if weight_name not in name:

409 continue

410 name = name.replace(weight_name, param_name)

411 + name = _resolve_scale_name(name)

412

413 param = params_dict[name]

414 weight_loader = param.weight_loader

@@ -447,6 +463,7 @@ class DeepSeekV4MTP(nn.Module):

463 )

464 if name.endswith(".ffn.gate.bias"):

465 name = name.replace(".bias", ".e_score_correction_bias")

466 + name = _resolve_scale_name(name)

467 param = params_dict[name]

468 weight_loader = getattr(

469 param, "weight_loader", default_weight_loader

vllm/models/deepseek_v4/quant_config.py

+37 -2

@@ -4,7 +4,7 @@

5 from __future__ import annotations

7 -from typing import TYPE_CHECKING

7 +from typing import TYPE_CHECKING, cast

9 from vllm.config import get_current_vllm_config

10 from vllm.model_executor.layers.fused_moe import (

@@ -117,13 +117,29 @@ class DeepseekV4FP8Config(Fp8Config):

117 def get_name(cls) -> QuantizationMethods:

118 return "deepseek_v4_fp8"

119

120 + @staticmethod

121 + def _is_quark_mxfp4_ocp(hf_quant_cfg: dict) -> bool:

122 + """True for AMD-Quark exports whose global scheme is MXFP4."""

123 + weight = (hf_quant_cfg.get("global_quant_config") or {}).get("weight") or {}

124 + return (

125 + weight.get("dtype") == "fp4"

126 + and weight.get("qscheme") == "per_group"

127 + and weight.get("group_size") == 32

128 + )

129 +

130 @classmethod

131 def override_quantization_method(

132 cls, hf_quant_cfg, user_quant, hf_config=None

133 ) -> QuantizationMethods | None:

134 if not (

135 isinstance(hf_quant_cfg, dict)

126 - and hf_quant_cfg.get("quant_method") in ("fp8", "deepseek_v4_fp8")

136 + and (

137 + hf_quant_cfg.get("quant_method") in ("fp8", "deepseek_v4_fp8")

138 + or (

139 + hf_quant_cfg.get("quant_method") == "quark"

140 + and cls._is_quark_mxfp4_ocp(hf_quant_cfg)

141 + )

142 + )

143 ):

144 return None

145 model_type = getattr(hf_config, "model_type", None)

@@ -131,6 +147,25 @@ class DeepseekV4FP8Config(Fp8Config):

147 return "deepseek_v4_fp8"

148 return None

149

150 + @classmethod

151 + def from_config(cls, config: dict) -> DeepseekV4FP8Config:

152 + # Reroute AMD-Quark fused shared expert MXFP4 checkpoints onto the fp8

153 + # path: the runtime layout matches the DeepSeek-native fp8 checkpoint,

154 + # so translate the schema into format Fp8Config.from_config expects.

155 + if config.get("quant_method") == "quark":

156 + quark_exclude = config.get("exclude") or []

157 + config = {

158 + "quant_method": "fp8",

159 + "activation_scheme": "dynamic",

160 + "fmt": "e4m3",

161 + "scale_fmt": "ue8m0",

162 + "weight_block_size": [128, 128],

163 + "ignored_layers": [

164 + name for name in quark_exclude if isinstance(name, str)

165 + ],

166 + }

167 + return cast("DeepseekV4FP8Config", super().from_config(config))

168 +

169 def get_quant_method(self, layer, prefix):

170 if isinstance(layer, RoutedExperts):

171 if is_layer_skipped(