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 LlavaNextImageInputs to TensorSchema (#21774)
Signed-off-by: Benji Beck <benjibeck@meta.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Benji Beck committed
Aug 10, 2025 at 09:05 UTC
b4e2916721463b43f3b06ccc980050dfb37b615a
2 files changed
+34
-63
vllm/model_executor/models/llava_next.py
+31
-63
@@ -3,7 +3,7 @@
3
4
from abc import abstractmethod
5
from collections.abc import Iterable, Mapping
6
-from typing import (Final, Literal, Optional, Protocol, TypedDict, TypeVar,
6
+from typing import (Annotated, Final, Literal, Optional, Protocol, TypeVar,
7
Union)
8
9
import torch
@@ -11,7 +11,6 @@ import torch.nn as nn
11
from transformers import BatchFeature, LlavaNextConfig, LlavaNextProcessor
12
from transformers.models.llava_next.modeling_llava_next import (
13
get_anyres_image_grid_shape, unpad_image)
14
-from typing_extensions import NotRequired
14
15
from vllm.config import VllmConfig
16
from vllm.model_executor.sampling_metadata import SamplingMetadata
@@ -19,6 +18,7 @@ from vllm.multimodal import MULTIMODAL_REGISTRY
18
from vllm.multimodal.inputs import MultiModalFieldConfig
19
from vllm.multimodal.parse import ImageSize
20
from vllm.sequence import IntermediateTensors
21
+from vllm.utils.tensor_schema import TensorSchema, TensorShape
22
23
from .clip import CLIPVisionModel
24
from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP
@@ -30,32 +30,36 @@ from .utils import (AutoWeightsLoader, WeightsMapper, embed_multimodal,
30
flatten_bn, init_vllm_registered_model, maybe_prefix)
31
32
33
-class LlavaNextImagePixelInputs(TypedDict):
34
- type: Literal["pixel_values"]
35
- pixel_values: Union[torch.Tensor, list[torch.Tensor]]
33
+class LlavaNextImagePixelInputs(TensorSchema):
34
"""
37
- Shape:
38
- `(batch_size * num_images, 1 + num_patches, num_channels, height, width)`
39
-
35
+ Dimensions:
36
+ - bn: Batch size * number of images
37
+ - np: Number of patches + 1
38
+ - c: Number of channels (3)
39
+ - h: Height
40
+ - w: Width
41
+
42
Note that `num_patches` may be different per batch and image,
43
in which case the data is passed as a list instead of a batched tensor.
44
"""
45
+ type: Literal["pixel_values"] = "pixel_values"
46
+ pixel_values: Annotated[
47
+ Union[torch.Tensor, list[torch.Tensor]],
48
+ TensorShape("bn", "np", 3, "h", "w", dynamic_dims={"np"})]
49
44
- image_sizes: NotRequired[torch.Tensor]
45
- """
46
- Shape: `(batch_size * num_images, 2)`
47
-
48
- This should be in `(height, width)` format.
49
- """
50
-
50
+ image_sizes: Annotated[Optional[torch.Tensor], TensorShape("bn", 2)]
51
+ # This should be in `(height, width)` format.
52
52
-class LlavaNextImageEmbeddingInputs(TypedDict):
53
- type: Literal["image_embeds"]
54
- data: torch.Tensor
55
- """Shape: `(batch_size * num_images, image_feature_size, hidden_size)`
53
57
- `hidden_size` must match the hidden size of language model backbone.
54
+class LlavaNextImageEmbeddingInputs(TensorSchema):
55
+ """
56
+ Dimensions:
57
+ - bn: Batch size * number of images
58
+ - ifs: Image feature size
59
+ - hs: Hidden size (must match language model backbone)
60
"""
61
+ type: Literal["image_embeds"] = "image_embeds"
62
+ data: Annotated[torch.Tensor, TensorShape("bn", "ifs", "hs")]
63
64
65
LlavaNextImageInputs = Union[LlavaNextImagePixelInputs,
@@ -269,44 +273,6 @@ class LlavaNextForConditionalGeneration(nn.Module, SupportsMultiModal,
273
self.make_empty_intermediate_tensors = (
274
self.language_model.make_empty_intermediate_tensors)
275
272
- def _validate_image_sizes(self, data: torch.Tensor) -> torch.Tensor:
273
- expected_dims = (2, )
274
-
275
- def _validate_shape(d: torch.Tensor):
276
- actual_dims = tuple(d.shape)
277
-
278
- if actual_dims != expected_dims:
279
- expected_expr = str(expected_dims)
280
- raise ValueError(
281
- f"The expected shape of image sizes per image per batch "
282
- f"is {expected_expr}. You supplied {tuple(d.shape)}.")
283
-
284
- for d in data:
285
- _validate_shape(d)
286
-
287
- return data
288
-
289
- def _validate_pixel_values(
290
- self, data: Union[torch.Tensor, list[torch.Tensor]]
291
- ) -> Union[torch.Tensor, list[torch.Tensor]]:
292
-
293
- h = w = self.config.vision_config.image_size
294
- expected_dims = (3, h, w)
295
-
296
- def _validate_shape(d: torch.Tensor):
297
- actual_dims = tuple(d.shape[1:])
298
-
299
- if actual_dims != expected_dims:
300
- expected_expr = ("num_patches", *map(str, expected_dims))
301
- raise ValueError(
302
- "The expected shape of pixel values per image per batch "
303
- f"is {expected_expr}. You supplied {tuple(d.shape)}.")
304
-
305
- for d in data:
306
- _validate_shape(d)
307
-
308
- return data
309
-
276
def _parse_and_validate_image_input(
277
self, **kwargs: object) -> Optional[LlavaNextImageInputs]:
278
pixel_values = kwargs.pop("pixel_values", None)
@@ -325,13 +291,15 @@ class LlavaNextForConditionalGeneration(nn.Module, SupportsMultiModal,
291
raise ValueError("Incorrect type of image sizes. "
292
f"Got type: {type(image_sizes)}")
293
294
+ expected_h = expected_w = self.config.vision_config.image_size
295
return LlavaNextImagePixelInputs(
296
type="pixel_values",
330
- pixel_values=self._validate_pixel_values(
331
- flatten_bn(pixel_values)),
332
- image_sizes=self._validate_image_sizes(
333
- flatten_bn(image_sizes, concat=True)),
334
- )
297
+ pixel_values=flatten_bn(pixel_values),
298
+ image_sizes=flatten_bn(image_sizes, concat=True),
299
+ resolve_bindings={
300
+ "h": expected_h,
301
+ "w": expected_w,
302
+ })
303
304
if image_embeds is not None:
305
if not isinstance(image_embeds, torch.Tensor):
vllm/utils/tensor_schema.py
+3
@@ -60,6 +60,9 @@ class TensorSchema:
60
def __getitem__(self, item) -> Any:
61
return getattr(self, item)
62
63
+ def get(self, item, default=None) -> Any:
64
+ return getattr(self, item, default)
65
+
66
def _match_shape_with_dynamic(self, actual: tuple[int, ...],
67
reference: tuple[int, ...],
68
expected_shape: tuple[Union[int, str], ...],