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 5, 04:45 UTC
[Model] Use mm_features for Ernie-4.5 VL M-RoPE (#39753)
Signed-off-by: Lalit Laxminarayan Bangad <lalitbangad@gmail.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>
lalit10 committed
Apr 14, 2026 at 01:11 UTC
0008729abfbd380d29e0a67a882c8f01220acd73
2 files changed
+196
-123
tests/model_executor/test_ernie45_vl_mrope.py
new
+143
@@ -0,0 +1,143 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from dataclasses import dataclass
4
+
5
+import pytest
6
+import torch
7
+
8
+from vllm.model_executor.models.ernie45_vl import (
9
+ Ernie4_5_VLMoeForConditionalGeneration,
10
+)
11
+from vllm.multimodal.inputs import (
12
+ MultiModalFeatureSpec,
13
+ MultiModalFieldElem,
14
+ MultiModalKwargsItem,
15
+ PlaceholderRange,
16
+)
17
+
18
+pytestmark = pytest.mark.skip_global_cleanup
19
+
20
+
21
+@pytest.fixture(autouse=True, scope="module")
22
+def _force_cpu_default_device():
23
+ original = torch.get_default_device()
24
+ torch.set_default_device("cpu")
25
+ yield
26
+ torch.set_default_device(original)
27
+
28
+
29
+@dataclass
30
+class DummyConfig:
31
+ spatial_conv_size: int = 2
32
+ temporal_conv_size: int = 2
33
+
34
+
35
+def make_model(config: DummyConfig) -> Ernie4_5_VLMoeForConditionalGeneration:
36
+ model = object.__new__(Ernie4_5_VLMoeForConditionalGeneration)
37
+ model.config = config
38
+ return model
39
+
40
+
41
+def make_mm_feature(
42
+ *,
43
+ modality: str,
44
+ offset: int,
45
+ length: int,
46
+ grid_thw: tuple[int, int, int],
47
+) -> MultiModalFeatureSpec:
48
+ field_name = "image_grid_thw" if modality == "image" else "video_grid_thw"
49
+ return MultiModalFeatureSpec(
50
+ data=MultiModalKwargsItem(
51
+ {
52
+ field_name: MultiModalFieldElem(
53
+ data=torch.tensor(grid_thw),
54
+ field=None, # HACK.
55
+ ),
56
+ }
57
+ ),
58
+ modality=modality,
59
+ identifier="DUMMY",
60
+ mm_position=PlaceholderRange(offset=offset, length=length),
61
+ )
62
+
63
+
64
+def test_get_mrope_input_positions_text_only():
65
+ model = make_model(DummyConfig())
66
+
67
+ positions, delta = model.get_mrope_input_positions(
68
+ input_tokens=[11, 12, 13, 14, 15],
69
+ mm_features=[],
70
+ )
71
+
72
+ expected = torch.tensor(
73
+ [
74
+ [0, 1, 2, 3, 4],
75
+ [0, 1, 2, 3, 4],
76
+ [0, 1, 2, 3, 4],
77
+ ]
78
+ )
79
+
80
+ assert torch.equal(positions, expected)
81
+ assert delta == 0
82
+
83
+
84
+def test_get_mrope_input_positions_single_image():
85
+ model = make_model(DummyConfig())
86
+ mm_features = [
87
+ make_mm_feature(
88
+ modality="image",
89
+ offset=1,
90
+ length=4,
91
+ grid_thw=(1, 4, 4),
92
+ )
93
+ ]
94
+
95
+ positions, delta = model.get_mrope_input_positions(
96
+ input_tokens=[10, 20, 21, 22, 23, 30, 31],
97
+ mm_features=mm_features,
98
+ )
99
+
100
+ expected = torch.tensor(
101
+ [
102
+ [0, 1, 1, 1, 1, 3, 4],
103
+ [0, 1, 1, 2, 2, 3, 4],
104
+ [0, 1, 2, 1, 2, 3, 4],
105
+ ]
106
+ )
107
+
108
+ assert torch.equal(positions, expected)
109
+ assert delta == -2
110
+
111
+
112
+def test_get_mrope_input_positions_interleaved_image_and_video():
113
+ model = make_model(DummyConfig())
114
+ mm_features = [
115
+ make_mm_feature(
116
+ modality="image",
117
+ offset=1,
118
+ length=4,
119
+ grid_thw=(1, 4, 4),
120
+ ),
121
+ make_mm_feature(
122
+ modality="video",
123
+ offset=7,
124
+ length=2,
125
+ grid_thw=(2, 4, 2),
126
+ ),
127
+ ]
128
+
129
+ positions, delta = model.get_mrope_input_positions(
130
+ input_tokens=[10, 20, 21, 22, 23, 30, 31, 40, 41, 50, 51],
131
+ mm_features=mm_features,
132
+ )
133
+
134
+ expected = torch.tensor(
135
+ [
136
+ [0, 1, 1, 1, 1, 3, 4, 5, 5, 7, 8],
137
+ [0, 1, 1, 2, 2, 3, 4, 5, 6, 7, 8],
138
+ [0, 1, 2, 1, 2, 3, 4, 5, 5, 7, 8],
139
+ ]
140
+ )
141
+
142
+ assert torch.equal(positions, expected)
143
+ assert delta == -2
vllm/model_executor/models/ernie45_vl.py
+53
-123
@@ -23,9 +23,8 @@
23
# limitations under the License.
24
"""Inference-only Ernie VL model compatible with HuggingFace weights."""
25
26
-import itertools
26
import math
28
-from collections.abc import Callable, Iterable, Mapping, Sequence
27
+from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence
28
from functools import partial
29
from typing import Annotated, Any, Literal
30
@@ -1401,131 +1400,62 @@ class Ernie4_5_VLMoeForConditionalGeneration(
1400
input_tokens: list[int],
1401
mm_features: list[MultiModalFeatureSpec],
1402
) -> tuple[torch.Tensor, int]:
1404
- kwargs = MultiModalFeatureSpec.gather_kwargs(
1405
- mm_features,
1406
- {"image_grid_thw", "video_grid_thw"},
1407
- )
1408
- image_grid_thw = [item.tolist() for item in kwargs.get("image_grid_thw", [])]
1409
- video_grid_thw = [item.tolist() for item in kwargs.get("video_grid_thw", [])]
1410
-
1411
- hf_config = self.config
1412
- image_token_id = hf_config.im_patch_id
1413
- video_start_token_id = hf_config.video_start_token_id
1414
- video_end_token_id = hf_config.video_end_token_id
1415
- spatial_conv_size = hf_config.spatial_conv_size
1416
- temporal_conv_size = hf_config.temporal_conv_size
1403
llm_pos_ids_list: list = []
1404
+ st = 0
1405
+
1406
+ for (
1407
+ offset,
1408
+ llm_grid_t,
1409
+ llm_grid_h,
1410
+ llm_grid_w,
1411
+ ) in self.iter_mm_grid_thw(mm_features):
1412
+ text_len = offset - st
1413
+ st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
1414
+ llm_pos_ids_list.append(
1415
+ np.broadcast_to(np.arange(text_len), (3, text_len)) + st_idx
1416
+ )
1417
1419
- if image_grid_thw or video_grid_thw:
1420
- input_token_type: list[str] = []
1421
- video_check_flg = False
1422
- for token in input_tokens:
1423
- if token == video_start_token_id:
1424
- video_check_flg = True
1425
- elif token == video_end_token_id:
1426
- video_check_flg = False
1427
-
1428
- if (token == image_token_id) and (video_check_flg is False):
1429
- input_token_type.append("image")
1430
- elif (token == image_token_id) and (video_check_flg is True):
1431
- input_token_type.append("video")
1432
- else:
1433
- input_token_type.append("text")
1434
-
1435
- input_type_group: list[tuple[str, int, int]] = []
1436
- for key, group_iter in itertools.groupby(
1437
- enumerate(input_token_type), lambda x: x[1]
1438
- ):
1439
- group_list = list(group_iter)
1440
- start_index = group_list[0][0]
1441
- end_index = group_list[-1][0] + 1
1442
- input_type_group.append((key, start_index, end_index))
1443
-
1444
- video_frame_num = 1
1445
- mm_data_idx = 0
1446
- for modality_type, start_idx, end_idx in input_type_group:
1447
- st_idx = (
1448
- llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
1449
- )
1450
- if modality_type == "image":
1451
- t, h, w = image_grid_thw[mm_data_idx]
1452
- llm_grid_t, llm_grid_h, llm_grid_w = (
1453
- t,
1454
- h // spatial_conv_size,
1455
- w // spatial_conv_size,
1456
- )
1457
-
1458
- t_index = (
1459
- torch.arange(llm_grid_t)
1460
- .view(-1, 1)
1461
- .expand(-1, llm_grid_h * llm_grid_w)
1462
- .flatten()
1463
- )
1464
- h_index = (
1465
- torch.arange(llm_grid_h)
1466
- .view(1, -1, 1)
1467
- .expand(llm_grid_t, -1, llm_grid_w)
1468
- .flatten()
1469
- )
1470
- w_index = (
1471
- torch.arange(llm_grid_w)
1472
- .view(1, 1, -1)
1473
- .expand(llm_grid_t, llm_grid_h, -1)
1474
- .flatten()
1475
- )
1476
- llm_pos_ids_list.append(
1477
- torch.stack([t_index, h_index, w_index]) + st_idx
1478
- )
1479
- mm_data_idx += 1
1480
-
1481
- elif modality_type == "video":
1482
- t, h, w = video_grid_thw[mm_data_idx]
1483
- llm_grid_t, llm_grid_h, llm_grid_w = (
1484
- t // temporal_conv_size,
1485
- h // spatial_conv_size,
1486
- w // spatial_conv_size,
1487
- )
1488
-
1489
- for t_idx in range(llm_grid_t):
1490
- t_index = (
1491
- torch.tensor(t_idx)
1492
- .view(-1, 1)
1493
- .expand(-1, llm_grid_h * llm_grid_w)
1494
- .flatten()
1495
- )
1496
- h_index = (
1497
- torch.arange(llm_grid_h)
1498
- .view(1, -1, 1)
1499
- .expand(1, -1, llm_grid_w)
1500
- .flatten()
1501
- )
1502
- w_index = (
1503
- torch.arange(llm_grid_w)
1504
- .view(1, 1, -1)
1505
- .expand(1, llm_grid_h, -1)
1506
- .flatten()
1507
- )
1508
- llm_pos_ids_list.append(
1509
- torch.stack([t_index, h_index, w_index]) + st_idx
1510
- )
1511
-
1512
- mm_data_idx += 1
1513
- video_frame_num += 1
1514
-
1515
- else:
1516
- text_len = end_idx - start_idx
1517
- llm_pos_ids_list.append(
1518
- torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx
1519
- )
1520
- video_frame_num = 1
1521
-
1522
- else:
1523
- text_len = len(input_tokens)
1524
- llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1))
1418
+ grid_indices = np.indices((llm_grid_t, llm_grid_h, llm_grid_w)).reshape(
1419
+ 3, -1
1420
+ )
1421
+ llm_pos_ids_list.append(grid_indices + text_len + st_idx)
1422
+ st = offset + llm_grid_t * llm_grid_h * llm_grid_w
1423
+
1424
+ if st < len(input_tokens):
1425
+ text_len = len(input_tokens) - st
1426
+ st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
1427
+ llm_pos_ids_list.append(
1428
+ np.broadcast_to(np.arange(text_len), (3, text_len)) + st_idx
1429
+ )
1430
1526
- llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1)
1431
+ llm_positions = np.concatenate(llm_pos_ids_list, axis=1).reshape(3, -1)
1432
mrope_position_delta = (llm_positions.max() + 1 - len(input_tokens)).item()
1528
- return llm_positions, mrope_position_delta
1433
+ return torch.from_numpy(llm_positions), mrope_position_delta
1434
+
1435
+ def iter_mm_grid_thw(
1436
+ self, mm_features: list[MultiModalFeatureSpec]
1437
+ ) -> Iterator[tuple[int, int, int, int]]:
1438
+ spatial_conv_size = self.config.spatial_conv_size
1439
+ temporal_conv_size = self.config.temporal_conv_size
1440
+
1441
+ for mm_feature in sorted(mm_features, key=lambda f: f.mm_position.offset):
1442
+ if mm_feature.data is None:
1443
+ raise ValueError("M-RoPE calculation requires multimodal feature data")
1444
+
1445
+ offset = mm_feature.mm_position.offset
1446
+ if mm_feature.modality == "image":
1447
+ t, h, w = mm_feature.data["image_grid_thw"].data.tolist()
1448
+ yield offset, t, h // spatial_conv_size, w // spatial_conv_size
1449
+ elif mm_feature.modality == "video":
1450
+ t, h, w = mm_feature.data["video_grid_thw"].data.tolist()
1451
+ yield (
1452
+ offset,
1453
+ t // temporal_conv_size,
1454
+ h // spatial_conv_size,
1455
+ w // spatial_conv_size,
1456
+ )
1457
+ else:
1458
+ raise ValueError(f"Unsupported modality: {mm_feature.modality}")
1459
1460
def _parse_and_validate_image_input(
1461
self, **kwargs: object