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, 23:45 UTC
Migrate KeyeImageInputs and KeyeVideoInputs to TensorSchema (#21686)
Signed-off-by: Benji Beck <benjibeck@meta.com>
Benji Beck committed
Jul 28, 2025 at 01:16 UTC
d128d0d554057ac91f93abe8f06e671203b121e2
1 file changed
+41
-65
vllm/model_executor/models/keye.py
+41
-65
@@ -3,7 +3,7 @@
3
import math
4
from collections.abc import Iterable, Mapping, Sequence
5
from functools import partial
6
-from typing import Any, Literal, Optional, TypedDict, Union
6
+from typing import Annotated, Any, Literal, Optional, Union
7
8
import numpy as np
9
import torch
@@ -46,6 +46,7 @@ from vllm.sequence import IntermediateTensors
46
from vllm.transformers_utils.config import uses_mrope
47
from vllm.transformers_utils.processor import (
48
cached_image_processor_from_config)
49
+from vllm.utils.tensor_schema import TensorSchema, TensorShape
50
51
from .interfaces import (MultiModalEmbeddings, SupportsLoRA,
52
SupportsMultiModal, SupportsPP)
@@ -102,77 +103,62 @@ def smart_resize(
103
return h_bar, w_bar
104
105
105
-class KeyeImagePixelInputs(TypedDict):
106
- type: Literal["pixel_values"]
107
- pixel_values: torch.Tensor
108
- """Shape:
109
- `(num_patches, num_channels * patch_size * patch_size)`
106
+class KeyeImagePixelInputs(TensorSchema):
107
"""
111
-
112
- image_grid_thw: torch.Tensor
113
- """Shape: `(num_images, 3)`
114
- This should be in `(grid_t, grid_h, grid_w)` format.
108
+ Dimensions:
109
+ - np: Number of patches
110
+ - cps: Number of channels * patch_size * patch_size
111
+ - ni: Number of images
112
+ - g: Grid dimensions (3 for t, h, w)
113
"""
114
+ type: Literal["pixel_values"]
115
+ pixel_values: Annotated[torch.Tensor, TensorShape("np", "cps")]
116
+ image_grid_thw: Annotated[torch.Tensor, TensorShape("ni", 3)]
117
118
118
-class KeyeImageEmbeddingInputs(TypedDict):
119
- type: Literal["image_embeds"]
120
- image_embeds: torch.Tensor
121
- """Supported types:
122
- - list[`torch.Tensor`]: A list of tensors holding all images' features.
123
- Each tensor holds an image's features.
124
- - `torch.Tensor`: A tensor holding all images' features
125
- (concatenation of all images' feature tensors).
126
-
127
- Tensor shape: `(num_image_features, hidden_size)`
128
- - `num_image_features` varies based on
129
- the number and resolution of the images.
130
- - `hidden_size` must match the hidden size of language model backbone.
119
+class KeyeImageEmbeddingInputs(TensorSchema):
120
"""
132
-
133
- image_grid_thw: torch.Tensor
134
- """Shape: `(num_images, 3)`
135
- This should be in `(grid_t, grid_h, grid_w)` format.
121
+ Dimensions:
122
+ - nf: Number of image features
123
+ - hs: Hidden size (must match the hidden size of language model
124
+ backbone)
125
+ - ni: Number of images
126
+ - g: Grid dimensions (3 for t, h, w)
127
"""
128
+ type: Literal["image_embeds"]
129
+ image_embeds: Annotated[torch.Tensor, TensorShape("nf", "hs")]
130
+ image_grid_thw: Annotated[torch.Tensor, TensorShape("ni", 3)]
131
132
133
KeyeImageInputs = Union[KeyeImagePixelInputs, KeyeImageEmbeddingInputs]
134
135
142
-class KeyeVideoPixelInputs(TypedDict):
143
- type: Literal["pixel_values_videos"]
144
- pixel_values_videos: torch.Tensor
145
- """Shape:
146
- `(num_patches,
147
- num_channels * temporal_patch_size * patch_size * patch_size)`
136
+class KeyeVideoPixelInputs(TensorSchema):
137
"""
149
-
150
- video_grid_thw: torch.Tensor
151
- """Shape: `(num_videos, 3)`
152
-
153
- This should be in `(grid_t, grid_h, grid_w)` format.
138
+ Dimensions:
139
+ - np: Number of patches
140
+ - ctps: Number of channels * temporal_patch_size * patch_size *
141
+ patch_size
142
+ - nv: Number of videos
143
+ - g: Grid dimensions (3 for t, h, w)
144
"""
145
+ type: Literal["pixel_values_videos"]
146
+ pixel_values_videos: Annotated[torch.Tensor, TensorShape("np", "ctps")]
147
+ video_grid_thw: Annotated[torch.Tensor, TensorShape("nv", 3)]
148
149
157
-class KeyeVideoEmbeddingInputs(TypedDict):
158
- type: Literal["video_embeds"]
159
- video_embeds: torch.Tensor
160
- """Supported types:
161
- - list[`torch.Tensor`]: A list of tensors holding all videos' features.
162
- Each tensor holds an video's features.
163
- - `torch.Tensor`: A tensor holding all videos' features
164
- (concatenation of all videos' feature tensors).
165
-
166
- Tensor shape: `(num_image_features, hidden_size)`
167
- - `num_image_features` varies based on
168
- the number and resolution of the videos.
169
- - `hidden_size` must match the hidden size of language model backbone.
150
+class KeyeVideoEmbeddingInputs(TensorSchema):
151
"""
171
-
172
- video_grid_thw: torch.Tensor
173
- """Shape: `(num_videos, 3)`
174
- This should be in `(grid_t, grid_h, grid_w)` format.
152
+ Dimensions:
153
+ - nf: Number of video features
154
+ - hs: Hidden size (must match the hidden size of language model
155
+ backbone)
156
+ - nv: Number of videos
157
+ - g: Grid dimensions (3 for t, h, w)
158
"""
159
+ type: Literal["video_embeds"]
160
+ video_embeds: Annotated[torch.Tensor, TensorShape("nf", "hs")]
161
+ video_grid_thw: Annotated[torch.Tensor, TensorShape("nv", 3)]
162
163
164
KeyeVideoInputs = Union[KeyeVideoPixelInputs, KeyeVideoEmbeddingInputs]
@@ -1420,10 +1406,6 @@ class KeyeForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsLoRA,
1406
image_grid_thw = self._validate_and_reshape_mm_tensor(
1407
image_grid_thw, "image grid_thw")
1408
1423
- if not isinstance(pixel_values, (torch.Tensor, list)):
1424
- raise ValueError("Incorrect type of image pixel values. "
1425
- f"Got type: {type(pixel_values)}")
1426
-
1409
return KeyeImagePixelInputs(
1410
type="pixel_values",
1411
pixel_values=pixel_values,
@@ -1436,9 +1418,6 @@ class KeyeForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsLoRA,
1418
image_grid_thw = self._validate_and_reshape_mm_tensor(
1419
image_grid_thw, "image grid_thw")
1420
1439
- if not isinstance(image_embeds, torch.Tensor):
1440
- raise ValueError("Incorrect type of image embeddings. "
1441
- f"Got type: {type(image_embeds)}")
1421
return KeyeImageEmbeddingInputs(
1422
type="image_embeds",
1423
image_embeds=image_embeds,
@@ -1474,9 +1453,6 @@ class KeyeForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsLoRA,
1453
video_grid_thw = self._validate_and_reshape_mm_tensor(
1454
video_grid_thw, "video grid_thw")
1455
1477
- if not isinstance(video_embeds, torch.Tensor):
1478
- raise ValueError("Incorrect type of video embeddings. "
1479
- f"Got type: {type(video_embeds)}")
1456
return KeyeVideoEmbeddingInputs(
1457
type="video_embeds",
1458
video_embeds=video_embeds,