@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 287471b994

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 3, 09:15 UTC

[Model Refactoring] Migrate DeepSeek V4 to vllm/models/ [1/N] (#43004)

Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>

Woosuk Kwon committed May 18, 2026 at 19:50 UTC 287471b99442b44c5a16c4d70b0f3e178dd52732

12 files changed +189 -126

tests/models/test_deepseek_v4_mega_moe.py

+1 -1

@@ -6,7 +6,7 @@ from types import SimpleNamespace

6 import pytest

7 import torch

9 -from vllm.model_executor.models.deepseek_v4 import (

9 +from vllm.models.deepseek_v4.nvidia.deepseek_v4 import (

10 DeepseekV4MegaMoEExperts,

11 _stage_deepseek_v4_mega_moe_inputs,

12 make_deepseek_v4_expert_params_mapping,

vllm/model_executor/layers/quantization/__init__.py

+1 -1

@@ -113,7 +113,7 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]:

113 # lazy import to avoid triggering `torch.compile` too early

114 from vllm.config.quantization import _ONLINE_SHORTHANDS

115 from vllm.model_executor.layers.quantization.quark.quark import QuarkConfig

116 - from vllm.model_executor.models.deepseek_v4 import DeepseekV4FP8Config

116 + from vllm.models.deepseek_v4 import DeepseekV4FP8Config

117

118 from .auto_gptq import AutoGPTQConfig

119 from .awq import AWQConfig

vllm/model_executor/models/registry.py

+26 -5

@@ -6,6 +6,7 @@ Whenever you add an architecture to this page, please also update

6 """

8 import importlib

9 +import importlib.util

10 import json

11 import os

12 import pickle

@@ -97,7 +98,7 @@ _TEXT_GENERATION_MODELS = {

98 "DeepseekV2ForCausalLM": ("deepseek_v2", "DeepseekV2ForCausalLM"),

99 "DeepseekV3ForCausalLM": ("deepseek_v2", "DeepseekV3ForCausalLM"),

100 "DeepseekV32ForCausalLM": ("deepseek_v2", "DeepseekV3ForCausalLM"),

100 - "DeepseekV4ForCausalLM": ("deepseek_v4", "DeepseekV4ForCausalLM"),

101 + "DeepseekV4ForCausalLM": ("vllm.models.deepseek_v4", "DeepseekV4ForCausalLM"),

102 "Dots1ForCausalLM": ("dots1", "Dots1ForCausalLM"),

103 "Ernie4_5ForCausalLM": ("ernie45", "Ernie4_5ForCausalLM"),

104 "Ernie4_5_MoeForCausalLM": ("ernie45_moe", "Ernie4_5_MoeForCausalLM"),

@@ -611,7 +612,7 @@ _SPECULATIVE_DECODING_MODELS = {

612 "Eagle3DeepseekV3ForCausalLM": ("deepseek_eagle3", "Eagle3DeepseekV2ForCausalLM"),

613 "EagleDeepSeekMTPModel": ("deepseek_eagle", "EagleDeepseekV3ForCausalLM"),

614 "DeepSeekMTPModel": ("deepseek_mtp", "DeepSeekMTP"),

614 - "DeepSeekV4MTPModel": ("deepseek_v4_mtp", "DeepSeekV4MTP"),

615 + "DeepSeekV4MTPModel": ("vllm.models.deepseek_v4", "DeepSeekV4MTP"),

616 "Gemma4MTPModel": ("gemma4_mtp", "Gemma4MTP"),

617 "ErnieMTPModel": ("ernie_mtp", "ErnieMTP"),

618 "ExaoneMoeMTP": ("exaone_moe_mtp", "ExaoneMoeMTP"),

@@ -870,10 +871,21 @@ class _LazyRegisteredModel(_BaseRegisteredModel):

871

872 @logtime(logger=logger, msg="Registry inspect model class")

873 def inspect_model_cls(self) -> _ModelInfo:

873 - model_path = Path(__file__).parent / f"{self.module_name.split('.')[-1]}.py"

874 + # Modules registered with a non-default location (e.g. the

875 + # hardware-isolated ``vllm.models.<name>`` layout) live outside

876 + # ``vllm/model_executor/models``. Resolve the module spec directly

877 + # so the file-hash cache stays warm for them.

878 + if self.module_name.startswith("vllm.model_executor.models."):

879 + model_path = Path(__file__).parent / f"{self.module_name.split('.')[-1]}.py"

880 + else:

881 + try:

882 + spec = importlib.util.find_spec(self.module_name)

883 + except (ImportError, ValueError):

884 + spec = None

885 + model_path = Path(spec.origin) if spec is not None and spec.origin else None

886 module_hash = None

887

876 - if model_path.exists():

888 + if model_path is not None and model_path.exists():

889 with open(model_path, "rb") as f:

890 module_hash = safe_hash(f.read(), usedforsecurity=False).hexdigest()

891

@@ -1328,10 +1340,19 @@ class _ModelRegistry:

1340 return model_cls.supports_transcription_only

1341

1342

1343 +def _resolve_module_name(mod_relname: str) -> str:

1344 + # Allow registry entries to point at fully-qualified module paths (e.g.

1345 + # ``vllm.models.deepseek_v4``) for models that live outside the legacy

1346 + # ``vllm.model_executor.models`` flat layout.

1347 + if mod_relname.startswith("vllm."):

1348 + return mod_relname

1349 + return f"vllm.model_executor.models.{mod_relname}"

1350 +

1351 +

1352 ModelRegistry = _ModelRegistry(

1353 {

1354 model_arch: _LazyRegisteredModel(

1334 - module_name=f"vllm.model_executor.models.{mod_relname}",

1355 + module_name=_resolve_module_name(mod_relname),

1356 class_name=cls_name,

1357 )

1358 for model_arch, (mod_relname, cls_name) in _VLLM_MODELS.items()

vllm/models/__init__.py new

@@ -0,0 +1,2 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

vllm/models/deepseek_v4/__init__.py new

+30

@@ -0,0 +1,30 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +"""DeepSeek V4 model — hardware-isolated entry point.

4 +

5 +The actual implementation lives under ``nvidia/`` and ``amd/``; this module

6 +picks the right one for the current platform and re-exports the public

7 +classes used by the model registry and quantization config lookup.

8 +"""

9 +

10 +from typing import TYPE_CHECKING

11 +

12 +from vllm.platforms import current_platform

13 +

14 +from .quant_config import DeepseekV4FP8Config

15 +

16 +# Pick the per-platform implementation. The NVIDIA branch is the static

17 +# default that mypy sees; the ROCm branch overrides it at runtime and is

18 +# kept type-compatible via ``# type: ignore[assignment]``.

19 +if TYPE_CHECKING or not current_platform.is_rocm():

20 + from .nvidia.deepseek_v4 import DeepseekV4ForCausalLM

21 + from .nvidia.deepseek_v4_mtp import DeepSeekV4MTP

22 +else:

23 + from .amd.deepseek_v4 import DeepseekV4ForCausalLM # type: ignore[assignment]

24 + from .amd.deepseek_v4_mtp import DeepSeekV4MTP # type: ignore[assignment]

25 +

26 +__all__ = [

27 + "DeepSeekV4MTP",

28 + "DeepseekV4FP8Config",

29 + "DeepseekV4ForCausalLM",

30 +]

vllm/models/deepseek_v4/amd/__init__.py new

@@ -0,0 +1,2 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

vllm/models/deepseek_v4/amd/deepseek_v4.py new

@@ -0,0 +1 @@

1 +../nvidia/deepseek_v4.py

\ No newline at end of file

vllm/models/deepseek_v4/amd/deepseek_v4_mtp.py new

@@ -0,0 +1 @@

1 +../nvidia/deepseek_v4_mtp.py

\ No newline at end of file

vllm/models/deepseek_v4/nvidia/__init__.py new

@@ -0,0 +1,2 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

vllm/models/deepseek_v4/nvidia/deepseek_v4.py renamed

+11 -114

@@ -9,7 +9,7 @@ import torch

9 import torch.nn as nn

11 from vllm.compilation.decorators import support_torch_compile

12 -from vllm.config import VllmConfig, get_current_vllm_config

12 +from vllm.config import VllmConfig

13 from vllm.distributed import (

14 get_ep_group,

15 get_pp_group,

@@ -24,7 +24,6 @@ from vllm.model_executor.layers.deepseek_v4_attention import (

24 DeepseekV4MultiHeadLatentAttentionWrapper,

25 )

26 from vllm.model_executor.layers.fused_moe import FusedMoE

27 -from vllm.model_executor.layers.fused_moe.layer import UnquantizedFusedMoEMethod

27 from vllm.model_executor.layers.fused_moe.router.fused_topk_bias_router import (

28 fused_topk_bias,

29 )

@@ -44,15 +43,7 @@ from vllm.model_executor.layers.mhc import (

43 MHCPostOp,

44 MHCPreOp,

45 )

47 -from vllm.model_executor.layers.quantization import (

48 - QuantizationConfig,

49 - QuantizationMethods,

50 -)

51 -from vllm.model_executor.layers.quantization.fp8 import Fp8Config

52 -from vllm.model_executor.layers.quantization.mxfp4 import Mxfp4MoEMethod

53 -from vllm.model_executor.layers.quantization.utils.quant_utils import (

54 - is_layer_skipped,

55 -)

46 +from vllm.model_executor.layers.quantization import QuantizationConfig

47 from vllm.model_executor.layers.rotary_embedding import get_rope

48 from vllm.model_executor.layers.vocab_parallel_embedding import (

49 ParallelLMHead,

@@ -60,13 +51,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (

51 )

52 from vllm.model_executor.model_loader.weight_utils import default_weight_loader

53 from vllm.model_executor.models.interfaces import SupportsPP

63 -from vllm.model_executor.utils import set_weight_attrs

64 -from vllm.platforms import current_platform

65 -from vllm.sequence import IntermediateTensors

66 -from vllm.triton_utils import tl, triton

67 -from vllm.utils.torch_utils import direct_register_custom_op

68 -

69 -from .utils import (

54 +from vllm.model_executor.models.utils import (

55 AutoWeightsLoader,

56 PPMissingLayer,

57 WeightsMapper,

@@ -75,8 +60,11 @@ from .utils import (

60 make_layers,

61 maybe_prefix,

62 )

78 -

79 -_DEEPSEEK_V4_EXPERT_DTYPES = ("fp4", "fp8")

63 +from vllm.model_executor.utils import set_weight_attrs

64 +from vllm.platforms import current_platform

65 +from vllm.sequence import IntermediateTensors

66 +from vllm.triton_utils import tl, triton

67 +from vllm.utils.torch_utils import direct_register_custom_op

70 class DeepseekV4MLP(nn.Module):

@@ -130,97 +118,6 @@ class DeepseekV4MLP(nn.Module):

118 return x

119

120

133 -class DeepseekV4FP8Config(Fp8Config):

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

135 -

136 - DeepSeek V4 checkpoints always use FP8 block quantization for

137 - linear/attention layers. The MoE expert weights vary by checkpoint:

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

139 - with ue8m0 (e8m0fnu) FP8 linear scales.

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

141 - experts with float32 FP8 linear scales.

142 -

143 - The dispatch and the linear scale dtype are both keyed off

144 - ``expert_dtype`` from the model's hf_config; missing values default

145 - to ``"fp4"`` so existing FP4 checkpoints stay unchanged.

146 -

147 - NOTE: ``expert_dtype`` is resolved lazily because this config is

148 - constructed during VllmConfig setup, before ``set_current_vllm_config``

149 - is active. Reading hf_config eagerly in ``__init__`` would always see

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

151 - """

152 -

153 - def __init__(self, *args, **kwargs):

154 - super().__init__(*args, **kwargs)

155 - self._resolved_expert_dtype: str | None = None

156 - # ``is_scale_e8m0`` is a property that resolves on first read,

157 - # by which time the current vllm_config has been set.

158 -

159 - @property

160 - def expert_dtype(self) -> str:

161 - if self._resolved_expert_dtype is None:

162 - try:

163 - hf_config = get_current_vllm_config().model_config.hf_config

164 - except Exception:

165 - # vllm_config not yet set; defer the decision until a

166 - # later call lands inside set_current_vllm_config.

167 - return "fp4"

168 - expert_dtype = getattr(hf_config, "expert_dtype", "fp4")

169 - if expert_dtype not in _DEEPSEEK_V4_EXPERT_DTYPES:

170 - raise ValueError(

171 - f"Unsupported DeepSeek V4 expert_dtype={expert_dtype!r}; "

172 - f"expected one of {_DEEPSEEK_V4_EXPERT_DTYPES}."

173 - )

174 - self._resolved_expert_dtype = expert_dtype

175 - from vllm.logger import init_logger

176 -

177 - init_logger(__name__).info_once(

178 - "DeepSeek V4 expert_dtype resolved to %r", expert_dtype

179 - )

180 - return self._resolved_expert_dtype

181 -

182 - @property

183 - def is_scale_e8m0(self) -> bool:

184 - # FP4 checkpoints store FP8 linear scales as e8m0fnu; FP8 expert

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

186 - return self.expert_dtype == "fp4"

187 -

188 - @classmethod

189 - def get_name(cls) -> QuantizationMethods:

190 - return "deepseek_v4_fp8"

191 -

192 - @classmethod

193 - def override_quantization_method(

194 - cls, hf_quant_cfg, user_quant, hf_config=None

195 - ) -> QuantizationMethods | None:

196 - if not (

197 - isinstance(hf_quant_cfg, dict)

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

199 - ):

200 - return None

201 - model_type = getattr(hf_config, "model_type", None)

202 - if model_type == "deepseek_v4" or user_quant == "deepseek_v4_fp8":

203 - return "deepseek_v4_fp8"

204 - return None

205 -

206 - def get_quant_method(self, layer, prefix):

207 - if isinstance(layer, FusedMoE):

208 - if is_layer_skipped(

209 - prefix=prefix,

210 - ignored_layers=self.ignored_layers,

211 - fused_mapping=self.packed_modules_mapping,

212 - ):

213 - return UnquantizedFusedMoEMethod(layer.moe_config)

214 - if self.expert_dtype == "fp4":

215 - return Mxfp4MoEMethod(layer.moe_config)

216 - # expert_dtype == "fp8": fall through to Fp8Config which

217 - # returns Fp8MoEMethod with block-wise float32 scales.

218 - return super().get_quant_method(layer, prefix)

219 -

220 - def is_mxfp4_quant(self, prefix, layer):

221 - return isinstance(layer, FusedMoE) and self.expert_dtype == "fp4"

222 -

223 -

121 @triton.jit

122 def _deepseek_v4_stage_mega_moe_inputs_kernel(

123 hidden_states,

@@ -1535,7 +1432,7 @@ class DeepseekV4Model(nn.Module):

1432 ):

1433 loaded_weight = loaded_weight.view(torch.uint8)

1434 for mapping in expert_mapping:

1538 - param_name, weight_name, expert_id, shard_id = mapping

1435 + param_name, weight_name, expert_id, expert_shard_id = mapping

1436 if weight_name not in name:

1437 continue

1438 name_mapped = name.replace(weight_name, param_name)

@@ -1552,7 +1449,7 @@ class DeepseekV4Model(nn.Module):

1449 param,

1450 loaded_weight,

1451 name_mapped,

1555 - shard_id=shard_id,

1452 + shard_id=expert_shard_id,

1453 expert_id=expert_id,

1454 return_success=True,

1455 )

@@ -1673,7 +1570,7 @@ class DeepseekV4ForCausalLM(nn.Module, SupportsPP):

1570 else:

1571 self.lm_head = PPMissingLayer()

1572 self.logits_processor = LogitsProcessor(config.vocab_size)

1676 - self.make_empty_intermediate_tensors = (

1573 + self.make_empty_intermediate_tensors = ( # type: ignore[method-assign]

1574 self.model.make_empty_intermediate_tensors

1575 )

1576

vllm/models/deepseek_v4/nvidia/deepseek_v4_mtp.py renamed

+6 -5

@@ -34,16 +34,16 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (

34 VocabParallelEmbedding,

35 )

36 from vllm.model_executor.model_loader.weight_utils import default_weight_loader

37 +from vllm.model_executor.models.deepseek_mtp import SharedHead

38 +from vllm.model_executor.models.deepseek_v2 import get_spec_layer_idx_from_weight_name

39 +from vllm.model_executor.models.utils import maybe_prefix

40 from vllm.platforms import current_platform

41 from vllm.sequence import IntermediateTensors

40 -from .deepseek_mtp import SharedHead

41 -from .deepseek_v2 import get_spec_layer_idx_from_weight_name

43 from .deepseek_v4 import (

44 DeepseekV4DecoderLayer,

45 make_deepseek_v4_expert_params_mapping,

46 )

46 -from .utils import maybe_prefix

48 logger = init_logger(__name__)

@@ -68,6 +68,7 @@ class DeepSeekV4MultiTokenPredictorLayer(nn.Module):

68 ) -> None:

69 super().__init__()

71 + assert vllm_config.speculative_config is not None

72 config = vllm_config.speculative_config.draft_model_config.hf_config

73 self.config = config

74 quant_config = vllm_config.quant_config

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

408 ):

