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 MolmoImageInputs to TensorSchema (#22022)
Signed-off-by: Benji Beck <benjibeck@meta.com>
Benji Beck committed
Aug 21, 2025 at 09:54 UTC
a482e4e7696141cba45180d7463ddde691d512ae
1 file changed
+23
-33
vllm/model_executor/models/molmo.py
+23
-33
@@ -5,7 +5,7 @@ import math
5
from collections.abc import Iterable, Mapping, Sequence
6
from dataclasses import dataclass
7
from functools import cached_property, partial
8
-from typing import Optional, TypedDict, Union
8
+from typing import Annotated, Optional, Union
9
10
import numpy as np
11
import torch
@@ -51,6 +51,7 @@ from vllm.multimodal.processing import (BaseMultiModalProcessor,
51
PromptUpdateDetails)
52
from vllm.multimodal.profiling import BaseDummyInputsBuilder
53
from vllm.sequence import IntermediateTensors
54
+from vllm.utils.tensor_schema import TensorSchema, TensorShape
55
56
from .interfaces import (MultiModalEmbeddings, SupportsLoRA,
57
SupportsMultiModal, SupportsPP, SupportsQuant)
@@ -70,23 +71,25 @@ IM_END_TOKEN = "<im_end>"
71
POOLING_SIZE = 2
72
73
73
-class MolmoImageInputs(TypedDict):
74
- images: Union[torch.Tensor, list[torch.Tensor]]
75
- """Shape: `(batch_size * num_images, num_crops, num_patch, patch_dim)`"""
76
-
77
- image_masks: Optional[Union[torch.Tensor, list[torch.Tensor]]]
78
- """Shape: `(batch_size * num_images, num_crops, num_patch)`"""
79
-
80
- feat_is_patch: Union[torch.Tensor, list[torch.Tensor]]
74
+class MolmoImageInputs(TensorSchema):
75
"""
82
- A boolean mask indicating which image features correspond
83
- to patch tokens.
84
-
85
- Shape: `(batch_size * num_images, num_crops, num_patch)`
76
+ Dimensions:
77
+ - bn: Batch size * number of images
78
+ - nc: Number of crops
79
+ - np: Number of patches
80
+ - pd: Patch dimension
81
"""
82
+ images: Annotated[Union[torch.Tensor, list[torch.Tensor]],
83
+ TensorShape("bn", "nc", "np", "pd")]
84
+
85
+ image_masks: Annotated[Optional[Union[torch.Tensor, list[torch.Tensor]]],
86
+ TensorShape("bn", "nc", "np")]
87
+
88
+ feat_is_patch: Annotated[Union[torch.Tensor, list[torch.Tensor]],
89
+ TensorShape("bn", "nc", "np")]
90
+ # A boolean mask indicating which image features correspond to patch tokens.
91
88
- num_crops: torch.Tensor
89
- """Shape: `(batch_size * num_images)`"""
92
+ num_crops: Annotated[torch.Tensor, TensorShape("bn")]
93
94
95
@dataclass
@@ -1410,28 +1413,17 @@ class MolmoForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA,
1413
**kwargs: object,
1414
) -> Optional[MolmoImageInputs]:
1415
images = kwargs.pop("images", None)
1413
- if images is None:
1414
- return None
1415
-
1416
- if not isinstance(images, (torch.Tensor, list)):
1417
- raise ValueError("Incorrect type of images. "
1418
- f"Got type: {type(images)}")
1419
-
1416
image_masks = kwargs.pop("image_masks", None)
1421
- if not (image_masks is None or isinstance(image_masks,
1422
- (torch.Tensor, list))):
1423
- raise ValueError("Incorrect type of image_masks. "
1424
- f"Got type: {type(image_masks)}")
1425
-
1417
feat_is_patch = kwargs.pop("feat_is_patch", None)
1427
- if not isinstance(feat_is_patch, (torch.Tensor, list)):
1428
- raise ValueError("Incorrect type of feat_is_patch. "
1429
- f"Got type: {type(feat_is_patch)}")
1430
-
1418
num_crops = kwargs.pop("num_crops", None)
1419
+
1420
+ if images is None:
1421
+ return None
1422
+
1423
if not isinstance(num_crops, (torch.Tensor, list)):
1424
raise ValueError("Incorrect type of num_crops. "
1425
f"Got type: {type(num_crops)}")
1426
+ num_crops = flatten_bn(num_crops, concat=True)
1427
1428
img_patch_id = kwargs.pop("img_patch_id", None)
1429
if not isinstance(img_patch_id, torch.Tensor):
@@ -1439,8 +1431,6 @@ class MolmoForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA,
1431
f"Got type: {type(img_patch_id)}")
1432
self.img_patch_id = img_patch_id.flatten().unique().item()
1433
1442
- num_crops = flatten_bn(num_crops, concat=True)
1443
-
1434
return MolmoImageInputs(
1435
images=images,
1436
image_masks=image_masks,