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 OvisImagePatchInputs to TensorSchema (#22024)
Signed-off-by: Benji Beck <benjibeck@meta.com>
Benji Beck committed
Sep 1, 2025 at 21:01 UTC
1fa1d6a9a0784855b420a75c2b42d6600d62cb41
1 file changed
+20
-19
vllm/model_executor/models/ovis.py
+20
-19
@@ -19,7 +19,7 @@
19
""" PyTorch Ovis model."""
20
import math
21
from collections.abc import Iterable, Mapping
22
-from typing import Literal, Optional, TypedDict, Union
22
+from typing import Annotated, Literal, Optional, Union
23
24
import torch
25
import torch.nn as nn
@@ -49,6 +49,7 @@ from vllm.multimodal.processing import (BaseMultiModalProcessor,
49
from vllm.multimodal.profiling import BaseDummyInputsBuilder
50
from vllm.sequence import IntermediateTensors
51
from vllm.transformers_utils.processors.ovis import OvisProcessor
52
+from vllm.utils.tensor_schema import TensorSchema, TensorShape
53
54
from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP
55
from .utils import merge_multimodal_embeddings
@@ -201,25 +202,22 @@ class VisualTokenizer(torch.nn.Module):
202
return tokens
203
204
204
-class OvisImagePatchInputs(TypedDict):
205
- type: Literal["image_patches"]
206
- flat_data: torch.Tensor
207
- """
208
- Shape:
209
- `(batch_size * num_patches, patch_size_x * patch_size_y * num_channels)`
210
- """
211
-
212
- indicator_tokens: torch.Tensor
205
+class OvisImagePatchInputs(TensorSchema):
206
"""
214
- Shape:
215
- `(batch_size * (num_patches + 1))`
216
- """
217
-
218
- patches_per_image: list[int]
219
- """
220
- List of number of total patches for each image in the batch.
221
- This is used to restore the first two dimensions of `flat_data`.
207
+ Dimensions:
208
+ - batch_patches: Batch size * number of patches
209
+ - patch_size: patch_size_x * patch_size_y * num_channels
210
+ - patch_indicators: Batch size * (number of patches + 1)
211
+ - patches_per_image: List of number of total patches for each image
212
+ in the batch.
213
"""
214
+ type: Literal["image_patches"]
215
+ flat_data: Annotated[torch.Tensor,
216
+ TensorShape("batch_patches", "patch_size")]
217
+ indicator_tokens: Annotated[torch.Tensor, TensorShape("patch_indicators")]
218
+ patches_per_image: Annotated[list[int],
219
+ TensorShape("num_patches_per_image")]
220
+ # This is used to restore the first two dimensions of `flat_data`.
221
222
223
class VisualEmbedding(torch.nn.Embedding):
@@ -458,9 +456,12 @@ class Ovis(nn.Module, SupportsMultiModal, SupportsPP):
456
raise ValueError("Incorrect type of indicator_tokens. "
457
f"Got type: {type(pixel_values)}")
458
459
+ flat_data = flatten_bn(pixel_values, concat=True)
460
+ if flat_data.ndim >= 3:
461
+ flat_data = flat_data.flatten(start_dim=1)
462
return OvisImagePatchInputs(
463
type="image_patches",
463
- flat_data=flatten_bn(flatten_bn(pixel_values), concat=True),
464
+ flat_data=flat_data,
465
patches_per_image=[
466
x.shape[0] for x in flatten_bn(pixel_values)
467
],