409 loaded_weight = loaded_weight.view(torch.uint8)

410 for mapping in expert_mapping:

410 - param_name, weight_name, expert_id, shard_id = mapping

411 + param_name, weight_name, expert_id, expert_shard_id = mapping

412 if weight_name not in name:

413 continue

414 name_mapped = name.replace(weight_name, param_name)

@@ -422,7 +423,7 @@ class DeepSeekV4MTP(nn.Module):

423 param,

424 loaded_weight,

425 name_mapped,

425 - shard_id=shard_id,

426 + shard_id=expert_shard_id,

427 expert_id=expert_id,

428 return_success=True,

429 )

vllm/models/deepseek_v4/quant_config.py new

+106

@@ -0,0 +1,106 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +"""Quantization config for DeepSeek V4."""

4 +

5 +from vllm.config import get_current_vllm_config

6 +from vllm.model_executor.layers.fused_moe import FusedMoE

7 +from vllm.model_executor.layers.fused_moe.layer import UnquantizedFusedMoEMethod

8 +from vllm.model_executor.layers.quantization import QuantizationMethods

9 +from vllm.model_executor.layers.quantization.fp8 import Fp8Config

10 +from vllm.model_executor.layers.quantization.mxfp4 import Mxfp4MoEMethod

11 +from vllm.model_executor.layers.quantization.utils.quant_utils import (

12 + is_layer_skipped,

13 +)

