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
Migrate Qwen inputs to TensorSchema (#23473)
Signed-off-by: Benji Beck <benjibeck@meta.com>
Benji Beck committed
Aug 27, 2025 at 19:43 UTC
a69693e38f27f12e5a5d05b6792e590b520ca27b
1 file changed
+25
-26
vllm/model_executor/models/qwen_vl.py
+25
-26
@@ -11,7 +11,7 @@ import math
11
import unicodedata
12
from collections.abc import Collection, Mapping, Sequence, Set
13
from functools import lru_cache, partial
14
-from typing import Callable, Literal, Optional, TypedDict, Union
14
+from typing import Annotated, Callable, Literal, Optional, Union
15
16
import regex as re
17
import torch
@@ -40,6 +40,7 @@ from vllm.multimodal.processing import (BaseMultiModalProcessor,
40
PromptUpdate, PromptUpdateDetails)
41
from vllm.multimodal.profiling import BaseDummyInputsBuilder
42
from vllm.sequence import IntermediateTensors
43
+from vllm.utils.tensor_schema import TensorSchema, TensorShape
44
45
from .interfaces import (MultiModalEmbeddings, SupportsLoRA,
46
SupportsMultiModal, SupportsPP)
@@ -47,26 +48,34 @@ from .qwen import QWenBaseModel, QWenModel
48
from .utils import flatten_bn, merge_multimodal_embeddings
49
50
50
-class QwenImagePixelInputs(TypedDict):
51
- type: Literal["pixel_values"]
52
- data: torch.Tensor
51
+class QwenImagePixelInputs(TensorSchema):
52
"""
54
- Shape: `(batch_size * num_images, 3, image_size, image_size)`
55
-
53
+ Dimensions:
54
+ - bn: Batch size * number of images
55
+ - c: Number of channels (3)
56
+ - h: Height
57
+ - w: Width
58
+
59
Note that image_size is the value in the vision config to which we resize
60
the image to in the normalization transform. Currently multi-image support
61
can only be leveraged by passing image embeddings directly.
62
"""
63
+ type: Literal["pixel_values"] = "pixel_values"
64
+ data: Annotated[torch.Tensor, TensorShape("bn", 3, "h", "w")]
65
66
62
-class QwenImageEmbeddingInputs(TypedDict):
63
- type: Literal["image_embeds"]
64
- data: torch.Tensor
65
- """Shape: `(batch_size * num_images, 256, hidden_size)`
66
-
67
+class QwenImageEmbeddingInputs(TensorSchema):
68
+ """
69
+ Dimensions:
70
+ - bn: Batch size * number of images
71
+ - ifs: Image feature size (256)
72
+ - hs: Hidden size
73
+
74
`hidden_size` must match the hidden size of the language model backbone
75
and is stored in the visual config of the model if we have one.
76
"""
77
+ type: Literal["image_embeds"] = "image_embeds"
78
+ data: Annotated[torch.Tensor, TensorShape("bn", 256, "hs")]
79
80
81
QwenImageInputs = Union[QwenImagePixelInputs, QwenImageEmbeddingInputs]
@@ -697,19 +706,6 @@ class QwenVLForConditionalGeneration(QWenBaseModel, SupportsPP, SupportsLoRA,
706
707
self.transformer: QwenVLModel
708
700
- def _validate_pixel_values(self, data: torch.Tensor) -> torch.Tensor:
701
- h = w = self.config.visual["image_size"]
702
- expected_dims = (3, h, w)
703
- actual_dims = tuple(data.shape[1:])
704
-
705
- if actual_dims != expected_dims:
706
- expected_expr = ("batch_size", *map(str, expected_dims))
707
- raise ValueError(
708
- f"The expected shape of pixel values is {expected_expr}. "
709
- f"You supplied {tuple(data.shape)}.")
710
-
711
- return data
712
-
709
def _parse_and_validate_image_input(
710
self, **kwargs: object) -> Optional[QwenImageInputs]:
711
pixel_values = kwargs.pop("pixel_values", None)
@@ -720,10 +716,13 @@ class QwenVLForConditionalGeneration(QWenBaseModel, SupportsPP, SupportsLoRA,
716
raise ValueError("Incorrect type of pixel values. "
717
f"Got type: {type(pixel_values)}")
718
719
+ expected_h = expected_w = self.config.visual["image_size"]
720
+ resolve_bindings = {"h": expected_h, "w": expected_w}
721
+
722
return QwenImagePixelInputs(
723
type="pixel_values",
725
- data=self._validate_pixel_values(
726
- flatten_bn(pixel_values, concat=True)),
724
+ data=flatten_bn(pixel_values, concat=True),
725
+ resolve_bindings=resolve_bindings,
726
)
727
728
if image_embeds is not None: