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 MllamaImagePixelInputs to TensorSchema (#22020)
Signed-off-by: Benji Beck <benjibeck@meta.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>
Benji Beck committed
Aug 21, 2025 at 20:28 UTC
0b9cc56fac3ee7ba0f7a117ad34809efd3f9d179
1 file changed
+26
-10
vllm/model_executor/models/mllama.py
+26
-10
@@ -17,7 +17,7 @@
17
"""PyTorch Mllama model."""
18
import math
19
from collections.abc import Iterable, Mapping, Sequence
20
-from typing import Literal, Optional, TypedDict, Union
20
+from typing import Annotated, Literal, Optional, Union
21
22
import numpy as np
23
import torch
@@ -64,6 +64,7 @@ from vllm.multimodal.processing import (BaseProcessingInfo,
64
EncDecMultiModalProcessor,
65
PromptReplacement, PromptUpdate)
66
from vllm.multimodal.profiling import BaseDummyInputsBuilder
67
+from vllm.utils.tensor_schema import TensorSchema, TensorShape
68
69
from .clip import CLIPMLP
70
from .interfaces import SupportsMultiModal, SupportsV0Only
@@ -73,15 +74,30 @@ from .utils import AutoWeightsLoader, WeightsMapper, maybe_prefix
74
logger = init_logger(__name__)
75
76
76
-class MllamaImagePixelInputs(TypedDict):
77
- type: Literal["pixel_values"]
78
- data: torch.Tensor
79
- """Shape: """
80
- """(batch_size, max_num_image, max_num_chunk, num_channel, height, width)"""
81
- aspect_ratio_ids: torch.Tensor
82
- """Shape: `(batch_size, max_num_image)`"""
83
- aspect_ratio_mask: torch.Tensor
84
- """Shape: `(batch_size, max_num_image, max_num_tiles)`"""
77
+class MllamaImagePixelInputs(TensorSchema):
78
+ """
79
+ Dimensions:
80
+ - batch_size: Batch size
81
+ - max_num_image: Max number of images
82
+ - max_num_chunk: Max number of chunks
83
+ - max_num_tiles: Max number of tiles per image
84
+ - num_channel: Number of channels
85
+ - height: Height
86
+ - width: Width
87
+ """
88
+
89
+ type: Literal["pixel_values"] = "pixel_values"
90
+
91
+ data: Annotated[torch.Tensor,
92
+ TensorShape("batch_size", "max_num_image", "max_num_chunk",
93
+ "num_channel", "height", "width")]
94
+
95
+ aspect_ratio_ids: Annotated[torch.Tensor,
96
+ TensorShape("batch_size", "max_num_image")]
97
+
98
+ aspect_ratio_mask: Annotated[
99
+ torch.Tensor,
100
+ TensorShape("batch_size", "max_num_image", "max_num_tiles")]
101
102
103
# TODO: support LlamaImageEmbeddingInputs