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 incorrect Molmo image processing (#26563)
Signed-off-by: sanghol <sanghol@allenai.org>
sangho.lee committed
Oct 11, 2025 at 00:28 UTC
55392bc87932da63b888e58f371fe4b67b438496
1 file changed
+23
-17
vllm/model_executor/models/molmo.py
+23
-17
@@ -114,11 +114,11 @@ class MolmoImageInputs(TensorSchema):
114
TensorShape("bn", "nc", "np", dynamic_dims={"nc"}),
115
]
116
117
- feat_is_patch: Annotated[
117
+ image_input_idx: Annotated[
118
Union[torch.Tensor, list[torch.Tensor]],
119
TensorShape("bn", "nc", "tp", dynamic_dims={"nc"}),
120
]
121
- # A boolean mask indicating which image features correspond to patch tokens.
121
+ # An index tensor that maps image features to their corresponding patch tokens.
122
num_crops: Annotated[torch.Tensor, TensorShape("bn")]
123
124
@@ -1177,7 +1177,7 @@ class MolmoProcessorWrapper:
1177
num_crops = torch.tensor(tilings).prod(-1) + 1
1178
assert num_crops.sum() == len(feat_is_patch)
1179
1180
- outputs["feat_is_patch"] = feat_is_patch
1180
+ outputs["image_input_idx"] = image_input_idx
1181
outputs["num_crops"] = num_crops
1182
outputs["img_patch_id"] = self.image_patch_id
1183
@@ -1211,8 +1211,9 @@ class MolmoProcessingInfo(BaseProcessingInfo):
1211
image_token_length_w = processor.image_token_length_w
1212
image_token_length_h = processor.image_token_length_h
1213
1214
- extra = image_token_length_w * image_token_length_h
1215
- joint = ((ncols + 1) // pooling_size) * ((nrows + 1) // pooling_size)
1214
+ # Calculate total tokens: 2 for start/end + (w+1)*h for column separators
1215
+ extra = 2 + (image_token_length_w + 1) * image_token_length_h
1216
+ joint = 2 + ((ncols + 1) // pooling_size + 1) * ((nrows + 1) // pooling_size)
1217
1218
return extra + joint
1219
@@ -1299,7 +1300,7 @@ class MolmoMultiModalProcessor(BaseMultiModalProcessor[MolmoProcessingInfo]):
1300
return dict(
1301
images=MultiModalFieldConfig.flat_from_sizes("image", num_crops),
1302
image_masks=MultiModalFieldConfig.flat_from_sizes("image", num_crops),
1302
- feat_is_patch=MultiModalFieldConfig.flat_from_sizes("image", num_crops),
1303
+ image_input_idx=MultiModalFieldConfig.flat_from_sizes("image", num_crops),
1304
num_crops=MultiModalFieldConfig.batched("image"),
1305
img_patch_id=MultiModalFieldConfig.shared("image", num_images),
1306
)
@@ -1444,7 +1445,7 @@ class MolmoForCausalLM(
1445
) -> Optional[MolmoImageInputs]:
1446
images = kwargs.pop("images", None)
1447
image_masks = kwargs.pop("image_masks", None)
1447
- feat_is_patch = kwargs.pop("feat_is_patch", None)
1448
+ image_input_idx = kwargs.pop("image_input_idx", None)
1449
num_crops = kwargs.pop("num_crops", None)
1450
1451
if images is None:
@@ -1466,7 +1467,7 @@ class MolmoForCausalLM(
1467
return MolmoImageInputs(
1468
images=images,
1469
image_masks=image_masks,
1469
- feat_is_patch=feat_is_patch,
1470
+ image_input_idx=image_input_idx,
1471
num_crops=num_crops,
1472
)
1473
@@ -1476,7 +1477,7 @@ class MolmoForCausalLM(
1477
) -> list[torch.Tensor]:
1478
images = image_input["images"]
1479
image_masks = image_input["image_masks"]
1479
- feat_is_patch = image_input["feat_is_patch"]
1480
+ image_input_idx = image_input["image_input_idx"]
1481
num_crops = image_input["num_crops"]
1482
1483
# Call the vision backbone on the whole batch at once
@@ -1484,7 +1485,7 @@ class MolmoForCausalLM(
1485
image_masks_flat = (
1486
None if image_masks is None else flatten_bn(image_masks, concat=True)
1487
)
1487
- feat_is_patch_flat = flatten_bn(feat_is_patch, concat=True)
1488
+ image_input_idx_flat = flatten_bn(image_input_idx, concat=True)
1489
1490
image_features_flat = self.vision_backbone(
1491
images=images_flat.unsqueeze(0),
@@ -1494,13 +1495,18 @@ class MolmoForCausalLM(
1495
).squeeze(0)
1496
1497
# Only the features corresponding to patch tokens are relevant
1497
- return [
1498
- feats[f_is_patch]
1499
- for feats, f_is_patch in zip(
1500
- image_features_flat.split(num_crops.tolist()),
1501
- feat_is_patch_flat.split(num_crops.tolist()),
1502
- )
1503
- ]
1498
+ # Re-order the features using the image_input_idx tensor
1499
+ results = []
1500
+ num_crops_list = num_crops.tolist()
1501
+ for feats, img_idx in zip(
1502
+ image_features_flat.split(num_crops_list),
1503
+ image_input_idx_flat.split(num_crops_list),
1504
+ ):
1505
+ is_valid = img_idx >= 0
1506
+ valid_img_idx = img_idx[is_valid]
1507
+ order = torch.argsort(valid_img_idx)
1508
+ results.append(feats[is_valid][order])
1509
+ return results
1510
1511
def get_language_model(self) -> torch.nn.Module:
1512
return self.model