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 Mistral3ImagePixelInputs to TensorSchema (#21945)
Signed-off-by: Benji Beck <benjibeck@meta.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>
Benji Beck committed
Aug 20, 2025 at 10:37 UTC
c4477f55e581e5ef5f52bbe39cba6e0de1956444
1 file changed
+17
-21
vllm/model_executor/models/mistral3.py
+17
-21
@@ -3,7 +3,7 @@
3
4
from abc import abstractmethod
5
from collections.abc import Iterable, Mapping, Sequence
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
@@ -32,6 +32,7 @@ from vllm.multimodal.processing import (BaseMultiModalProcessor,
32
PromptUpdateDetails)
33
from vllm.multimodal.profiling import BaseDummyInputsBuilder
34
from vllm.sequence import IntermediateTensors
35
+from vllm.utils.tensor_schema import TensorSchema, TensorShape
36
37
from .interfaces import (MultiModalEmbeddings, SupportsLoRA,
38
SupportsMultiModal, SupportsPP)
@@ -42,16 +43,24 @@ from .utils import (AutoWeightsLoader, WeightsMapper, flatten_bn,
43
from .vision import get_vision_encoder_info
44
45
45
-class Mistral3ImagePixelInputs(TypedDict):
46
- type: Literal["pixel_values_pixtral"]
47
- pixel_values: Union[torch.Tensor, list[torch.Tensor]]
46
+class Mistral3ImagePixelInputs(TensorSchema):
47
"""
49
- Shape: `(batch_size * num_images, num_channels, height, width)`
50
-
51
- Note that `height` or `width` may be different per batch and image,
52
- in which case the data is passed as a list instead of a batched tensor.
48
+ Dimensions:
49
+ - bn: Batch size * number of images
50
+ - c: Number of channels (3)
51
+ - h: Height of each image
52
+ - w: Width of each image
53
"""
54
55
+ type: Literal["pixel_values_pixtral"] = "pixel_values_pixtral"
56
+
57
+ # Note that `height` or `width` may be different per batch and image,
58
+ # in which case the data is passed as a list instead of a batched tensor.
59
+ pixel_values: Annotated[
60
+ Union[torch.Tensor, list[torch.Tensor]],
61
+ TensorShape("bn", 3, "h", "w", dynamic_dims={"h", "w"}),
62
+ ]
63
+
64
65
class Mistral3PatchMerger(nn.Module):
66
"""
@@ -456,19 +465,6 @@ class Mistral3ForConditionalGeneration(nn.Module, SupportsLoRA,
465
self.make_empty_intermediate_tensors = (
466
self.language_model.make_empty_intermediate_tensors)
467
459
- def _validate_pixel_values(self, data: torch.Tensor) -> torch.Tensor:
460
- h = w = self.config.vision_config.image_size
461
- expected_dims = (3, h, w)
462
- actual_dims = tuple(data.shape[1:])
463
-
464
- if actual_dims != expected_dims:
465
- expected_expr = ("batch_size", *map(str, expected_dims))
466
- raise ValueError(
467
- f"The expected shape of pixel values is {expected_expr}. "
468
- f"You supplied {tuple(data.shape)}.")
469
-
470
- return data
471
-
468
def _parse_and_validate_image_input(
469
self, **kwargs: object) -> Optional[Mistral3ImagePixelInputs]:
470
pixel_values = kwargs.pop("pixel_values", None)