@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 2c8b76c5cb

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 4, 15:00 UTC

[Model][DSV4] Support base model (#41006)

Signed-off-by: Jee Jee Li <pandaleefree@gmail.com>

Jee Jee Li committed Apr 28, 2026 at 08:16 UTC 2c8b76c5cb2683f05650f20d90a63f3d9799e909

2 files changed +111 -23

vllm/model_executor/models/deepseek_v4.py

+93 -19

@@ -10,7 +10,7 @@ import torch.nn as nn

10 import torch.nn.functional as F

12 from vllm.compilation.decorators import support_torch_compile

13 -from vllm.config import VllmConfig

13 +from vllm.config import VllmConfig, get_current_vllm_config

14 from vllm.distributed import (

15 get_ep_group,

16 get_tensor_model_parallel_rank,

@@ -65,6 +65,8 @@ from .utils import (

65 maybe_prefix,

66 )

68 +_DEEPSEEK_V4_EXPERT_DTYPES = ("fp4", "fp8")

69 +

71 class DeepseekV4MLP(nn.Module):

72 def __init__(

@@ -118,16 +120,59 @@ class DeepseekV4MLP(nn.Module):

120

121

122 class DeepseekV4FP8Config(Fp8Config):

121 - """FP8 config that routes MoE layers to MXFP4 quantization.

122 -

123 - DeepSeek V4 checkpoints use FP8 for linear/attention layers but

124 - MXFP4 for MoE expert weights. This config inherits standard FP8

125 - behavior and overrides only the MoE dispatch.

123 + """FP8 config for DeepSeek V4 with expert-dtype-aware MoE dispatch.

124 +

125 + DeepSeek V4 checkpoints always use FP8 block quantization for

126 + linear/attention layers. The MoE expert weights vary by checkpoint:

127 + - ``expert_dtype="fp4"`` (e.g. DeepSeek-V4-Flash): MXFP4 experts

128 + with ue8m0 (e8m0fnu) FP8 linear scales.

129 + - ``expert_dtype="fp8"`` (e.g. DeepSeek-V4-Flash-Base): FP8 block

130 + experts with float32 FP8 linear scales.

131 +

132 + The dispatch and the linear scale dtype are both keyed off

133 + ``expert_dtype`` from the model's hf_config; missing values default

134 + to ``"fp4"`` so existing FP4 checkpoints stay unchanged.

135 +

136 + NOTE: ``expert_dtype`` is resolved lazily because this config is

137 + constructed during VllmConfig setup, before ``set_current_vllm_config``

138 + is active. Reading hf_config eagerly in ``__init__`` would always see

139 + the default ``"fp4"`` and silently misroute Flash-Base checkpoints.

140 """

141

142 def __init__(self, *args, **kwargs):

143 super().__init__(*args, **kwargs)

130 - self.is_scale_e8m0: bool = True

144 + self._resolved_expert_dtype: str | None = None

145 + # ``is_scale_e8m0`` is a property that resolves on first read,

146 + # by which time the current vllm_config has been set.

147 +

148 + @property

149 + def expert_dtype(self) -> str:

150 + if self._resolved_expert_dtype is None:

151 + try:

152 + hf_config = get_current_vllm_config().model_config.hf_config

153 + except Exception:

154 + # vllm_config not yet set; defer the decision until a

155 + # later call lands inside set_current_vllm_config.

156 + return "fp4"

157 + expert_dtype = getattr(hf_config, "expert_dtype", "fp4")

158 + if expert_dtype not in _DEEPSEEK_V4_EXPERT_DTYPES:

159 + raise ValueError(

160 + f"Unsupported DeepSeek V4 expert_dtype={expert_dtype!r}; "

161 + f"expected one of {_DEEPSEEK_V4_EXPERT_DTYPES}."

162 + )

163 + self._resolved_expert_dtype = expert_dtype

164 + from vllm.logger import init_logger

165 +

166 + init_logger(__name__).info_once(

167 + "DeepSeek V4 expert_dtype resolved to %r", expert_dtype

168 + )

169 + return self._resolved_expert_dtype

170 +

171 + @property

172 + def is_scale_e8m0(self) -> bool:

173 + # FP4 checkpoints store FP8 linear scales as e8m0fnu; FP8 expert

174 + # checkpoints (Flash-Base) store them as float32.

175 + return self.expert_dtype == "fp4"

176

177 @classmethod

178 def get_name(cls) -> QuantizationMethods:

@@ -155,11 +200,14 @@ class DeepseekV4FP8Config(Fp8Config):

200 fused_mapping=self.packed_modules_mapping,

201 ):

202 return UnquantizedFusedMoEMethod(layer.moe_config)

158 - return Mxfp4MoEMethod(layer.moe_config)

203 + if self.expert_dtype == "fp4":

204 + return Mxfp4MoEMethod(layer.moe_config)

205 + # expert_dtype == "fp8": fall through to Fp8Config which

206 + # returns Fp8MoEMethod with block-wise float32 scales.

207 return super().get_quant_method(layer, prefix)

208

209 def is_mxfp4_quant(self, prefix, layer):

162 - return isinstance(layer, FusedMoE)

210 + return isinstance(layer, FusedMoE) and self.expert_dtype == "fp4"

211

212

213 @triton.jit

@@ -689,6 +737,12 @@ class DeepseekV4MoE(nn.Module):

737 raise NotImplementedError(

738 "DeepSeek V4 MegaMoE currently supports sqrtsoftplus routing only."

739 )

740 + if self.use_mega_moe and getattr(config, "expert_dtype", "fp4") != "fp4":

741 + raise NotImplementedError(

742 + "DeepSeek V4 MegaMoE only supports fp4 experts; got expert_dtype="

743 + f"{config.expert_dtype!r}. Drop --kernel-config moe_backend="

744 + "deep_gemm_mega_moe for this checkpoint."

745 + )

746

747 self.gate = GateLinear(

748 config.hidden_size,

@@ -1410,10 +1464,24 @@ def hc_head(

1464 return y.to(dtype)

1465

1466

1413 -class DeepseekV4ForCausalLM(nn.Module):

1414 - model_cls = DeepseekV4Model

1415 -

1416 - hf_to_vllm_mapper = WeightsMapper(

1467 +def _make_deepseek_v4_weights_mapper(expert_dtype: str) -> WeightsMapper:

1468 + if expert_dtype == "fp4":

1469 + # MXFP4 experts use Mxfp4MoEMethod, which registers scales as

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

1471 + # shared experts use Fp8LinearMethod's block scales, which

1472 + # register as ``weight_scale_inv``.

1473 + scale_regex = {

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

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

1476 + }

1477 + else:

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

1479 + # scales as ``w{13,2}_weight_scale_inv``. Map all ``.scale`` keys

1480 + # there.

1481 + scale_regex = {

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

1483 + }

1484 + return WeightsMapper(

1485 orig_to_new_prefix={

1486 "layers.": "model.layers.",

1487 "embed.": "model.embed.",

@@ -1421,12 +1489,7 @@ class DeepseekV4ForCausalLM(nn.Module):

1489 "hc_head": "model.hc_head",

1490 "mtp.": "model.mtp.",

1491 },

1424 - orig_to_new_regex={

1425 - # Routed MoE expert scales: experts.N.wX.scale -> .weight_scale

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

1427 - # Everything else (FP8 linear + shared experts): .scale -> .weight_scale_inv

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

1429 - },

1492 + orig_to_new_regex=scale_regex,

1493 orig_to_new_suffix={

1494 "head.weight": "lm_head.weight",

1495 "embed.weight": "embed_tokens.weight",

@@ -1438,11 +1501,22 @@ class DeepseekV4ForCausalLM(nn.Module):

1501 },

1502 )

1503

1504 +

1505 +class DeepseekV4ForCausalLM(nn.Module):

1506 + model_cls = DeepseekV4Model

1507 +

1508 + # Default mapper assumes the original FP4-expert checkpoint layout.

1509 + # Overridden per-instance in __init__ when expert_dtype != "fp4".

1510 + hf_to_vllm_mapper = _make_deepseek_v4_weights_mapper("fp4")

1511 +

1512 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

1513 super().__init__()

1514

1515 config = vllm_config.model_config.hf_config

1516 self.config = config

1517 + expert_dtype = getattr(config, "expert_dtype", "fp4")

1518 + if expert_dtype != "fp4":

1519 + self.hf_to_vllm_mapper = _make_deepseek_v4_weights_mapper(expert_dtype)

1520

1521 self.model = self.model_cls(

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

vllm/model_executor/models/deepseek_v4_mtp.py

+18 -4

@@ -48,9 +48,14 @@ from .utils import maybe_prefix

49 logger = init_logger(__name__)

51 -# MoE expert scales are fused into per-layer w13/w2 tensors; other FP8 linear

52 -# scales use `.weight_scale_inv`. Mirrors the regex in

53 -# DeepseekV4ForCausalLM.hf_to_vllm_mapper.

51 +# MoE expert scales are fused into per-layer w13/w2 tensors. The exact

52 +# parameter suffix depends on which FusedMoE method handles the experts:

53 +# - fp4 experts (Mxfp4MoEMethod) register ``w{1,2,3}_weight_scale``;

54 +# - fp8 experts (Fp8MoEMethod with block_quant=True) register

55 +# ``w{1,2,3}_weight_scale_inv``.

56 +# Other FP8 linear scales (including shared experts) always use

57 +# ``.weight_scale_inv``. Mirrors the per-instance mapper built by

58 +# ``_make_deepseek_v4_weights_mapper`` in deepseek_v4.py.

59 _EXPERT_SCALE_RE = re.compile(r"\.experts\.\d+\.w[123]\.scale$")

@@ -326,6 +331,15 @@ class DeepSeekV4MTP(nn.Module):

331 num_experts=self.config.n_routed_experts,

332 )

333

334 + # FP8 experts register ``..._weight_scale_inv`` (block_quant) while

335 + # FP4/MXFP4 experts register ``..._weight_scale``. Choose the suffix

336 + # for the rename below based on the model's expert dtype.

337 + expert_scale_suffix = (

338 + ".weight_scale"

339 + if getattr(self.config, "expert_dtype", "fp4") == "fp4"

340 + else ".weight_scale_inv"

341 + )

342 +

343 for name, loaded_weight in weights:

344 mtp_layer_idx = _find_mtp_layer_idx(name)

345 # V4 checkpoints store MTP weights as `mtp.{i}.*`; remap to

@@ -347,7 +361,7 @@ class DeepSeekV4MTP(nn.Module):

361 continue

362 if name.endswith(".scale"):

363 suffix = (

350 - ".weight_scale"

364 + expert_scale_suffix

365 if _EXPERT_SCALE_RE.search(name)

366 else ".weight_scale_inv"

367 )