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
[BugFix][Multi Modal] Fix TensorSchema shape mismatch in Molmo (#24559)
Signed-off-by: wwl2755 <wangwenlong2755@gmail.com>
Wenlong Wang committed
Sep 10, 2025 at 06:14 UTC
4c04eef706fcd26c8e046cd797fa9b14a5fb361d
1 file changed
+8
-6
vllm/model_executor/models/molmo.py
+8
-6
@@ -76,20 +76,22 @@ class MolmoImageInputs(TensorSchema):
76
"""
77
Dimensions:
78
- bn: Batch size * number of images
79
- - nc: Number of crops
79
+ - nc: Number of crops (dynamic)
80
- np: Number of patches
81
+ - tp: Token sequence positions
82
- pd: Patch dimension
83
"""
84
images: Annotated[Union[torch.Tensor, list[torch.Tensor]],
84
- TensorShape("bn", "nc", "np", "pd")]
85
+ TensorShape("bn", "nc", "np", "pd", dynamic_dims={"nc"})]
86
+ # Number of crops may vary per batch and image, so pass it as a list.
87
88
image_masks: Annotated[Optional[Union[torch.Tensor, list[torch.Tensor]]],
87
- TensorShape("bn", "nc", "np")]
89
+ TensorShape("bn", "nc", "np", dynamic_dims={"nc"})]
90
89
- feat_is_patch: Annotated[Union[torch.Tensor, list[torch.Tensor]],
90
- TensorShape("bn", "nc", "np")]
91
+ feat_is_patch: Annotated[
92
+ Union[torch.Tensor, list[torch.Tensor]],
93
+ TensorShape("bn", "nc", "tp", dynamic_dims={"nc"})]
94
# A boolean mask indicating which image features correspond to patch tokens.
92
-
95
num_crops: Annotated[torch.Tensor, TensorShape("bn")]
96
97