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, 02:15 UTC
[Model] Use mm_features to compute mrope positions for PaddleOCR-VL (#39888)
Signed-off-by: grYe99 <guorongye99@gmail.com> Co-authored-by: grYe99 <guorongye99@gmail.com>
grYe99 committed
Apr 16, 2026 at 21:14 UTC
4269b794091c057491b2ee5ac855d5a268959cb7
2 files changed
+266
-94
tests/model_executor/test_paddleocr_vl_mrope.py
new
+210
@@ -0,0 +1,210 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from dataclasses import dataclass, field
4
+
5
+import pytest
6
+import torch
7
+
8
+from vllm.model_executor.models.paddleocr_vl import (
9
+ PaddleOCRVLForConditionalGeneration,
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 DummyVisionConfig:
31
+ spatial_merge_size: int = 2
32
+ patch_size: int = 14
33
+
34
+
35
+@dataclass
36
+class DummyConfig:
37
+ image_token_id: int = 151655
38
+ video_token_id: int = 151654
39
+ vision_start_token_id: int = 151652
40
+ vision_end_token_id: int = 151653
41
+ vision_config: DummyVisionConfig = field(default_factory=DummyVisionConfig)
42
+
43
+
44
+def make_model(config: DummyConfig) -> PaddleOCRVLForConditionalGeneration:
45
+ model = object.__new__(PaddleOCRVLForConditionalGeneration)
46
+ model.config = config
47
+ return model
48
+
49
+
50
+def make_mm_feature(
51
+ *,
52
+ offset: int,
53
+ length: int,
54
+ image_grid_thw: tuple[int, int, int],
55
+) -> MultiModalFeatureSpec:
56
+ return MultiModalFeatureSpec(
57
+ data=MultiModalKwargsItem(
58
+ {
59
+ "image_grid_thw": MultiModalFieldElem(
60
+ data=torch.tensor(image_grid_thw),
61
+ field=None,
62
+ ),
63
+ }
64
+ ),
65
+ modality="image",
66
+ identifier="DUMMY",
67
+ mm_position=PlaceholderRange(offset=offset, length=length),
68
+ )
69
+
70
+
71
+def test_get_mrope_input_positions_text_only():
72
+ model = make_model(DummyConfig())
73
+ input_tokens = [11, 12, 13, 14, 15]
74
+ positions, delta = model.get_mrope_input_positions(
75
+ input_tokens=input_tokens,
76
+ mm_features=[],
77
+ )
78
+ expected = torch.tensor(
79
+ [
80
+ [0, 1, 2, 3, 4],
81
+ [0, 1, 2, 3, 4],
82
+ [0, 1, 2, 3, 4],
83
+ ]
84
+ )
85
+ assert torch.equal(positions, expected)
86
+ assert delta == 0
87
+
88
+
89
+def test_get_mrope_input_positions_single_image():
90
+ model = make_model(DummyConfig())
91
+ spatial_merge_size = model.config.vision_config.spatial_merge_size
92
+
93
+ t, h, w = 1, 2, 2
94
+ num_image_tokens = t * h * w
95
+
96
+ input_tokens = (
97
+ [10]
98
+ + [model.config.vision_start_token_id]
99
+ + [model.config.image_token_id] * num_image_tokens
100
+ + [model.config.vision_end_token_id]
101
+ + [30, 31]
102
+ )
103
+
104
+ mm_features = [
105
+ make_mm_feature(
106
+ offset=2, # 1 (text) + 1 (vision_start)
107
+ length=num_image_tokens,
108
+ image_grid_thw=(t, h * spatial_merge_size, w * spatial_merge_size),
109
+ )
110
+ ]
111
+
112
+ positions, delta = model.get_mrope_input_positions(
113
+ input_tokens=input_tokens,
114
+ mm_features=mm_features,
115
+ )
116
+
117
+ expected = torch.tensor(
118
+ [
119
+ [0, 1, 2, 2, 2, 2, 4, 5, 6],
120
+ [0, 1, 2, 2, 3, 3, 4, 5, 6],
121
+ [0, 1, 2, 3, 2, 3, 4, 5, 6],
122
+ ]
123
+ )
124
+
125
+ assert torch.equal(positions, expected)
126
+ expected_delta = (positions.max().item() + 1) - len(input_tokens)
127
+ assert delta == expected_delta
128
+
129
+
130
+def test_get_mrope_input_positions_multiple_images():
131
+ model = make_model(DummyConfig())
132
+ spatial_merge_size = model.config.vision_config.spatial_merge_size
133
+
134
+ t1, h1, w1 = 1, 2, 2
135
+ num1 = t1 * h1 * w1
136
+
137
+ t2, h2, w2 = 1, 1, 3
138
+ num2 = t2 * h2 * w2
139
+
140
+ input_tokens = (
141
+ [10]
142
+ + [model.config.vision_start_token_id]
143
+ + [model.config.image_token_id] * num1
144
+ + [model.config.vision_end_token_id]
145
+ + [20, 21]
146
+ + [model.config.vision_start_token_id]
147
+ + [model.config.image_token_id] * num2
148
+ + [model.config.vision_end_token_id]
149
+ + [30]
150
+ )
151
+
152
+ mm_features = [
153
+ make_mm_feature(
154
+ offset=2,
155
+ length=num1,
156
+ image_grid_thw=(t1, h1 * spatial_merge_size, w1 * spatial_merge_size),
157
+ ),
158
+ make_mm_feature(
159
+ offset=2 + num1 + 1 + 2 + 1,
160
+ length=num2,
161
+ image_grid_thw=(t2, h2 * spatial_merge_size, w2 * spatial_merge_size),
162
+ ),
163
+ ]
164
+
165
+ positions, delta = model.get_mrope_input_positions(
166
+ input_tokens=input_tokens,
167
+ mm_features=mm_features,
168
+ )
169
+
170
+ assert positions.shape == (3, 15)
171
+ assert not torch.equal(positions[:, 2:6], torch.arange(4).expand(3, 4) + 2)
172
+ assert not torch.equal(positions[:, 10:13], torch.arange(3).expand(3, 3) + 10)
173
+
174
+
175
+def test_get_mrope_input_positions_image_at_start():
176
+ model = make_model(DummyConfig())
177
+ spatial_merge_size = model.config.vision_config.spatial_merge_size
178
+
179
+ t, h, w = 1, 2, 2
180
+ num_tokens = t * h * w
181
+
182
+ input_tokens = (
183
+ [model.config.vision_start_token_id]
184
+ + [model.config.image_token_id] * num_tokens
185
+ + [model.config.vision_end_token_id]
186
+ + [10, 11]
187
+ )
188
+
189
+ mm_features = [
190
+ make_mm_feature(
191
+ offset=1, # start token at index 0
192
+ length=num_tokens,
193
+ image_grid_thw=(t, h * spatial_merge_size, w * spatial_merge_size),
194
+ )
195
+ ]
196
+
197
+ positions, delta = model.get_mrope_input_positions(
198
+ input_tokens=input_tokens,
199
+ mm_features=mm_features,
200
+ )
201
+
202
+ expected = torch.tensor(
203
+ [
204
+ [0, 1, 1, 1, 1, 3, 4, 5],
205
+ [0, 1, 1, 2, 2, 3, 4, 5],
206
+ [0, 1, 2, 1, 2, 3, 4, 5],
207
+ ]
208
+ )
209
+
210
+ assert torch.equal(positions, expected)
vllm/model_executor/models/paddleocr_vl.py
+56
-94
@@ -15,7 +15,7 @@
15
# limitations under the License.
16
17
import math
18
-from collections.abc import Iterable, Mapping, Sequence
18
+from collections.abc import Iterable, Iterator, Mapping, Sequence
19
from functools import partial
20
from typing import Annotated, Literal
21
@@ -1056,121 +1056,83 @@ class PaddleOCRVLForConditionalGeneration(nn.Module, SupportsMultiModal, Support
1056
) -> torch.Tensor | None:
1057
return self.language_model.compute_logits(hidden_states)
1058
1059
+ def iter_mm_grid_thw(
1060
+ self, mm_features: list[MultiModalFeatureSpec]
1061
+ ) -> Iterator[tuple[int, int, int, int, float]]:
1062
+ """
1063
+ Iterate over multimodal features and yield grid information.
1064
+
1065
+ Args:
1066
+ mm_features: List of multimodal feature specifications
1067
+
1068
+ Yields:
1069
+ Tuple of (offset, grid_t, grid_h, grid_w, t_factor) for each frame/image
1070
+ """
1071
+ spatial_merge_size = self.config.vision_config.spatial_merge_size
1072
+ tokens_per_second = getattr(self.config.vision_config, "tokens_per_second", 1.0)
1073
+ for mm_feature in sorted(mm_features, key=lambda f: f.mm_position.offset):
1074
+ offset = mm_feature.mm_position.offset
1075
+ if mm_feature.modality == "image":
1076
+ t, h, w = mm_feature.data["image_grid_thw"].data.tolist()
1077
+ assert t == 1, f"Image must have 1 frame, got {t}"
1078
+ yield offset, 1, h // spatial_merge_size, w // spatial_merge_size, 1.0
1079
+ elif mm_feature.modality == "video":
1080
+ t, h, w = mm_feature.data["video_grid_thw"].data.tolist()
1081
+ second_per_grid_ts = 1.0
1082
+ if mm_feature.data.get("second_per_grid_ts", None):
1083
+ second_per_grid_ts = mm_feature.data[
1084
+ "second_per_grid_ts"
1085
+ ].data.item()
1086
+ t_factor = second_per_grid_ts * tokens_per_second
1087
+ yield (
1088
+ offset,
1089
+ t,
1090
+ h // spatial_merge_size,
1091
+ w // spatial_merge_size,
1092
+ t_factor,
1093
+ )
1094
+ else:
1095
+ raise ValueError(f"Unsupported modality: {mm_feature.modality}")
1096
+
1097
def get_mrope_input_positions(
1098
self,
1099
input_tokens: list[int],
1100
mm_features: list[MultiModalFeatureSpec],
1101
) -> tuple[torch.Tensor, int]:
1064
- kwargs = MultiModalFeatureSpec.gather_kwargs(
1065
- mm_features,
1066
- {"image_grid_thw", "video_grid_thw", "second_per_grid_ts"},
1067
- )
1068
- image_grid_thw = [item.tolist() for item in kwargs.get("image_grid_thw", [])]
1069
- video_grid_thw = [item.tolist() for item in kwargs.get("video_grid_thw", [])]
1070
- second_per_grid_ts = kwargs.get("second_per_grid_ts", [])
1071
-
1072
- hf_config = self.config
1073
- image_token_id = hf_config.image_token_id
1074
- video_token_id = hf_config.video_token_id
1075
- vision_start_token_id = hf_config.vision_start_token_id
1076
- spatial_merge_size = hf_config.vision_config.spatial_merge_size
1077
- tokens_per_second = getattr(hf_config.vision_config, "tokens_per_second", 1.0)
1078
-
1079
- input_tokens_tensor = torch.tensor(input_tokens)
1080
- vision_start_indices = torch.argwhere(
1081
- input_tokens_tensor == vision_start_token_id
1082
- ).squeeze(1)
1083
- vision_tokens = input_tokens_tensor[vision_start_indices + 1]
1084
- image_nums = (vision_tokens == image_token_id).sum()
1085
- video_nums = (vision_tokens == video_token_id).sum()
1102
llm_pos_ids_list: list = []
1087
-
1103
st = 0
1089
- remain_images, remain_videos = image_nums, video_nums
1090
-
1091
- image_index, video_index = 0, 0
1092
- for _ in range(image_nums + video_nums):
1093
- video_second_per_grid_t = 0.0
1094
- if remain_images > 0:
1095
- try:
1096
- ed_image = input_tokens.index(image_token_id, st)
1097
- except ValueError:
1098
- ed_image = len(input_tokens) + 1
1099
- else:
1100
- ed_image = len(input_tokens) + 1
1101
- if remain_videos > 0:
1102
- try:
1103
- ed_video = input_tokens.index(video_token_id, st)
1104
- except ValueError:
1105
- ed_video = len(input_tokens) + 1
1106
- else:
1107
- ed_video = len(input_tokens) + 1
1108
- if ed_image < ed_video:
1109
- t, h, w = image_grid_thw[image_index]
1110
- image_index += 1
1111
- remain_images -= 1
1112
- ed = ed_image
1113
- else:
1114
- t, h, w = video_grid_thw[video_index]
1115
- video_second_per_grid_t = 1.0
1116
- if second_per_grid_ts:
1117
- video_second_per_grid_t = second_per_grid_ts[video_index]
1118
- video_index += 1
1119
- remain_videos -= 1
1120
- ed = ed_video
1121
-
1122
- llm_grid_t, llm_grid_h, llm_grid_w = (
1123
- t,
1124
- h // spatial_merge_size,
1125
- w // spatial_merge_size,
1126
- )
1127
- text_len = ed - st
1104
1105
+ for (
1106
+ offset,
1107
+ llm_grid_t,
1108
+ llm_grid_h,
1109
+ llm_grid_w,
1110
+ t_factor,
1111
+ ) in self.iter_mm_grid_thw(mm_features):
1112
+ text_len = offset - st
1113
st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
1114
llm_pos_ids_list.append(
1131
- torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx
1115
+ np.broadcast_to(np.arange(text_len), (3, text_len)) + st_idx
1116
)
1117
1134
- t_index = (
1135
- (
1136
- torch.arange(llm_grid_t)
1137
- .view(-1, 1)
1138
- .expand(-1, llm_grid_h * llm_grid_w)
1139
- * video_second_per_grid_t
1140
- * tokens_per_second
1141
- )
1142
- .long()
1143
- .flatten()
1144
- )
1118
+ grid_indices = np.indices((llm_grid_t, llm_grid_h, llm_grid_w))
1119
+ if t_factor != 1.0:
1120
+ grid_indices[0] = (grid_indices[0] * t_factor).astype(np.int64)
1121
1146
- h_index = (
1147
- torch.arange(llm_grid_h)
1148
- .view(1, -1, 1)
1149
- .expand(llm_grid_t, -1, llm_grid_w)
1150
- .flatten()
1151
- )
1152
- w_index = (
1153
- torch.arange(llm_grid_w)
1154
- .view(1, 1, -1)
1155
- .expand(llm_grid_t, llm_grid_h, -1)
1156
- .flatten()
1157
- )
1158
- llm_pos_ids_list.append(
1159
- torch.stack([t_index, h_index, w_index]) + text_len + st_idx
1160
- )
1161
- st = ed + llm_grid_t * llm_grid_h * llm_grid_w
1122
+ llm_pos_ids_list.append(grid_indices.reshape(3, -1) + text_len + st_idx)
1123
+ st = offset + llm_grid_t * llm_grid_h * llm_grid_w
1124
1125
if st < len(input_tokens):
1126
st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
1127
text_len = len(input_tokens) - st
1128
llm_pos_ids_list.append(
1167
- torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx
1129
+ np.broadcast_to(np.arange(text_len), (3, text_len)) + st_idx
1130
)
1131
1170
- llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1)
1132
+ llm_positions = np.concatenate(llm_pos_ids_list, axis=1).reshape(3, -1)
1133
mrope_position_delta = (llm_positions.max() + 1 - len(input_tokens)).item()
1134
1173
- return llm_positions, mrope_position_delta
1135
+ return torch.from_numpy(llm_positions), mrope_position_delta
1136
1137
def _parse_and_validate_image_input(
1138
self, **kwargs: object