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 14, 05:30 UTC
[Model] Use `merge_by_field_config` for MM models (D-F) (#26076)
Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>
Cyrus Leung committed
Oct 2, 2025 at 23:17 UTC
cc253b73d3d6bdbd46b336065b148d344c51c20d
4 files changed
+99
-177
vllm/model_executor/models/deepseek_vl2.py
+25
-33
@@ -20,8 +20,7 @@ from vllm.model_executor.model_loader.utils import set_default_torch_dtype
20
from vllm.model_executor.models.transformers import replace_linear_class
21
from vllm.multimodal import MULTIMODAL_REGISTRY
22
from vllm.multimodal.inputs import (MultiModalDataDict, MultiModalFieldConfig,
23
- MultiModalKwargsItems, MultiModalUUIDDict,
24
- NestedTensors)
23
+ MultiModalKwargsItems, MultiModalUUIDDict)
24
from vllm.multimodal.parse import (ImageEmbeddingItems, ImageProcessorItems,
25
ImageSize, MultiModalDataItems)
26
from vllm.multimodal.processing import (BaseMultiModalProcessor,
@@ -40,7 +39,7 @@ from vllm.utils import is_list_of
39
from vllm.utils.tensor_schema import TensorSchema, TensorShape
40
41
from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP
43
-from .utils import (AutoWeightsLoader, WeightsMapper, flatten_bn,
42
+from .utils import (AutoWeightsLoader, WeightsMapper,
43
init_vllm_registered_model, maybe_prefix)
44
45
# The image token id may be various
@@ -50,15 +49,15 @@ _IMAGE_TOKEN = "<image>"
49
class DeepseekVL2ImagePixelInputs(TensorSchema):
50
"""
51
Dimensions:
53
- - bn: Batch size * number of images
52
+ - bnp: Batch size * number of images * number of patches
53
- p: Number of patches
54
- c: Number of channels (3)
55
- h: Height of each image
56
- w: Width of each image
57
"""
58
type: Literal["pixel_values"]
60
- data: Annotated[Union[torch.Tensor, list[torch.Tensor]],
61
- TensorShape("bn", "p", 3, "h", "w", dynamic_dims={"p"})]
59
+ data: Annotated[torch.Tensor,
60
+ TensorShape("bnp", 3, "h", "w", dynamic_dims={"bnp"})]
61
images_spatial_crop: Annotated[torch.Tensor, TensorShape("bn", 2)]
62
63
@@ -228,12 +227,8 @@ class DeepseekVL2MultiModalProcessor(
227
tok_kwargs=tok_kwargs,
228
)
229
231
- pixel_values = processed_outputs["pixel_values"]
232
- # split pixel values into patches corresponding to each image
233
- images_spatial_crop = processed_outputs["images_spatial_crop"]
234
- patches_per_image = [x.prod().item() + 1 for x in images_spatial_crop]
235
- pixel_values = pixel_values.split(patches_per_image)
236
- processed_outputs["pixel_values"] = pixel_values
230
+ processed_outputs["num_patches"] = (
231
+ processed_outputs["images_spatial_crop"].prod(-1) + 1)
232
233
return processed_outputs
234
@@ -242,8 +237,11 @@ class DeepseekVL2MultiModalProcessor(
237
hf_inputs: BatchFeature,
238
hf_processor_mm_kwargs: Mapping[str, object],
239
) -> Mapping[str, MultiModalFieldConfig]:
240
+ num_patches = hf_inputs.get("num_patches", torch.empty(0))
241
+
242
return dict(
246
- pixel_values=MultiModalFieldConfig.batched("image"),
243
+ pixel_values=MultiModalFieldConfig.flat_from_sizes(
244
+ "image", num_patches),
245
images_spatial_crop=MultiModalFieldConfig.batched("image"),
246
image_embeds=MultiModalFieldConfig.batched("image"),
247
)
@@ -318,6 +316,7 @@ class DeepseekVL2MultiModalProcessor(
316
info=DeepseekVL2ProcessingInfo,
317
dummy_inputs=DeepseekVL2DummyInputsBuilder)
318
class DeepseekVLV2ForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):
319
+ merge_by_field_config = True
320
321
hf_to_vllm_mapper = WeightsMapper(orig_to_new_prefix={
322
"language.": "language_model.",
@@ -460,37 +459,30 @@ class DeepseekVLV2ForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):
459
460
if pixel_values is not None:
461
expected_h = expected_w = self.vision_config.image_size
463
- return DeepseekVL2ImagePixelInputs(type="pixel_values",
464
- data=flatten_bn(pixel_values),
465
- images_spatial_crop=flatten_bn(
466
- images_spatial_crop,
467
- concat=True),
468
- resolve_bindings={
469
- "h": expected_h,
470
- "w": expected_w,
471
- })
462
+ return DeepseekVL2ImagePixelInputs(
463
+ type="pixel_values",
464
+ data=pixel_values,
465
+ images_spatial_crop=images_spatial_crop,
466
+ resolve_bindings={
467
+ "h": expected_h,
468
+ "w": expected_w,
469
+ })
470
471
if image_embeds is not None:
472
return DeepseekVL2VImageEmbeddingInputs(
473
type="image_embeds",
476
- data=flatten_bn(image_embeds),
474
+ data=image_embeds,
475
)
476
477
raise AssertionError("This line should be unreachable.")
478
479
def _pixel_values_to_embedding(
480
self,
483
- pixel_values: NestedTensors,
481
+ pixel_values: torch.Tensor,
482
images_spatial_crop: torch.Tensor,
485
- ) -> NestedTensors:
486
- # Pixel_values: n_image * batch_size * [patch_per_img, 3, height, width]
487
- total_tiles = [x for x in pixel_values]
488
-
489
- # [batch_all_tiles, 3, height, width]
490
- total_tiles = torch.cat(total_tiles, dim=0)
491
-
483
+ ) -> list[torch.Tensor]:
484
# [batch_all_tiles, vit_seq_len, c]
493
- images_feature = self.vision.forward_features(total_tiles)
485
+ images_feature = self.vision.forward_features(pixel_values)
486
487
# [batch_all_tiles, hw, D]
488
images_embeds = self.projector(images_feature)
@@ -573,7 +565,7 @@ class DeepseekVLV2ForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):
565
return vision_embeddings
566
567
def _process_image_input(
576
- self, image_input: DeepseekVL2ImageInputs) -> torch.Tensor:
568
+ self, image_input: DeepseekVL2ImageInputs) -> list[torch.Tensor]:
569
if image_input["type"] == "image_embeds":
570
image_data = image_input["data"]
571
if is_list_of(image_data, torch.Tensor):
vllm/model_executor/models/dots_ocr.py
+24
-51
@@ -1,7 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
from collections.abc import Iterable, Mapping
4
-from typing import Literal, Optional, TypedDict, Union
4
+from typing import Annotated, Literal, Optional, Union
5
6
import torch
7
import torch.nn as nn
@@ -42,34 +42,38 @@ from vllm.platforms import _Backend
42
from vllm.sequence import IntermediateTensors
43
from vllm.transformers_utils.configs.dotsocr import (DotsOCRConfig,
44
DotsVisionConfig)
45
+from vllm.utils.tensor_schema import TensorSchema, TensorShape
46
47
from .vision import run_dp_sharded_mrope_vision_model
48
49
IMAGE_TOKEN = "<|imgpad|>"
50
51
51
-class DotsOCRImagePixelInputs(TypedDict):
52
- type: Literal["pixel_values", "image_grid_thw"]
52
+class DotsOCRImagePixelInputs(TensorSchema):
53
+ """
54
+ Dimensions:
55
+ - np: The total number of patches over each image over each prompt in
56
+ the batch
57
+ - ni: Number of images
58
+ - cps: Number of channels * patch_size * patch_size
59
+ """
60
+ type: Literal["pixel_values"]
61
54
- pixel_values: torch.Tensor
55
- image_grid_thw: torch.Tensor
62
+ pixel_values: Annotated[torch.Tensor, TensorShape("np", "cps")]
63
+ image_grid_thw: Annotated[torch.Tensor, TensorShape("ni", 3)]
64
65
58
-class DotsOCRImageEmbeddingInputs(TypedDict):
59
- type: Literal["image_embeds", "image_grid_thw"]
60
- image_embeds: torch.Tensor
61
- """Supported types:
62
- - List[`torch.Tensor`]: A list of tensors holding all images' features.
63
- Each tensor holds an image's features.
64
- - `torch.Tensor`: A tensor holding all images' features
65
- (concatenation of all images' feature tensors).
66
- Tensor shape: `(num_image_features, hidden_size)`
67
- - `num_image_features` varies based on
68
- the number and resolution of the images.
69
- - `hidden_size` must match the hidden size of language model backbone.
66
+class DotsOCRImageEmbeddingInputs(TensorSchema):
67
"""
68
+ Dimensions:
69
+ - nf: Number of image features
70
+ - hs: Hidden size
71
+ - ni: Number of images
72
+ """
73
+ type: Literal["image_embeds"]
74
72
- image_grid_thw: torch.Tensor
75
+ image_embeds: Annotated[torch.Tensor, TensorShape("nf", "hs")]
76
+ image_grid_thw: Annotated[torch.Tensor, TensorShape("ni", 3)]
77
78
79
DotsOCRImageInputs = Union[DotsOCRImagePixelInputs,
@@ -654,6 +658,8 @@ class DotsVisionTransformer(nn.Module):
658
)
659
class DotsOCRForCausalLM(nn.Module, SupportsMultiModal, SupportsPP,
660
SupportsLoRA):
661
+ merge_by_field_config = True
662
+
663
hf_to_vllm_mapper = WeightsMapper(
664
orig_to_new_substr={
665
".attn.qkv_proj.": ".attn.qkv.",
@@ -709,22 +715,6 @@ class DotsOCRForCausalLM(nn.Module, SupportsMultiModal, SupportsPP,
715
architectures=["Qwen2ForCausalLM"],
716
)
717
712
- def _validate_and_reshape_mm_tensor(self, mm_input: object,
713
- name: str) -> torch.Tensor:
714
- if not isinstance(mm_input, (torch.Tensor, list)):
715
- raise ValueError(f"Incorrect type of {name}. "
716
- f"Got type: {type(mm_input)}")
717
- if isinstance(mm_input, torch.Tensor):
718
- if mm_input.ndim == 2:
719
- return mm_input
720
- if mm_input.ndim != 3:
721
- raise ValueError(f"{name} should be 2D or batched 3D tensor. "
722
- f"Got ndim: {mm_input.ndim} "
723
- f"(shape={mm_input.shape})")
724
- return torch.concat(list(mm_input))
725
- else:
726
- return torch.concat(mm_input)
727
-
718
def _parse_and_validate_image_input(
719
self, **kwargs: object) -> Optional[DotsOCRImageInputs]:
720
pixel_values = kwargs.pop("pixel_values", None)
@@ -735,28 +725,11 @@ class DotsOCRForCausalLM(nn.Module, SupportsMultiModal, SupportsPP,
725
return None
726
727
if pixel_values is not None:
738
- pixel_values = self._validate_and_reshape_mm_tensor(
739
- pixel_values, "image pixel values")
740
- image_grid_thw = self._validate_and_reshape_mm_tensor(
741
- image_grid_thw, "image grid_thw")
742
-
743
- if not isinstance(pixel_values, (torch.Tensor, list)):
744
- raise ValueError("Incorrect type of image pixel values. "
745
- f"Got type: {type(pixel_values)}")
746
-
728
return DotsOCRImagePixelInputs(type="pixel_values",
729
pixel_values=pixel_values,
730
image_grid_thw=image_grid_thw)
731
732
if image_embeds is not None:
752
- image_embeds = self._validate_and_reshape_mm_tensor(
753
- image_embeds, "image embeds")
754
- image_grid_thw = self._validate_and_reshape_mm_tensor(
755
- image_grid_thw, "image grid_thw")
756
-
757
- if not isinstance(image_embeds, torch.Tensor):
758
- raise ValueError("Incorrect type of image embeddings. "
759
- f"Got type: {type(image_embeds)}")
733
return DotsOCRImageEmbeddingInputs(type="image_embeds",
734
image_embeds=image_embeds,
735
image_grid_thw=image_grid_thw)
vllm/model_executor/models/ernie45_vl.py
+24
-52
@@ -25,7 +25,7 @@
25
import math
26
from collections.abc import Iterable, Mapping, Sequence
27
from functools import partial
28
-from typing import Any, Callable, Literal, Optional, TypedDict, Union
28
+from typing import Annotated, Any, Callable, Literal, Optional, Union
29
30
import numpy as np
31
import torch
@@ -56,6 +56,7 @@ from vllm.multimodal.processing import (BaseMultiModalProcessor,
56
from vllm.multimodal.profiling import BaseDummyInputsBuilder
57
from vllm.platforms import _Backend, current_platform
58
from vllm.sequence import IntermediateTensors
59
+from vllm.utils.tensor_schema import TensorSchema, TensorShape
60
61
from .ernie45_vl_moe import Ernie4_5_VLMoeForCausalLM
62
from .interfaces import (MultiModalEmbeddings, SupportsLoRA,
@@ -579,38 +580,38 @@ class Ernie4_5_VisionTransformer(nn.Module):
580
# === Vision Inputs === #
581
582
582
-class Ernie4_5_VLImagePixelInputs(TypedDict):
583
- type: Literal["pixel_values"]
584
- pixel_values: torch.Tensor
585
- """Shape:
586
- `(num_patches, num_channels * patch_size * patch_size)`
583
+class Ernie4_5_VLImagePixelInputs(TensorSchema):
584
"""
588
-
589
- grid_thw: torch.Tensor
590
- """Shape: `(num_images, 3)`
591
- This should be in `(grid_t, grid_h, grid_w)` format.
585
+ Dimensions:
586
+ - np: The total number of patches over each image over each prompt in
587
+ the batch
588
+ - ni: Number of images
589
+ - cps: Number of channels * patch_size * patch_size
590
"""
591
+ type: Literal["pixel_values"]
592
+
593
+ pixel_values: Annotated[torch.Tensor, TensorShape("np", "cps")]
594
+ image_grid_thw: Annotated[torch.Tensor, TensorShape("ni", 3)]
595
596
597
Ernie4_5_VLImageInputs = Ernie4_5_VLImagePixelInputs
598
599
598
-class Ernie4_5_VLVideoPixelInputs(TypedDict):
599
- type: Literal["pixel_values_videos"]
600
- pixel_values_videos: torch.Tensor
601
- """Shape:
602
- `(num_patches,
603
- num_channels * temporal_patch_size * patch_size * patch_size)`
600
+class Ernie4_5_VLVideoPixelInputs(TensorSchema):
601
"""
605
-
606
- video_grid_thw: torch.Tensor
607
- """Shape: `(num_videos, 3)`
608
-
609
- This should be in `(grid_t, grid_h, grid_w)` format.
602
+ Dimensions:
603
+ - np: The total number of patches over each image over each prompt in
604
+ the batch
605
+ - ni: Number of images
606
+ - cps: Number of channels * temporal_patch_size * patch_size *
607
+ patch_size
608
"""
609
+ type: Literal["pixel_values_videos"]
610
+ pixel_values_videos: Annotated[torch.Tensor, TensorShape("np", "cps")]
611
+ video_grid_thw: Annotated[torch.Tensor, TensorShape("ni", 3)]
612
613
613
-Ernie4_5_VLVideoInputs = Ernie4_5_VLImagePixelInputs
614
+Ernie4_5_VLVideoInputs = Ernie4_5_VLVideoPixelInputs
615
616
# === Vision Processor === #
617
@@ -1213,6 +1214,7 @@ class Ernie4_5_VLDummyInputsBuilder(
1214
dummy_inputs=Ernie4_5_VLDummyInputsBuilder)
1215
class Ernie4_5_VLMoeForConditionalGeneration(nn.Module, SupportsMultiModal,
1216
SupportsLoRA, SupportsPP):
1217
+ merge_by_field_config = True
1218
1219
packed_modules_mapping = {
1220
"qkv_proj": [
@@ -1325,22 +1327,6 @@ class Ernie4_5_VLMoeForConditionalGeneration(nn.Module, SupportsMultiModal,
1327
def get_language_model(self) -> torch.nn.Module:
1328
return self.language_model
1329
1328
- def _validate_and_reshape_mm_tensor(self, mm_input: object,
1329
- name: str) -> torch.Tensor:
1330
- if not isinstance(mm_input, (torch.Tensor, list)):
1331
- raise ValueError(f"Incorrect type of {name}. "
1332
- f"Got type: {type(mm_input)}")
1333
- if isinstance(mm_input, torch.Tensor):
1334
- if mm_input.ndim == 2:
1335
- return mm_input
1336
- if mm_input.ndim != 3:
1337
- raise ValueError(f"{name} should be 2D or batched 3D tensor. "
1338
- f"Got ndim: {mm_input.ndim} "
1339
- f"(shape={mm_input.shape})")
1340
- return mm_input.reshape(-1, mm_input.shape[-1])
1341
- else:
1342
- return torch.concat(mm_input)
1343
-
1330
def _parse_and_validate_image_input(
1331
self, **kwargs: object) -> Optional[Ernie4_5_VLImageInputs]:
1332
pixel_values = kwargs.pop("pixel_values", None)
@@ -1350,15 +1336,6 @@ class Ernie4_5_VLMoeForConditionalGeneration(nn.Module, SupportsMultiModal,
1336
return None
1337
1338
if pixel_values is not None:
1353
- pixel_values = self._validate_and_reshape_mm_tensor(
1354
- pixel_values, "image pixel values")
1355
- image_grid_thw = self._validate_and_reshape_mm_tensor(
1356
- image_grid_thw, "image grid_thw")
1357
-
1358
- if not isinstance(pixel_values, (torch.Tensor, list)):
1359
- raise ValueError("Incorrect type of image pixel values. "
1360
- f"Got type: {type(pixel_values)}")
1361
-
1339
return Ernie4_5_VLImagePixelInputs(type="pixel_values",
1340
pixel_values=pixel_values,
1341
image_grid_thw=image_grid_thw)
@@ -1372,11 +1349,6 @@ class Ernie4_5_VLMoeForConditionalGeneration(nn.Module, SupportsMultiModal,
1349
return None
1350
1351
if pixel_values_videos is not None:
1375
- pixel_values_videos = self._validate_and_reshape_mm_tensor(
1376
- pixel_values_videos, "video pixel values")
1377
- video_grid_thw = self._validate_and_reshape_mm_tensor(
1378
- video_grid_thw, "video grid_thw")
1379
-
1352
return Ernie4_5_VLVideoPixelInputs(
1353
type="pixel_values_videos",
1354
pixel_values_videos=pixel_values_videos,
vllm/model_executor/models/fuyu.py
+26
-41
@@ -59,17 +59,14 @@ class FuyuImagePatchInputs(TensorSchema):
59
60
type: Literal["image_patches"] = "image_patches"
61
62
- flat_data: Annotated[
63
- torch.Tensor,
64
- TensorShape("bnp", "fn"),
65
- ]
62
+ image_patches_flat: Annotated[torch.Tensor, TensorShape("bnp", "fn")]
63
64
patches_per_image: Annotated[list[int], TensorShape("bn")]
65
"""
66
The number of total patches for each image in the batch.
67
68
This is used to split the embeddings which has the first two dimensions
72
- flattened just like `flat_data`.
69
+ flattened just like `image_patches_flat`.
70
"""
71
72
@@ -174,28 +171,10 @@ class FuyuMultiModalProcessor(BaseMultiModalProcessor[FuyuProcessingInfo]):
171
tok_kwargs=tok_kwargs,
172
)
173
177
- image_patches = processed_outputs.get("image_patches")
178
- if image_patches is not None:
179
- images = mm_data["images"]
180
- assert isinstance(images, list)
181
-
182
- # Original output: (1, num_images, Pn, Px * Py * C)
183
- # New output: (num_images, Pn, Px * Py * C)
184
- # image_patches is a list with shape:
185
- # (1, num_images, Pn, Px * Py * C)
186
- # before Transformers 4.53
187
- if isinstance(image_patches, list):
188
- assert len(image_patches) == 1
189
- assert (isinstance(image_patches[0], torch.Tensor)
190
- and len(image_patches[0]) == len(images))
191
- processed_outputs["image_patches"] = image_patches[0]
192
- # image_patches is a tensor with shape:
193
- # (num_images, Pn, Px * Py * C)
194
- # after Transformers 4.53
195
- elif isinstance(image_patches, torch.Tensor):
196
- assert len(image_patches) == len(images)
197
- else:
198
- raise AssertionError("This line should be unreachable.")
174
+ image_patches = processed_outputs["image_patches"]
175
+ processed_outputs["image_patches"] = flatten_bn(image_patches)
176
+ processed_outputs["patches_per_image"] = torch.tensor(
177
+ [len(p) for p in image_patches])
178
179
return processed_outputs
180
@@ -218,7 +197,13 @@ class FuyuMultiModalProcessor(BaseMultiModalProcessor[FuyuProcessingInfo]):
197
hf_inputs: BatchFeature,
198
hf_processor_mm_kwargs: Mapping[str, object],
199
) -> Mapping[str, MultiModalFieldConfig]:
221
- return dict(image_patches=MultiModalFieldConfig.batched("image"))
200
+ patches_per_image = hf_inputs.get("patches_per_image", torch.empty(0))
201
+
202
+ return dict(
203
+ image_patches=MultiModalFieldConfig.flat_from_sizes(
204
+ "image", patches_per_image),
205
+ patches_per_image=MultiModalFieldConfig.batched("image"),
206
+ )
207
208
def _get_prompt_updates(
209
self,
@@ -263,6 +248,7 @@ class FuyuMultiModalProcessor(BaseMultiModalProcessor[FuyuProcessingInfo]):
248
info=FuyuProcessingInfo,
249
dummy_inputs=FuyuDummyInputsBuilder)
250
class FuyuForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):
251
+ merge_by_field_config = True
252
253
hf_to_vllm_mapper = WeightsMapper(
254
orig_to_new_prefix={
@@ -306,29 +292,28 @@ class FuyuForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):
292
def _parse_and_validate_image_input(
293
self, **kwargs: object) -> Optional[FuyuImagePatchInputs]:
294
image_patches = kwargs.pop("image_patches", None)
309
- if image_patches is not None:
310
- image_patches_flat = flatten_bn(image_patches)
311
- flat_data = flatten_bn(image_patches_flat, concat=True)
312
-
313
- return FuyuImagePatchInputs(
314
- type="image_patches",
315
- flat_data=flat_data,
316
- patches_per_image=[x.size(0) for x in image_patches_flat],
317
- resolve_bindings={"fn": self.image_feature_size},
318
- )
295
+ patches_per_image = kwargs.pop("patches_per_image", None)
296
320
- return None
297
+ if image_patches is None:
298
+ return None
299
+
300
+ return FuyuImagePatchInputs(
301
+ type="image_patches",
302
+ image_patches_flat=image_patches,
303
+ patches_per_image=patches_per_image,
304
+ resolve_bindings={"fn": self.image_feature_size},
305
+ )
306
307
def _process_image_input(
308
self, image_input: FuyuImagePatchInputs) -> MultiModalEmbeddings:
324
- image_patches_flat = image_input["flat_data"]
309
+ image_patches_flat = image_input["image_patches_flat"]
310
patches_per_image = image_input["patches_per_image"]
311
312
assert self.vision_embed_tokens is not None
313
vision_embeddings_flat, _ = self.vision_embed_tokens(
314
image_patches_flat)
315
331
- return vision_embeddings_flat.split(patches_per_image, dim=0)
316
+ return vision_embeddings_flat.split(patches_per_image.tolist(), dim=0)
317
318
def get_language_model(self) -> torch.nn.Module:
319
return self.language_model