14 +

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

16 +

17 +

18 +class DeepseekV4FP8Config(Fp8Config):

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

20 +

21 + DeepSeek V4 checkpoints always use FP8 block quantization for

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

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

24 + with ue8m0 (e8m0fnu) FP8 linear scales.

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

26 + experts with float32 FP8 linear scales.

27 +

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

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

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

31 +

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

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

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

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

36 + """

37 +

38 + def __init__(self, *args, **kwargs):

39 + super().__init__(*args, **kwargs)

40 + self._resolved_expert_dtype: str | None = None

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

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

43 +

44 + @property

45 + def expert_dtype(self) -> str:

46 + if self._resolved_expert_dtype is None:

47 + try:

48 + hf_config = get_current_vllm_config().model_config.hf_config

49 + except Exception:

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

51 + # later call lands inside set_current_vllm_config.

52 + return "fp4"

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

54 + if expert_dtype not in _DEEPSEEK_V4_EXPERT_DTYPES:

55 + raise ValueError(

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

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

58 + )

59 + self._resolved_expert_dtype = expert_dtype

60 + from vllm.logger import init_logger

61 +

62 + init_logger(__name__).info_once(

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

64 + )

65 + return self._resolved_expert_dtype

66 +

67 + @property

68 + def is_scale_e8m0(self) -> bool:

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

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

71 + return self.expert_dtype == "fp4"

72 +

73 + @classmethod

74 + def get_name(cls) -> QuantizationMethods:

75 + return "deepseek_v4_fp8"

76 +

77 + @classmethod

78 + def override_quantization_method(

79 + cls, hf_quant_cfg, user_quant, hf_config=None

80 + ) -> QuantizationMethods | None:

81 + if not (

82 + isinstance(hf_quant_cfg, dict)

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

84 + ):

85 + return None

86 + model_type = getattr(hf_config, "model_type", None)

87 + if model_type == "deepseek_v4" or user_quant == "deepseek_v4_fp8":

88 + return "deepseek_v4_fp8"

89 + return None

90 +

91 + def get_quant_method(self, layer, prefix):

92 + if isinstance(layer, FusedMoE):

93 + if is_layer_skipped(

94 + prefix=prefix,

95 + ignored_layers=self.ignored_layers,

96 + fused_mapping=self.packed_modules_mapping,

97 + ):

98 + return UnquantizedFusedMoEMethod(layer.moe_config)

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

100 + return Mxfp4MoEMethod(layer.moe_config)

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

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

103 + return super().get_quant_method(layer, prefix)

104 +

105 + def is_mxfp4_quant(self, prefix, layer):

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