| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """DeepSeek-V4 vision variant (e.g. DeepSeek-V4-Flash-Vision-Exp). |
| 4 | |
| 5 | Thin multimodal wrapper around the text-only ``DeepseekV4ForCausalLM``: |
| 6 | |
| 7 | - ``vision`` ViT + ``aligner`` produce per-image embeddings for the IMAGE |
| 8 | sentinel positions; four learned vectors (``image_start`` / ``image_pad`` / |
| 9 | ``image_newline`` / ``image_end``) fill the remaining sentinel positions. |
| 10 | - Image placeholders (``<｜deepseek_image｜>``) are expanded by the processor |
| 11 | in ``common/mm_preprocess.py`` into sentinel blocks borrowing reserved |
| 12 | in-vocab tokens ``<|place_holder_mm_span_0431|>``..``_0435|>`` |
| 13 | (see ``common/vision.py`` for the tower itself). |
| 14 | - Merged embeddings enter the text model via ``inputs_embeds``, i.e. before |
| 15 | its hyper-connection stream expansion. Raw ``input_ids`` still flow into |
| 16 | the model so the MoE router can apply ``bias_vl`` to image tokens |
| 17 | (``requires_raw_input_tokens``). |
| 18 | """ |
| 19 | |
| 20 | from collections.abc import Iterable |
| 21 | |
| 22 | import torch |
| 23 | from torch import nn |
| 24 | |
| 25 | from vllm.model_executor.models.interfaces import ( |
| 26 | MultiModalEmbeddings, |
| 27 | SupportsEagle3, |
| 28 | SupportsLoRA, |
| 29 | SupportsMultiModal, |
| 30 | SupportsPP, |
| 31 | ) |
| 32 | from vllm.model_executor.models.module_mapping import MultiModelKeys |
| 33 | from vllm.model_executor.models.utils import ( |
| 34 | AutoWeightsLoader, |
| 35 | WeightsMapper, |
| 36 | init_vllm_registered_model, |
| 37 | maybe_prefix, |
| 38 | ) |
| 39 | from vllm.multimodal import MULTIMODAL_REGISTRY |
| 40 | from vllm.multimodal.inputs import MultiModalKwargsItem |
| 41 | |
| 42 | from .mm_preprocess import ( |
| 43 | IMAGE_PLACEHOLDER, |
| 44 | IMAGE_SENTINEL_BASE_ID, |
| 45 | DeepseekV4VLDummyInputsBuilder, |
| 46 | DeepseekV4VLMultiModalProcessor, |
| 47 | DeepseekV4VLProcessingInfo, |
| 48 | image_sentinel_mask, |
| 49 | ) |
| 50 | from .vision import DeepseekV4Aligner, DeepseekV4ViT |
| 51 | |
| 52 | |
| 53 | def _make_deepseek_v4_vl_weights_mapper( |
| 54 | text_mapper: WeightsMapper, image_enabled: bool |
| 55 | ) -> WeightsMapper: |
| 56 | """Text-checkpoint mapping rules re-rooted under ``language_model.``.""" |
| 57 | orig_to_new_prefix = { |
| 58 | src: None if dst is None else f"language_model.{dst}" |
| 59 | for src, dst in text_mapper.orig_to_new_prefix.items() |
| 60 | } |
| 61 | if not image_enabled: |
| 62 | orig_to_new_prefix.update({"vision.": None, "aligner.": None, "image_": None}) |
| 63 | return WeightsMapper( |
| 64 | orig_to_new_renaming=text_mapper.orig_to_new_renaming, |
| 65 | orig_to_new_prefix=orig_to_new_prefix, |
| 66 | orig_to_new_regex=text_mapper.orig_to_new_regex, |
| 67 | orig_to_new_stacked=text_mapper.orig_to_new_stacked, |
| 68 | orig_to_new_suffix={ |
| 69 | **text_mapper.orig_to_new_suffix, |
| 70 | "head.weight": "language_model.lm_head.weight", |
| 71 | }, |
| 72 | orig_to_new_substr={ |
| 73 | **text_mapper.orig_to_new_substr, |
| 74 | # Draft models load the checkpoint's MTP weights separately from |
| 75 | # the target model. |
| 76 | "mtp.": None, |
| 77 | }, |
| 78 | ) |
| 79 | |
| 80 | |
| 81 | @MULTIMODAL_REGISTRY.register_processor( |
| 82 | DeepseekV4VLMultiModalProcessor, |
| 83 | info=DeepseekV4VLProcessingInfo, |
| 84 | dummy_inputs=DeepseekV4VLDummyInputsBuilder, |
| 85 | ) |
| 86 | class DeepseekV4ForConditionalGeneration( |
| 87 | nn.Module, SupportsMultiModal, SupportsPP, SupportsEagle3, SupportsLoRA |
| 88 | ): |
| 89 | """Multimodal entry point for DeepSeek-V4 checkpoints with a vision tower. |
| 90 | |
| 91 | ``SupportsEagle3`` (aux hidden-state plumbing for MTP/DSpark drafters) |
| 92 | delegates through ``language_model`` via the protocol defaults. |
| 93 | """ |
| 94 | |
| 95 | packed_modules_mapping = { |
| 96 | "gate_up_proj": ["w1", "w3"], |
| 97 | "fused_wqa_wkv": ["wq_a", "wkv"], |
| 98 | "fused_wkv_wgate": ["wkv", "wgate"], |
| 99 | # for visual encoder |
| 100 | "wqkv": ["wqkv"], |
| 101 | "w1": ["w1"], |
| 102 | } |
| 103 | |
| 104 | # The MoE router needs raw token ids to detect image sentinel tokens |
| 105 | # (borrowed reserved ids, see common/mm_preprocess.py) and apply bias_vl. |
| 106 | requires_raw_input_tokens = True |
| 107 | supports_tower_connector_lora = True |
| 108 | |
| 109 | @classmethod |
| 110 | def get_placeholder_str(cls, modality: str, i: int) -> str | None: |
| 111 | if modality == "image": |
| 112 | return IMAGE_PLACEHOLDER |
| 113 | raise ValueError(f"Unsupported modality: {modality!r}") |
| 114 | |
| 115 | def __init__(self, *, vllm_config, prefix: str = "") -> None: |
| 116 | super().__init__() |
| 117 | model_config = vllm_config.model_config |
| 118 | config = model_config.hf_config |
| 119 | self.config = config |
| 120 | self.multimodal_config = model_config.multimodal_config |
| 121 | assert self.multimodal_config is not None |
| 122 | |
| 123 | image_enabled = ( |
| 124 | config.vision_n_layers > 0 |
| 125 | and self.multimodal_config.get_limit_per_prompt("image") > 0 |
| 126 | ) |
| 127 | with self._mark_tower_model(vllm_config, {"image"}): |
| 128 | self.vision: DeepseekV4ViT | None = None |
| 129 | self.aligner: DeepseekV4Aligner | None = None |
| 130 | self.image_start: nn.Parameter | None = None |
| 131 | self.image_end: nn.Parameter | None = None |
| 132 | self.image_newline: nn.Parameter | None = None |
| 133 | self.image_pad: nn.Parameter | None = None |
| 134 | if image_enabled: |
| 135 | self.vision = DeepseekV4ViT(config) |
| 136 | self.aligner = DeepseekV4Aligner(config) |
| 137 | for name in ( |
| 138 | "image_start", |
| 139 | "image_end", |
| 140 | "image_newline", |
| 141 | "image_pad", |
| 142 | ): |
| 143 | setattr( |
| 144 | self, |
| 145 | name, |
| 146 | nn.Parameter( |
| 147 | torch.empty(config.hidden_size, dtype=torch.float32) |
| 148 | ), |
| 149 | ) |
| 150 | self.vision.to(dtype=model_config.dtype) |
| 151 | self.aligner.to(dtype=model_config.dtype) |
| 152 | |
| 153 | with self._mark_language_model(vllm_config): |
| 154 | # The arch convertor routes any config with a vision tower to |
| 155 | # this wrapper class; mark the copy handed to the text backbone |
| 156 | # so it resolves to DeepseekV4ForCausalLM instead of recursing |
| 157 | # (with_hf_config deepcopies the config, the marker survives). |
| 158 | config._dsv4_vl_inner = True # type: ignore[attr-defined] |
| 159 | try: |
| 160 | self.language_model = init_vllm_registered_model( |
| 161 | vllm_config=vllm_config, |
| 162 | hf_config=config, |
| 163 | prefix=maybe_prefix(prefix, "language_model"), |
| 164 | architectures=["DeepseekV4ForCausalLM"], |
| 165 | ) |
| 166 | finally: |
| 167 | del config._dsv4_vl_inner # type: ignore[attr-defined] |
| 168 | # The outer mapper (see load_weights) fully resolves HF names into |
| 169 | # this wrapper's namespace before AutoWeightsLoader strips the |
| 170 | # "language_model." prefix and delegates to the child's load_weights, |
| 171 | # so the child's own mapper must be a no-op. Its suffix rules are not |
| 172 | # idempotent (e.g. "lm_head.weight".endswith("head.weight") would |
| 173 | # re-fire "head.weight" -> "lm_head.weight"). |
| 174 | text_mapper = self.language_model.hf_to_vllm_mapper |
| 175 | self.language_model.hf_to_vllm_mapper = WeightsMapper() |
| 176 | self.make_empty_intermediate_tensors = ( # type: ignore[method-assign] |
| 177 | self.language_model.make_empty_intermediate_tensors |
| 178 | ) |
| 179 | |
| 180 | self.hf_to_vllm_mapper = _make_deepseek_v4_vl_weights_mapper( |
| 181 | text_mapper, image_enabled |
| 182 | ) |
| 183 | self._weights_finalized = False |
| 184 | |
| 185 | def _parse_and_validate_image_input(self, **kwargs: object) -> dict | None: |
| 186 | patches = kwargs.pop("patches", None) |
| 187 | if patches is None: |
| 188 | return None |
| 189 | vit_grid = kwargs.pop("vit_grid", None) |
| 190 | llm_grid = kwargs.pop("llm_grid", None) |
| 191 | perm = kwargs.pop("perm", None) |
| 192 | assert vit_grid is not None and llm_grid is not None and perm is not None |
| 193 | return { |
| 194 | "patches": patches, |
| 195 | "vit_grid": vit_grid, |
| 196 | "llm_grid": llm_grid, |
| 197 | "perm": perm, |
| 198 | } |
| 199 | |
| 200 | def _encode_image( |
| 201 | self, |
| 202 | patches: torch.Tensor, |
| 203 | n_vit_h: int, |
| 204 | n_vit_w: int, |
| 205 | perm: torch.Tensor, |
| 206 | ) -> torch.Tensor: |
| 207 | assert self.vision is not None and self.aligner is not None |
| 208 | image_embeds = self.aligner( |
| 209 | self.vision(patches, n_vit_h, n_vit_w), n_vit_h, n_vit_w |
| 210 | ) |
| 211 | # Reorder into the N-layout block order used in the prompt. |
| 212 | return image_embeds[perm.to(image_embeds.device)] |
| 213 | |
| 214 | def _process_image_input( |
| 215 | self, |
| 216 | patches: torch.Tensor, |
| 217 | vit_grid: torch.Tensor, |
| 218 | llm_grid: torch.Tensor, |
| 219 | perm: torch.Tensor, |
| 220 | ) -> tuple[torch.Tensor, ...]: |
| 221 | assert self.vision is not None and self.aligner is not None |
| 222 | patches = patches.to(self.aligner.w1.weight.dtype) |
| 223 | |
| 224 | embeds: list[torch.Tensor] = [] |
| 225 | vit_offset = 0 |
| 226 | llm_offset = 0 |
| 227 | for (n_vit_h, n_vit_w), (n_llm_h, n_llm_w) in zip( |
| 228 | vit_grid.tolist(), llm_grid.tolist(), strict=True |
| 229 | ): |
| 230 | n_vit = n_vit_h * n_vit_w |
| 231 | n_llm = n_llm_h * n_llm_w |
| 232 | embeds.append( |
| 233 | self._encode_image( |
| 234 | patches[vit_offset : vit_offset + n_vit], |
| 235 | n_vit_h, |
| 236 | n_vit_w, |
| 237 | perm[llm_offset : llm_offset + n_llm], |
| 238 | ) |
| 239 | ) |
| 240 | vit_offset += n_vit |
| 241 | llm_offset += n_llm |
| 242 | return tuple(embeds) |
| 243 | |
| 244 | def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings: |
| 245 | image_input = self._parse_and_validate_image_input(**kwargs) |
| 246 | if image_input is None or self.vision is None: |
| 247 | return [] |
| 248 | return self._process_image_input( |
| 249 | image_input["patches"], |
| 250 | image_input["vit_grid"], |
| 251 | image_input["llm_grid"], |
| 252 | image_input["perm"], |
| 253 | ) |
| 254 | |
| 255 | def embed_input_ids( |
| 256 | self, |
| 257 | input_ids: torch.Tensor, |
| 258 | multimodal_embeddings: MultiModalEmbeddings | None = None, |
| 259 | *, |
| 260 | is_multimodal: torch.Tensor | None = None, |
| 261 | ) -> torch.Tensor: |
| 262 | from vllm.model_executor.models.utils import _merge_multimodal_embeddings |
| 263 | |
| 264 | # All ids are in-vocab here: image-block sentinels are borrowed |
| 265 | # reserved tokens (their embedding rows are always overwritten below). |
| 266 | inputs_embeds = self.language_model.embed_input_ids(input_ids) |
| 267 | |
| 268 | if self.image_start is not None: |
| 269 | # Branch-free sentinel overwrite: safe inside compiled/captured |
| 270 | # regions (no data-dependent control flow). |
| 271 | sentinel_mask = image_sentinel_mask(input_ids) |
| 272 | if is_multimodal is not None: |
| 273 | # IMAGE positions get vision embeddings via the merge below. |
| 274 | sentinel_mask = sentinel_mask & ~is_multimodal.to(input_ids.device) |
| 275 | table = torch.stack( |
| 276 | [ |
| 277 | self.image_start, |
| 278 | self.image_pad, |
| 279 | self.image_pad, |
| 280 | self.image_newline, |
| 281 | self.image_end, |
| 282 | ] |
| 283 | ).to(inputs_embeds.dtype) |
| 284 | idx = (input_ids - IMAGE_SENTINEL_BASE_ID).clamp(0, 4) |
| 285 | inputs_embeds = torch.where( |
| 286 | sentinel_mask.unsqueeze(-1), table[idx], inputs_embeds |
| 287 | ) |
| 288 | |
| 289 | if multimodal_embeddings is None or len(multimodal_embeddings) == 0: |
| 290 | return inputs_embeds |
| 291 | |
| 292 | assert is_multimodal is not None |
| 293 | return _merge_multimodal_embeddings( |
| 294 | inputs_embeds=inputs_embeds, |
| 295 | multimodal_embeddings=multimodal_embeddings, |
| 296 | is_multimodal=is_multimodal, |
| 297 | ) |
| 298 | |
| 299 | def forward( |
| 300 | self, |
| 301 | input_ids: torch.Tensor, |
| 302 | positions: torch.Tensor, |
| 303 | intermediate_tensors=None, |
| 304 | inputs_embeds: torch.Tensor | None = None, |
| 305 | **kwargs, |
| 306 | ) -> torch.Tensor: |
| 307 | return self.language_model( |
| 308 | input_ids, positions, intermediate_tensors, inputs_embeds |
| 309 | ) |
| 310 | |
| 311 | def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None: |
| 312 | return self.language_model.compute_logits(hidden_states) |
| 313 | |
| 314 | def compute_logits_local(self, hidden_states: torch.Tensor) -> torch.Tensor: |
| 315 | return self.language_model.compute_logits_local(hidden_states) |
| 316 | |
| 317 | def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: |
| 318 | return self.language_model.get_expert_mapping() |
| 319 | |
| 320 | def get_mtp_target_hidden_states(self) -> torch.Tensor | None: |
| 321 | """Pre-hc_head residual stream buffer for the MTP/DSpark draft model.""" |
| 322 | return self.language_model.get_mtp_target_hidden_states() |
| 323 | |
| 324 | def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: |
| 325 | child_finalizes = getattr( |
| 326 | self.language_model, "finalizes_weights_during_load", True |
| 327 | ) |
| 328 | mapped = self.hf_to_vllm_mapper.apply(weights) |
| 329 | if child_finalizes: |
| 330 | # A child which finalizes inside load_weights must see all of its |
| 331 | # weights in one contiguous delegation from AutoWeightsLoader. |
| 332 | mapped = iter(sorted(mapped, key=lambda x: x[0])) |
| 333 | loader = AutoWeightsLoader(self) |
| 334 | loaded_params = loader.load_weights(mapped) |
| 335 | self._weights_finalized = child_finalizes |
| 336 | return loaded_params |
| 337 | |
| 338 | def process_weights_after_loading(self) -> None: |
| 339 | # Backbones such as the ROCm implementation require this to run only |
| 340 | # after the loader's generic per-layer quantization finalization. |
| 341 | if getattr(self, "_weights_finalized", False): |
| 342 | return |
| 343 | self.language_model.process_weights_after_loading() |
| 344 | self._weights_finalized = True |
| 345 | |
| 346 | def get_mm_mapping(self) -> MultiModelKeys: |
| 347 | """Get the module prefixes in the multimodal model.""" |
| 348 | return MultiModelKeys.from_string_field( |
| 349 | language_model="language_model", |
| 350 | connector="aligner", |
| 351 | tower_model="vision.", |
| 352 | ) |
| 353 | |
| 354 | def get_mm_lora_token_counts( |
| 355 | self, |
| 356 | *, |
| 357 | modality: str, |
| 358 | mm_kwargs: MultiModalKwargsItem | None, |
| 359 | num_mm_embeds: int, |
| 360 | ) -> tuple[int, int | None]: |
| 361 | if modality != "image": |
| 362 | raise ValueError(f"Unsupported modality: {modality!r}") |
| 363 | |
| 364 | patches = mm_kwargs.get("patches") if mm_kwargs else None |
| 365 | if patches is not None and isinstance(patches.data, torch.Tensor): |
| 366 | tower_tokens = patches.data.shape[0] |
| 367 | else: |
| 368 | tower_tokens = num_mm_embeds * self.config.vision_downsample_ratio**2 |
| 369 | |
| 370 | return tower_tokens, num_mm_embeds |