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 Paligemma inputs to TensorSchema (#23470)
Signed-off-by: Benji Beck <benjibeck@meta.com>
Benji Beck committed
Aug 23, 2025 at 21:56 UTC
a75277285ba6fa178c9cb9185fec7ec5943fff6b
1 file changed
+27
-36
vllm/model_executor/models/paligemma.py
+27
-36
@@ -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, Sequence
4
-from typing import Literal, Optional, TypedDict, Union
4
+from typing import Annotated, Literal, Optional, Union
5
6
import torch
7
from torch import nn
@@ -21,6 +21,7 @@ from vllm.multimodal.processing import (BaseMultiModalProcessor,
21
PromptUpdateDetails)
22
from vllm.multimodal.profiling import BaseDummyInputsBuilder
23
from vllm.sequence import IntermediateTensors
24
+from vllm.utils.tensor_schema import TensorSchema, TensorShape
25
26
from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP
27
from .siglip import SiglipVisionModel
@@ -32,19 +33,27 @@ from .vision import get_vision_encoder_info
33
logger = init_logger(__name__)
34
35
35
-class PaliGemmaImagePixelInputs(TypedDict):
36
- type: Literal["pixel_values"]
37
- data: torch.Tensor
38
- """Shape: `(batch_size * num_images, num_channels, height, width)`"""
39
-
36
+class PaliGemmaImagePixelInputs(TensorSchema):
37
+ """
38
+ Dimensions:
39
+ - bn: Batch size * number of images
40
+ - c: Number of channels (3)
41
+ - h: Height
42
+ - w: Width
43
+ """
44
+ type: Literal["pixel_values"] = "pixel_values"
45
+ data: Annotated[torch.Tensor, TensorShape("bn", 3, "h", "w")]
46
41
-class PaliGemmaImageEmbeddingInputs(TypedDict):
42
- type: Literal["image_embeds"]
43
- data: torch.Tensor
44
- """Shape: `(batch_size * num_images, image_feature_size, hidden_size)`
47
46
- `hidden_size` must match the hidden size of language model backbone.
48
+class PaliGemmaImageEmbeddingInputs(TensorSchema):
49
"""
50
+ Dimensions:
51
+ - bn: Batch size * number of images
52
+ - ifs: Image feature size
53
+ - hs: Hidden size (must match language model backbone)
54
+ """
55
+ type: Literal["image_embeds"] = "image_embeds"
56
+ data: Annotated[torch.Tensor, TensorShape("bn", "ifs", "hs")]
57
58
59
PaliGemmaImageInputs = Union[PaliGemmaImagePixelInputs,
@@ -279,19 +288,6 @@ class PaliGemmaForConditionalGeneration(nn.Module, SupportsMultiModal,
288
self.make_empty_intermediate_tensors = (
289
self.language_model.make_empty_intermediate_tensors)
290
282
- def _validate_pixel_values(self, data: torch.Tensor) -> torch.Tensor:
283
- h = w = self.config.vision_config.image_size
284
- expected_dims = (3, h, w)
285
- actual_dims = tuple(data.shape[1:])
286
-
287
- if actual_dims != expected_dims:
288
- expected_expr = ("batch_size", *map(str, expected_dims))
289
- raise ValueError(
290
- f"The expected shape of pixel values is {expected_expr}. "
291
- f"You supplied {tuple(data.shape)}.")
292
-
293
- return data
294
-
291
def _parse_and_validate_image_input(
292
self, **kwargs: object) -> Optional[PaliGemmaImageInputs]:
293
pixel_values = kwargs.pop("pixel_values", None)
@@ -301,22 +297,17 @@ class PaliGemmaForConditionalGeneration(nn.Module, SupportsMultiModal,
297
return None
298
299
if pixel_values is not None:
304
- if not isinstance(pixel_values, (torch.Tensor, list)):
305
- raise ValueError("Incorrect type of pixel values. "
306
- f"Got type: {type(pixel_values)}")
307
-
300
pixel_values = flatten_bn(pixel_values, concat=True)
301
310
- return PaliGemmaImagePixelInputs(
311
- type="pixel_values",
312
- data=self._validate_pixel_values(pixel_values),
313
- )
302
+ h = w = self.config.vision_config.image_size
303
+ return PaliGemmaImagePixelInputs(type="pixel_values",
304
+ data=pixel_values,
305
+ resolve_bindings={
306
+ "h": h,
307
+ "w": w
308
+ })
309
310
if image_embeds is not None:
316
- if not isinstance(image_embeds, torch.Tensor):
317
- raise ValueError("Incorrect type of image embeddings. "
318
- f"Got type: {type(image_embeds)}")
319
-
311
image_embeds = flatten_bn(image_embeds, concat=True)
312
313
return PaliGemmaImageEmbeddingInputs(