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 tarsier inputs to TensorSchema (#23500)
Signed-off-by: Benji Beck <benjibeck@meta.com>
Benji Beck committed
Aug 24, 2025 at 21:42 UTC
99f8094400b7177d7ac758083905abde18ddfa64
1 file changed
+23
-21
vllm/model_executor/models/tarsier.py
+23
-21
@@ -3,7 +3,7 @@
3
4
import math
5
from collections.abc import Iterable, Mapping, Sequence
6
-from typing import (Final, Literal, Optional, Protocol, TypedDict, TypeVar,
6
+from typing import (Annotated, Final, Literal, Optional, Protocol, TypeVar,
7
Union, cast)
8
9
import torch
@@ -34,6 +34,7 @@ from vllm.multimodal.processing import (BaseMultiModalProcessor,
34
from vllm.multimodal.profiling import BaseDummyInputsBuilder
35
from vllm.sequence import IntermediateTensors
36
from vllm.utils.jsontree import json_map_leaves
37
+from vllm.utils.tensor_schema import TensorSchema, TensorShape
38
39
from .clip import CLIPVisionModel
40
from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP
@@ -43,14 +44,28 @@ from .utils import (AutoWeightsLoader, flatten_bn, init_vllm_registered_model,
44
from .vision import VisionEncoderInfo, get_vision_encoder_info
45
46
46
-class TarsierImagePixelInputs(TypedDict):
47
- type: Literal["pixel_values"]
48
- pixel_values: torch.Tensor
47
+class TarsierImagePixelInputs(TensorSchema):
48
+ """
49
+ Dimensions:
50
+ - bn: Batch size * number of images
51
+ - c: Number of channels (3)
52
+ - h: Height
53
+ - w: Width
54
+ """
55
+ type: Literal["pixel_values"] = "pixel_values"
56
+ pixel_values: Annotated[torch.Tensor, TensorShape("bn", 3, "h", "w")]
57
58
51
-class TarsierImageEmbeddingInputs(TypedDict):
52
- type: Literal["image_embeds"]
53
- data: torch.Tensor
59
+class TarsierImageEmbeddingInputs(TensorSchema):
60
+ """
61
+ Dimensions:
62
+ - bn: Batch size * number of images
63
+ - ifs: Image feature size
64
+ - hs: Hidden size (must match the hidden size of language model
65
+ backbone)
66
+ """
67
+ type: Literal["image_embeds"] = "image_embeds"
68
+ data: Annotated[torch.Tensor, TensorShape("bn", "ifs", "hs")]
69
70
71
TarsierImageInputs = Union[TarsierImagePixelInputs,
@@ -432,18 +447,6 @@ class TarsierForConditionalGeneration(nn.Module, SupportsMultiModal,
447
self.make_empty_intermediate_tensors = (
448
self.language_model.make_empty_intermediate_tensors)
449
435
- def _validate_pixel_values(self, data: torch.Tensor) -> torch.Tensor:
436
- h = w = self.config.vision_config.image_size
437
- expected_dims = (3, h, w) # Assuming 3 channels
438
- actual_dims = tuple(data.shape[1:])
439
-
440
- if actual_dims != expected_dims:
441
- expected_expr = ("batch_size", *map(str, expected_dims))
442
- raise ValueError(
443
- f"The expected shape of pixel values is {expected_expr}. "
444
- f"You supplied {tuple(data.shape)}.")
445
- return data
446
-
450
def _parse_and_validate_image_input(
451
self, **kwargs: object) -> Optional[TarsierImageInputs]:
452
pixel_values = kwargs.pop("pixel_values", None)
@@ -459,8 +462,7 @@ class TarsierForConditionalGeneration(nn.Module, SupportsMultiModal,
462
463
return TarsierImagePixelInputs(
464
type="pixel_values",
462
- pixel_values=self._validate_pixel_values(
463
- flatten_bn(pixel_values, concat=True)),
465
+ pixel_values=flatten_bn(pixel_values, concat=True),
466
)
467
468
if image_embeds is not None: