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 LlavaImageInputs to TensorSchema (#21770)
Signed-off-by: Benji Beck <benjibeck@meta.com>
Benji Beck committed
Aug 10, 2025 at 19:29 UTC
06da44f0cbf84da771a2a1e336e06432a09875c8
1 file changed
+35
-32
vllm/model_executor/models/llava.py
+35
-32
@@ -3,7 +3,7 @@
3
4
from abc import abstractmethod
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
@@ -33,6 +33,7 @@ from vllm.multimodal.processing import (BaseMultiModalProcessor,
33
PromptUpdateDetails)
34
from vllm.multimodal.profiling import BaseDummyInputsBuilder
35
from vllm.sequence import IntermediateTensors
36
+from vllm.utils.tensor_schema import TensorSchema, TensorShape
37
38
from .clip import CLIPVisionModel
39
from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP
@@ -44,35 +45,46 @@ from .utils import (AutoWeightsLoader, WeightsMapper, flatten_bn,
45
from .vision import get_vision_encoder_info
46
47
47
-class LlavaImagePixelInputs(TypedDict):
48
- type: Literal["pixel_values"]
49
- pixel_values: torch.Tensor
48
+class LlavaImagePixelInputs(TensorSchema):
49
"""
51
- Shape: `(batch_size * num_images, num_channels, height, width)`
52
-
50
+ Dimensions:
51
+ - bn: Batch size * number of images
52
+ - c: Number of channels (3)
53
+ - h: Height
54
+ - w: Width
55
+
56
Note that `height` or `width` may be different per batch and image,
57
in which case the data is passed as a list instead of a batched tensor.
58
"""
59
+ type: Literal["pixel_values"] = "pixel_values"
60
+ pixel_values: Annotated[torch.Tensor, TensorShape("bn", 3, "h", "w")]
61
62
58
-class PixtralHFImagePixelInputs(TypedDict):
59
- type: Literal["pixel_values_pixtral"]
60
- pixel_values: Union[torch.Tensor, list[torch.Tensor]]
63
+class PixtralHFImagePixelInputs(TensorSchema):
64
"""
62
- Shape: `(batch_size * num_images, num_channels, height, width)`
63
-
65
+ Dimensions:
66
+ - bn: Batch size * number of images
67
+ - c: Number of channels
68
+ - h: Height
69
+ - w: Width
70
+
71
Note that `height` or `width` may be different per batch and image,
72
in which case the data is passed as a list instead of a batched tensor.
73
"""
74
+ type: Literal["pixel_values_pixtral"] = "pixel_values_pixtral"
75
+ pixel_values: Annotated[Union[torch.Tensor, list[torch.Tensor]],
76
+ TensorShape("bn", "c", "h", "w")]
77
78
69
-class LlavaImageEmbeddingInputs(TypedDict):
70
- type: Literal["image_embeds"]
71
- data: torch.Tensor
72
- """Shape: `(batch_size * num_images, image_feature_size, hidden_size)`
73
-
74
- `hidden_size` must match the hidden size of language model backbone.
79
+class LlavaImageEmbeddingInputs(TensorSchema):
80
"""
81
+ Dimensions:
82
+ - bn: Batch size * number of images
83
+ - ifs: Image feature size
84
+ - hs: Hidden size (must match language model backbone)
85
+ """
86
+ type: Literal["image_embeds"] = "image_embeds"
87
+ data: Annotated[torch.Tensor, TensorShape("bn", "ifs", "hs")]
88
89
90
LlavaImageInputs = Union[LlavaImagePixelInputs, PixtralHFImagePixelInputs,
@@ -547,19 +559,6 @@ class LlavaForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP):
559
self.make_empty_intermediate_tensors = (
560
self.language_model.make_empty_intermediate_tensors)
561
550
- def _validate_pixel_values(self, data: torch.Tensor) -> torch.Tensor:
551
- h = w = self.config.vision_config.image_size
552
- expected_dims = (3, h, w)
553
- actual_dims = tuple(data.shape[1:])
554
-
555
- if actual_dims != expected_dims:
556
- expected_expr = ("batch_size", *map(str, expected_dims))
557
- raise ValueError(
558
- f"The expected shape of pixel values is {expected_expr}. "
559
- f"You supplied {tuple(data.shape)}.")
560
-
561
- return data
562
-
562
def _parse_and_validate_image_input(
563
self, **kwargs: object) -> Optional[LlavaImageInputs]:
564
pixel_values = kwargs.pop("pixel_values", None)
@@ -579,10 +578,14 @@ class LlavaForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP):
578
pixel_values=flatten_bn(pixel_values),
579
)
580
581
+ expected_h = expected_w = self.config.vision_config.image_size
582
return LlavaImagePixelInputs(
583
type="pixel_values",
584
- pixel_values=self._validate_pixel_values(
585
- flatten_bn(pixel_values, concat=True)),
584
+ pixel_values=flatten_bn(pixel_values, concat=True),
585
+ resolve_bindings={
586
+ "h": expected_h,
587
+ "w": expected_w
588
+ },
589
)
590
591
if image_embeds is not None: