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 Llama4ImagePatchInputs to TensorSchema (#22021)
Signed-off-by: Benji Beck <benjibeck@meta.com>
Benji Beck committed
Aug 28, 2025 at 10:29 UTC
f32a5bc5058afc2fb601dcb456b581e2fefa94dd
1 file changed
+23
-18
vllm/model_executor/models/mllama4.py
+23
-18
@@ -19,7 +19,7 @@
19
import math
20
from collections.abc import Iterable, Mapping
21
from itertools import tee
22
-from typing import Literal, Optional, TypedDict, Union
22
+from typing import Annotated, Literal, Optional, Union
23
24
import torch
25
from torch import nn
@@ -53,6 +53,7 @@ from vllm.multimodal.processing import (BaseMultiModalProcessor,
53
from vllm.multimodal.profiling import BaseDummyInputsBuilder
54
from vllm.multimodal.utils import run_dp_sharded_vision_model
55
from vllm.sequence import IntermediateTensors
56
+from vllm.utils.tensor_schema import TensorSchema, TensorShape
57
58
from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP
59
from .llama4 import Llama4ForCausalLM
@@ -60,28 +61,34 @@ from .utils import (AutoWeightsLoader, flatten_bn, maybe_prefix,
61
merge_multimodal_embeddings)
62
63
63
-class Llama4ImagePatchInputs(TypedDict):
64
- type: Literal["pixel_values"]
65
- flat_data: torch.Tensor
64
+class Llama4ImagePatchInputs(TensorSchema):
65
"""
67
- Shape:
68
- `(batch_size * num_chunks, num_channels, image size, image size)`
66
+ Dimensions:
67
+ - batch_size: Batch size
68
+ - total_num_chunks: Batch size * number of chunks
69
+ - num_channels: Number of channels
70
+ - image_size: Size of each image
71
"""
70
- patches_per_image: torch.Tensor
72
+
73
+ type: Literal["pixel_values"] = "pixel_values"
74
+
75
+ flat_data: Annotated[torch.Tensor,
76
+ TensorShape("total_num_chunks", "num_channels",
77
+ "image_size", "image_size")]
78
+
79
+ patches_per_image: Annotated[torch.Tensor, TensorShape("batch_size")]
80
"""
81
The number of total patches for each image in the batch.
73
-
82
+
83
This is used to split the embeddings which has the first two dimensions
84
flattened just like `flat_data`.
85
"""
86
78
- aspect_ratios: Union[torch.Tensor, list[torch.Tensor]]
87
+ aspect_ratios: Annotated[torch.Tensor, TensorShape("batch_size", 2)]
88
"""
89
A list of aspect ratios corresponding to the number of tiles
90
in each dimension that each image in the batch corresponds to.
82
-
83
- Shape:
84
- `(batch_size, ratio)` where ratio is a pair `(ratio_h, ratio_w)`
91
+ Each aspect ratio is a pair (ratio_h, ratio_w).
92
"""
93
94
@@ -623,7 +630,7 @@ class Mllama4MultiModalProcessor(BaseMultiModalProcessor[Mllama4ProcessingInfo]
630
for (r_h, r_w) in aspect_ratios
631
]
632
626
- processed_outputs["aspect_ratios"] = aspect_ratios
633
+ processed_outputs["aspect_ratios"] = torch.tensor(aspect_ratios)
634
processed_outputs["patches_per_image"] = torch.tensor(
635
patches_per_image)
636
@@ -770,11 +777,9 @@ class Llama4ForConditionalGeneration(nn.Module, SupportsMultiModal,
777
# TODO: confirm handling for variable lengths
778
flat_pixel_values = flatten_bn(pixel_values, concat=True)
779
patches_per_image = flatten_bn(kwargs.pop("patches_per_image"))
773
-
774
- aspect_ratios = kwargs.pop("aspect_ratios", None)
775
- if not isinstance(aspect_ratios, (torch.Tensor, list)):
776
- raise ValueError("Incorrect type of aspect_ratios. "
777
- f"Got type: {type(aspect_ratios)}")
780
+ aspect_ratios = kwargs.pop("aspect_ratios")
781
+ if aspect_ratios.ndim == 3:
782
+ aspect_ratios = aspect_ratios.squeeze(1)
783
784
return Llama4ImagePatchInputs(
785
type="pixel_values",