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 3, 22:45 UTC
Fix Molmo2 image token metadata (#42162)
Signed-off-by: Haoqi Wang <78337154+hqhq1025@users.noreply.github.com>
Haoqing Wang committed
May 11, 2026 at 09:19 UTC
879a8c318032ed32716e5e0b10af355a4ddbced2
2 files changed
+94
-10
tests/models/multimodal/processing/test_molmo2.py
new
+56
@@ -0,0 +1,56 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from types import SimpleNamespace
5
+
6
+import torch
7
+
8
+from vllm.model_executor.models.molmo2 import build_flat_image_bool_length
9
+
10
+
11
+def test_build_flat_image_bool_length_matches_molmoweb_processor_tokens():
12
+ hf_config = SimpleNamespace(
13
+ image_patch_id=151938,
14
+ low_res_image_start_token_id=151940,
15
+ image_start_token_id=151936,
16
+ image_col_id=151939,
17
+ image_end_token_id=151937,
18
+ )
19
+ image_grids = torch.tensor([[14, 14, 14, 23]], dtype=torch.long)
20
+
21
+ image_tokens, num_image_tokens = build_flat_image_bool_length(
22
+ image_grids,
23
+ hf_config,
24
+ image_use_col_tokens=True,
25
+ use_single_crop_col_tokens=None,
26
+ use_single_crop_start_token=False,
27
+ )
28
+
29
+ assert num_image_tokens.tolist() == [550]
30
+ assert len(image_tokens) == 550
31
+ assert image_tokens[0].item() == hf_config.image_start_token_id
32
+ assert (image_tokens == hf_config.image_col_id).sum().item() == 28
33
+
34
+
35
+def test_build_flat_image_bool_length_respects_disabled_col_tokens():
36
+ hf_config = SimpleNamespace(
37
+ image_patch_id=151938,
38
+ low_res_image_start_token_id=151940,
39
+ image_start_token_id=151936,
40
+ image_col_id=151939,
41
+ image_end_token_id=151937,
42
+ )
43
+ image_grids = torch.tensor([[2, 3, 5, 7]], dtype=torch.long)
44
+
45
+ image_tokens, num_image_tokens = build_flat_image_bool_length(
46
+ image_grids,
47
+ hf_config,
48
+ image_use_col_tokens=False,
49
+ use_single_crop_col_tokens=False,
50
+ use_single_crop_start_token=True,
51
+ )
52
+
53
+ assert num_image_tokens.tolist() == [45]
54
+ assert len(image_tokens) == 45
55
+ assert image_tokens[0].item() == hf_config.low_res_image_start_token_id
56
+ assert (image_tokens == hf_config.image_col_id).sum().item() == 0
vllm/model_executor/models/molmo2.py
+38
-10
@@ -1338,6 +1338,9 @@ def exif_transpose(
1338
def build_flat_image_bool_length(
1339
image_grids: torch.LongTensor,
1340
hf_config: PretrainedConfig,
1341
+ image_use_col_tokens: bool = True,
1342
+ use_single_crop_col_tokens: bool | None = None,
1343
+ use_single_crop_start_token: bool = True,
1344
) -> tuple[torch.LongTensor, torch.LongTensor]:
1345
image_patch_id = hf_config.image_patch_id
1346
low_res_image_start_id = hf_config.low_res_image_start_token_id
@@ -1353,7 +1356,17 @@ def build_flat_image_bool_length(
1356
h = image_grids[:, 2]
1357
w = image_grids[:, 3]
1358
1356
- lengths = resized_h * resized_w + h * (w + 1) + 4 # [B]
1359
+ low_res_use_col_tokens = (
1360
+ image_use_col_tokens
1361
+ if use_single_crop_col_tokens is None
1362
+ else use_single_crop_col_tokens
1363
+ )
1364
+ low_res_extra = int(low_res_use_col_tokens)
1365
+ high_res_extra = int(image_use_col_tokens)
1366
+
1367
+ lengths = (
1368
+ resized_h * (resized_w + low_res_extra) + h * (w + high_res_extra) + 4
1369
+ ) # [B]
1370
total_len = int(lengths.sum().item())
1371
1372
flat = torch.empty(total_len, dtype=torch.long, device=device)
@@ -1363,16 +1376,24 @@ def build_flat_image_bool_length(
1376
resized_h_i, resized_w_i, h_i, w_i = image_grids[i].tolist()
1377
L_i = int(lengths[i].item())
1378
1366
- num_low_res_patches = resized_h_i * resized_w_i
1367
-
1379
idx = offset
1380
1370
- flat[idx] = low_res_image_start_id
1381
+ flat[idx] = (
1382
+ low_res_image_start_id if use_single_crop_start_token else image_start_id
1383
+ )
1384
idx += 1
1385
1373
- if num_low_res_patches > 0:
1374
- flat[idx : idx + num_low_res_patches] = image_patch_id
1375
- idx += num_low_res_patches
1386
+ low_res_block_len = resized_w_i + low_res_extra
1387
+ if low_res_block_len > 0 and resized_h_i > 0:
1388
+ line = torch.empty(low_res_block_len, dtype=torch.long, device=device)
1389
+ if resized_w_i > 0:
1390
+ line[:resized_w_i] = image_patch_id
1391
+ if low_res_use_col_tokens:
1392
+ line[resized_w_i] = image_col_id
1393
+
1394
+ block = line.repeat(resized_h_i)
1395
+ flat[idx : idx + resized_h_i * low_res_block_len] = block
1396
+ idx += resized_h_i * low_res_block_len
1397
1398
flat[idx] = image_end_id
1399
idx += 1
@@ -1380,12 +1401,13 @@ def build_flat_image_bool_length(
1401
flat[idx] = image_start_id
1402
idx += 1
1403
1383
- block_len = w_i + 1
1404
+ block_len = w_i + high_res_extra
1405
if block_len > 0 and h_i > 0:
1406
line = torch.empty(block_len, dtype=torch.long, device=device)
1407
if w_i > 0:
1408
line[:w_i] = image_patch_id
1388
- line[w_i] = image_col_id
1409
+ if image_use_col_tokens:
1410
+ line[w_i] = image_col_id
1411
1412
block = line.repeat(h_i)
1413
flat[idx : idx + h_i * block_len] = block
@@ -2108,7 +2130,13 @@ class Molmo2MultiModalProcessor(BaseMultiModalProcessor[Molmo2ProcessingInfo]):
2130
(
2131
processed_outputs["image_tokens"],
2132
processed_outputs["num_image_tokens"],
2111
- ) = build_flat_image_bool_length(image_grids, hf_config)
2133
+ ) = build_flat_image_bool_length(
2134
+ image_grids,
2135
+ hf_config,
2136
+ image_use_col_tokens=hf_processor.image_use_col_tokens,
2137
+ use_single_crop_col_tokens=hf_processor.use_single_crop_col_tokens,
2138
+ use_single_crop_start_token=hf_processor.use_single_crop_start_token,
2139
+ )
2140
2141
return BatchFeature({**processed_outputs, **all_video_outputs})
2142