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 Pixtral inputs to TensorSchema (#23472)
Signed-off-by: Benji Beck <benjibeck@meta.com>
Benji Beck committed
Aug 23, 2025 at 21:55 UTC
053278a5dc7a81d751f8e63c1ed793062b32cbce
1 file changed
+13
-11
vllm/model_executor/models/pixtral.py
+13
-11
@@ -5,7 +5,7 @@ import math
5
from collections.abc import Iterable, Mapping, Sequence
6
from dataclasses import dataclass, fields
7
from functools import cached_property
8
-from typing import Literal, Optional, TypedDict, Union
8
+from typing import Annotated, Literal, Optional, Union
9
10
import torch
11
import torch.nn as nn
@@ -48,6 +48,7 @@ from vllm.platforms import current_platform
48
from vllm.sequence import IntermediateTensors
49
from vllm.transformers_utils.tokenizer import (MistralTokenizer,
50
cached_tokenizer_from_config)
51
+from vllm.utils.tensor_schema import TensorSchema, TensorShape
52
53
from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP
54
from .utils import (flatten_bn, init_vllm_registered_model, maybe_prefix,
@@ -68,15 +69,20 @@ except ImportError:
69
PATCH_MERGE = "patch_merge"
70
71
71
-class PixtralImagePixelInputs(TypedDict):
72
- type: Literal["pixel_values"]
73
-
74
- images: Union[torch.Tensor, list[torch.Tensor]]
72
+class PixtralImagePixelInputs(TensorSchema):
73
"""
76
- Shape: `(batch_size * num_images, num_channels, image_width, image_height)`
77
-
74
+ Dimensions:
75
+ - bn: Batch size * number of images
76
+ - c: Number of channels (3)
77
+ - h: Height of each image
78
+ - w: Width of each image
79
+
80
The result of stacking `ImageEncoding.tokens` from each prompt.
81
"""
82
+ type: Literal["pixel_values"] = "pixel_values"
83
+
84
+ images: Annotated[Union[torch.Tensor, list[torch.Tensor]],
85
+ TensorShape("bn", 3, "h", "w", dynamic_dims={"h", "w"})]
86
87
88
class PixtralProcessorAdapter:
@@ -381,10 +387,6 @@ class PixtralForConditionalGeneration(nn.Module, SupportsMultiModal,
387
if images is None:
388
return None
389
384
- if not isinstance(images, (torch.Tensor, list)):
385
- raise ValueError("Incorrect type of images. "
386
- f"Got type: {type(images)}")
387
-
390
return PixtralImagePixelInputs(
391
type="pixel_values",
392
images=flatten_bn(images),