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 LlavaOnevisionMultiInputs to TensorSchema (#21844)
Signed-off-by: Benji Beck <benjibeck@meta.com>
Benji Beck committed
Aug 19, 2025 at 10:02 UTC
a70d0bd0a39bfb278b7bda9c82d99df8f628d779
1 file changed
+56
-93
vllm/model_executor/models/llava_onevision.py
+56
-93
@@ -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, Union
6
+from typing import Annotated, Final, Literal, Optional, Protocol, Union
7
8
import torch
9
import torch.nn as nn
@@ -11,7 +11,6 @@ from transformers import (BatchFeature, LlavaOnevisionConfig,
11
LlavaOnevisionProcessor)
12
from transformers.models.llava_onevision.modeling_llava_onevision 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.layers.activation import get_act_fn
@@ -23,6 +22,7 @@ from vllm.multimodal.parse import (ImageSize, MultiModalDataItems,
22
VideoEmbeddingItems, VideoProcessorItems)
23
from vllm.multimodal.processing import PromptReplacement, PromptUpdate
24
from vllm.sequence import IntermediateTensors
25
+from vllm.utils.tensor_schema import TensorSchema, TensorShape
26
27
from .clip import CLIPVisionModel
28
from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP
@@ -38,44 +38,62 @@ from .utils import (AutoWeightsLoader, WeightsMapper, flatten_bn,
38
_MAX_FRAMES_PER_VIDEO = 16
39
40
41
-class LlavaOnevisionVideoPixelInputs(TypedDict):
42
- type: Literal["pixel_values_videos"]
43
- pixel_values_videos: Union[torch.Tensor, list[torch.Tensor]]
41
+class LlavaOnevisionVideoPixelInputs(TensorSchema):
42
"""
45
- Shape: `(batch_size * num_videos, num_frames, num_channels, height, width)`
46
-
47
- Note that `num_videos` may be different for each batch, and 'num_frames'
48
- may be different for each video, in which case the data is passed as a
49
- list instead of a batched tensor.
43
+ Dimensions:
44
+ - bn: Batch size * number of videos
45
+ - f: Number of frames
46
+ - c: Number of channels (3)
47
+ - h: Height
48
+ - w: Width
49
+
50
+ Note that `num_videos` may be different for each batch, and 'num_frames'
51
+ may be different for each video, in which case the data is passed as a
52
+ list instead of a batched tensor.
53
"""
54
+ type: Literal["pixel_values_videos"] = "pixel_values_videos"
55
56
+ pixel_values_videos: Annotated[
57
+ Union[torch.Tensor, list[torch.Tensor]],
58
+ TensorShape("bn", "f", 3, "h", "w", dynamic_dims={"f"}),
59
+ ]
60
53
-class LlavaOnevisionImagePixelInputs(TypedDict):
54
- type: Literal["pixel_values"]
55
- pixel_values: Union[torch.Tensor, list[torch.Tensor]]
56
- """
57
- Shape:
58
- `(batch_size * num_images, 1 + num_patches, num_channels, height, width)`
61
60
- Note that `num_patches` may be different per batch and image,
61
- in which case the data is passed as a list instead of a batched tensor.
62
+class LlavaOnevisionImagePixelInputs(TensorSchema):
63
"""
63
-
64
- image_sizes: NotRequired[torch.Tensor]
64
+ Dimensions:
65
+ - bn: Batch size * number of images
66
+ - np: Number of patches (1 + num_patches)
67
+ - c: Number of channels (3)
68
+ - h: Height
69
+ - w: Width
70
+
71
+ Note that `num_patches` may be different per batch and image,
72
+ in which case the data is passed as a list instead of a batched tensor.
73
"""
66
- Shape: `(batch_size * num_images, 2)`
74
+ type: Literal["pixel_values"] = "pixel_values"
75
68
- This should be in `(height, width)` format.
69
- """
76
+ pixel_values: Annotated[
77
+ Union[torch.Tensor, list[torch.Tensor]],
78
+ TensorShape("bn", "np", 3, "h", "w"),
79
+ ]
80
81
+ image_sizes: Annotated[Optional[torch.Tensor], TensorShape("bn", 2)]
82
72
-class LlavaOnevisionImageEmbeddingInputs(TypedDict):
73
- type: Literal["image_embeds"]
74
- data: torch.Tensor
75
- """Shape: `(batch_size * num_images, image_feature_size, hidden_size)`
83
77
- `hidden_size` must match the hidden size of language model backbone.
84
+class LlavaOnevisionImageEmbeddingInputs(TensorSchema):
85
"""
86
+ Dimensions:
87
+ - bn: Batch size * number of images
88
+ - ifs: Image feature size
89
+ - hs: Hidden size (must match language model backbone)
90
+ """
91
+ type: Literal["image_embeds"] = "image_embeds"
92
+
93
+ data: Annotated[
94
+ torch.Tensor,
95
+ TensorShape("bn", "ifs", "hs"),
96
+ ]
97
98
99
LlavaOnevisionImageInputs = Union[LlavaOnevisionImagePixelInputs,
@@ -482,44 +500,6 @@ class LlavaOnevisionForConditionalGeneration(nn.Module, SupportsMultiModal,
500
self.make_empty_intermediate_tensors = (
501
self.language_model.model.make_empty_intermediate_tensors)
502
485
- def _validate_image_sizes(self, data: torch.Tensor) -> torch.Tensor:
486
- expected_dims = (2, )
487
-
488
- def _validate_shape(d: torch.Tensor):
489
- actual_dims = tuple(d.shape)
490
-
491
- if actual_dims != expected_dims:
492
- expected_expr = str(expected_dims)
493
- raise ValueError(
494
- f"The expected shape of image sizes per image per batch "
495
- f"is {expected_expr}. You supplied {tuple(d.shape)}.")
496
-
497
- for d in data:
498
- _validate_shape(d)
499
-
500
- return data
501
-
502
- def _validate_image_pixel_values(
503
- self, data: Union[torch.Tensor, list[torch.Tensor]]
504
- ) -> Union[torch.Tensor, list[torch.Tensor]]:
505
-
506
- h = w = self.config.vision_config.image_size
507
- expected_dims = (3, h, w)
508
-
509
- def _validate_shape(d: torch.Tensor):
510
- actual_dims = tuple(d.shape[1:])
511
-
512
- if actual_dims != expected_dims:
513
- expected_expr = ("num_patches", *map(str, expected_dims))
514
- raise ValueError(
515
- "The expected shape of pixel values per image per batch "
516
- f"is {expected_expr}. You supplied {tuple(d.shape)}.")
517
-
518
- for d in data:
519
- _validate_shape(d)
520
-
521
- return data
522
-
503
def _parse_and_validate_image_input(
504
self, **kwargs: object) -> Optional[LlavaOnevisionImageInputs]:
505
pixel_values = kwargs.pop("pixel_values", None)
@@ -540,11 +520,12 @@ class LlavaOnevisionForConditionalGeneration(nn.Module, SupportsMultiModal,
520
521
return LlavaOnevisionImagePixelInputs(
522
type="pixel_values",
543
- pixel_values=self._validate_image_pixel_values(
544
- flatten_bn(pixel_values)),
545
- image_sizes=self._validate_image_sizes(
546
- flatten_bn(image_sizes, concat=True)),
547
- )
523
+ pixel_values=flatten_bn(pixel_values),
524
+ image_sizes=flatten_bn(image_sizes, concat=True),
525
+ resolve_bindings={
526
+ "h": self.config.vision_config.image_size,
527
+ "w": self.config.vision_config.image_size
528
+ })
529
530
if image_embeds is not None:
531
if not isinstance(image_embeds, torch.Tensor):
@@ -558,27 +539,6 @@ class LlavaOnevisionForConditionalGeneration(nn.Module, SupportsMultiModal,
539
540
raise AssertionError("This line should be unreachable.")
541
561
- def _validate_video_pixel_values(
562
- self, data: Union[torch.Tensor, list[torch.Tensor]]
563
- ) -> Union[torch.Tensor, list[torch.Tensor]]:
564
-
565
- h = w = self.config.vision_config.image_size
566
- expected_dims = (3, h, w)
567
-
568
- def _validate_shape(d: torch.Tensor):
569
- actual_dims = tuple(d.shape[2:])
570
-
571
- if actual_dims != expected_dims:
572
- expected_expr = ("num_frames", *map(str, expected_dims))
573
- raise ValueError(
574
- "The expected shape of pixel values in each video frame "
575
- f"is {expected_expr}. You supplied {tuple(d.shape)}.")
576
-
577
- for d in data:
578
- _validate_shape(d)
579
-
580
- return data
581
-
542
def _parse_and_validate_video_input(
543
self,
544
**kwargs: object) -> Optional[LlavaOnevisionVideoPixelInputs]:
@@ -600,7 +560,10 @@ class LlavaOnevisionForConditionalGeneration(nn.Module, SupportsMultiModal,
560
return LlavaOnevisionVideoPixelInputs(
561
type="pixel_values_videos",
562
pixel_values_videos=flatten_bn(pixel_values_videos),
603
- )
563
+ resolve_bindings={
564
+ "h": self.config.vision_config.image_size,
565
+ "w": self.config.vision_config.image_size
566
+ })
567
568
def _parse_and_validate_multimodal_inputs(self, **kwargs: object) -> dict:
569
mm_input_by_modality = {}