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, 23:45 UTC
Migrate DeepseekVL2ImageInputs to TensorSchema (#21658)
Signed-off-by: Benji Beck <benjibeck@meta.com>
Benji Beck committed
Jul 26, 2025 at 19:34 UTC
0b8caf909581a2fbd1b825c65478878c3ce3d0b9
1 file changed
+31
-72
vllm/model_executor/models/deepseek_vl2.py
+31
-72
@@ -5,7 +5,7 @@
5
"""Inference-only Deepseek-VL2 model compatible with HuggingFace weights."""
6
import math
7
from collections.abc import Iterable, Mapping, Sequence
8
-from typing import Literal, Optional, TypedDict, Union
8
+from typing import Annotated, Literal, Optional, Union
9
10
import torch
11
import torch.nn as nn
@@ -36,6 +36,7 @@ from vllm.transformers_utils.processors.deepseek_vl2 import (
36
DeepseekVLV2Processor)
37
from vllm.transformers_utils.tokenizer import cached_tokenizer_from_config
38
from vllm.utils import is_list_of
39
+from vllm.utils.tensor_schema import TensorSchema, TensorShape
40
41
from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP
42
from .utils import (AutoWeightsLoader, WeightsMapper, flatten_bn,
@@ -46,25 +47,30 @@ from .utils import (AutoWeightsLoader, WeightsMapper, flatten_bn,
47
_IMAGE_TOKEN = "<image>"
48
49
49
-class DeepseekVL2ImagePixelInputs(TypedDict):
50
- type: Literal["pixel_values"]
51
- data: Union[torch.Tensor, list[torch.Tensor]]
52
- """
53
- Shape: `(batch_size * num_images, num_channels, height, width)`
54
- """
55
- images_spatial_crop: torch.Tensor
50
+class DeepseekVL2ImagePixelInputs(TensorSchema):
51
"""
57
- Shape: `(batch_size * num_images, 2)`
52
+ Dimensions:
53
+ - bn: Batch size * number of images
54
+ - c: Number of channels (3)
55
+ - h: Height of each image
56
+ - w: Width of each image
57
"""
58
+ type: Literal["pixel_values"]
59
+ data: Annotated[Union[torch.Tensor, list[torch.Tensor]],
60
+ TensorShape("bn", 3, "h", "w")]
61
+ images_spatial_crop: Annotated[torch.Tensor, TensorShape("bn", 2)]
62
63
61
-class DeepseekVL2VImageEmbeddingInputs(TypedDict):
62
- type: Literal["image_embeds"]
63
- data: Union[torch.Tensor, list[torch.Tensor]]
64
- """Shape: `(batch_size * num_images, image_feature_size, hidden_size)`
65
-
66
- `hidden_size` must match the hidden size of language model backbone.
64
+class DeepseekVL2VImageEmbeddingInputs(TensorSchema):
65
"""
66
+ Dimensions:
67
+ - bn: Batch size * number of images
68
+ - f: Image feature size
69
+ - h: Hidden size (must match language model backbone)
70
+ """
71
+ type: Literal["image_embeds"]
72
+ data: Annotated[Union[torch.Tensor, list[torch.Tensor]],
73
+ TensorShape("bn", "f", "h")]
74
75
76
DeepseekVL2ImageInputs = Union[DeepseekVL2ImagePixelInputs,
@@ -439,46 +445,6 @@ class DeepseekVLV2ForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):
445
model = model.to(dtype=torch.get_default_dtype())
446
return model
447
442
- def _validate_pixel_values(
443
- self, data: Union[torch.Tensor, list[torch.Tensor]]
444
- ) -> Union[torch.Tensor, list[torch.Tensor]]:
445
-
446
- h = w = self.vision_config.image_size
447
- expected_dims = (3, h, w)
448
-
449
- def _validate_shape(d: torch.Tensor):
450
- actual_dims = tuple(d.shape[1:])
451
-
452
- if actual_dims != expected_dims:
453
- expected_expr = ("num_patches", *map(str, expected_dims))
454
- raise ValueError(
455
- "The expected shape of pixel values per image per batch "
456
- f"is {expected_expr}. You supplied {tuple(d.shape)}.")
457
-
458
- for d in data:
459
- _validate_shape(d)
460
-
461
- return data
462
-
463
- def _validate_images_spatial_crop(
464
- self, data: Union[torch.Tensor, list[torch.Tensor]]
465
- ) -> Union[torch.Tensor, list[torch.Tensor]]:
466
- expected_dims = 2
467
-
468
- def _validate_shape(d: torch.Tensor):
469
- actual_dims = d.size(-1)
470
-
471
- if actual_dims != expected_dims:
472
- expected_expr = str(expected_dims)
473
- raise ValueError(
474
- f"The expected shape of image sizes per image per batch "
475
- f"is {expected_expr}. You supplied {tuple(d.shape)}.")
476
-
477
- for d in data:
478
- _validate_shape(d)
479
-
480
- return data
481
-
448
def _parse_and_validate_image_input(
449
self, **kwargs: object) -> Optional[DeepseekVL2ImageInputs]:
450
pixel_values = kwargs.pop("pixel_values", None)
@@ -489,25 +455,18 @@ class DeepseekVLV2ForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):
455
return None
456
457
if pixel_values is not None:
492
- if not isinstance(pixel_values, (torch.Tensor, list)):
493
- raise ValueError("Incorrect type of pixel values. "
494
- f"Got type: {type(pixel_values)}")
495
-
496
- if not isinstance(images_spatial_crop, (torch.Tensor, list)):
497
- raise ValueError("Incorrect type of image sizes. "
498
- f"Got type: {type(images_spatial_crop)}")
499
-
500
- return DeepseekVL2ImagePixelInputs(
501
- type="pixel_values",
502
- data=self._validate_pixel_values(flatten_bn(pixel_values)),
503
- images_spatial_crop=self._validate_images_spatial_crop(
504
- flatten_bn(images_spatial_crop, concat=True)))
458
+ expected_h = expected_w = self.vision_config.image_size
459
+ return DeepseekVL2ImagePixelInputs(type="pixel_values",
460
+ data=flatten_bn(pixel_values),
461
+ images_spatial_crop=flatten_bn(
462
+ images_spatial_crop,
463
+ concat=True),
464
+ resolve_bindings={
465
+ "h": expected_h,
466
+ "w": expected_w,
467
+ })
468
469
if image_embeds is not None:
507
- if not isinstance(image_embeds, (torch.Tensor, list)):
508
- raise ValueError("Incorrect type of image embeddings. "
509
- f"Got type: {type(image_embeds)}")
510
-
470
return DeepseekVL2VImageEmbeddingInputs(
471
type="image_embeds",
472
data=flatten_bn(image_embeds),