From 4269b794091c057491b2ee5ac855d5a268959cb7 Mon Sep 17 00:00:00 2001 From: grYe99 <35474918+grYe99@users.noreply.github.com> Date: Thu, 16 Apr 2026 21:14:00 +0800 Subject: [PATCH] [Model] Use mm_features to compute mrope positions for PaddleOCR-VL (#39888) Signed-off-by: grYe99 Co-authored-by: grYe99 --- .../model_executor/test_paddleocr_vl_mrope.py | 210 ++++++++++++++++++ vllm/model_executor/models/paddleocr_vl.py | 150 +++++-------- 2 files changed, 266 insertions(+), 94 deletions(-) create mode 100644 tests/model_executor/test_paddleocr_vl_mrope.py diff --git a/tests/model_executor/test_paddleocr_vl_mrope.py b/tests/model_executor/test_paddleocr_vl_mrope.py new file mode 100644 index 00000000000..9090a9254d0 --- /dev/null +++ b/tests/model_executor/test_paddleocr_vl_mrope.py @@ -0,0 +1,210 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from dataclasses import dataclass, field + +import pytest +import torch + +from vllm.model_executor.models.paddleocr_vl import ( + PaddleOCRVLForConditionalGeneration, +) +from vllm.multimodal.inputs import ( + MultiModalFeatureSpec, + MultiModalFieldElem, + MultiModalKwargsItem, + PlaceholderRange, +) + +pytestmark = pytest.mark.skip_global_cleanup + + +@pytest.fixture(autouse=True, scope="module") +def _force_cpu_default_device(): + original = torch.get_default_device() + torch.set_default_device("cpu") + yield + torch.set_default_device(original) + + +@dataclass +class DummyVisionConfig: + spatial_merge_size: int = 2 + patch_size: int = 14 + + +@dataclass +class DummyConfig: + image_token_id: int = 151655 + video_token_id: int = 151654 + vision_start_token_id: int = 151652 + vision_end_token_id: int = 151653 + vision_config: DummyVisionConfig = field(default_factory=DummyVisionConfig) + + +def make_model(config: DummyConfig) -> PaddleOCRVLForConditionalGeneration: + model = object.__new__(PaddleOCRVLForConditionalGeneration) + model.config = config + return model + + +def make_mm_feature( + *, + offset: int, + length: int, + image_grid_thw: tuple[int, int, int], +) -> MultiModalFeatureSpec: + return MultiModalFeatureSpec( + data=MultiModalKwargsItem( + { + "image_grid_thw": MultiModalFieldElem( + data=torch.tensor(image_grid_thw), + field=None, + ), + } + ), + modality="image", + identifier="DUMMY", + mm_position=PlaceholderRange(offset=offset, length=length), + ) + + +def test_get_mrope_input_positions_text_only(): + model = make_model(DummyConfig()) + input_tokens = [11, 12, 13, 14, 15] + positions, delta = model.get_mrope_input_positions( + input_tokens=input_tokens, + mm_features=[], + ) + expected = torch.tensor( + [ + [0, 1, 2, 3, 4], + [0, 1, 2, 3, 4], + [0, 1, 2, 3, 4], + ] + ) + assert torch.equal(positions, expected) + assert delta == 0 + + +def test_get_mrope_input_positions_single_image(): + model = make_model(DummyConfig()) + spatial_merge_size = model.config.vision_config.spatial_merge_size + + t, h, w = 1, 2, 2 + num_image_tokens = t * h * w + + input_tokens = ( + [10] + + [model.config.vision_start_token_id] + + [model.config.image_token_id] * num_image_tokens + + [model.config.vision_end_token_id] + + [30, 31] + ) + + mm_features = [ + make_mm_feature( + offset=2, # 1 (text) + 1 (vision_start) + length=num_image_tokens, + image_grid_thw=(t, h * spatial_merge_size, w * spatial_merge_size), + ) + ] + + positions, delta = model.get_mrope_input_positions( + input_tokens=input_tokens, + mm_features=mm_features, + ) + + expected = torch.tensor( + [ + [0, 1, 2, 2, 2, 2, 4, 5, 6], + [0, 1, 2, 2, 3, 3, 4, 5, 6], + [0, 1, 2, 3, 2, 3, 4, 5, 6], + ] + ) + + assert torch.equal(positions, expected) + expected_delta = (positions.max().item() + 1) - len(input_tokens) + assert delta == expected_delta + + +def test_get_mrope_input_positions_multiple_images(): + model = make_model(DummyConfig()) + spatial_merge_size = model.config.vision_config.spatial_merge_size + + t1, h1, w1 = 1, 2, 2 + num1 = t1 * h1 * w1 + + t2, h2, w2 = 1, 1, 3 + num2 = t2 * h2 * w2 + + input_tokens = ( + [10] + + [model.config.vision_start_token_id] + + [model.config.image_token_id] * num1 + + [model.config.vision_end_token_id] + + [20, 21] + + [model.config.vision_start_token_id] + + [model.config.image_token_id] * num2 + + [model.config.vision_end_token_id] + + [30] + ) + + mm_features = [ + make_mm_feature( + offset=2, + length=num1, + image_grid_thw=(t1, h1 * spatial_merge_size, w1 * spatial_merge_size), + ), + make_mm_feature( + offset=2 + num1 + 1 + 2 + 1, + length=num2, + image_grid_thw=(t2, h2 * spatial_merge_size, w2 * spatial_merge_size), + ), + ] + + positions, delta = model.get_mrope_input_positions( + input_tokens=input_tokens, + mm_features=mm_features, + ) + + assert positions.shape == (3, 15) + assert not torch.equal(positions[:, 2:6], torch.arange(4).expand(3, 4) + 2) + assert not torch.equal(positions[:, 10:13], torch.arange(3).expand(3, 3) + 10) + + +def test_get_mrope_input_positions_image_at_start(): + model = make_model(DummyConfig()) + spatial_merge_size = model.config.vision_config.spatial_merge_size + + t, h, w = 1, 2, 2 + num_tokens = t * h * w + + input_tokens = ( + [model.config.vision_start_token_id] + + [model.config.image_token_id] * num_tokens + + [model.config.vision_end_token_id] + + [10, 11] + ) + + mm_features = [ + make_mm_feature( + offset=1, # start token at index 0 + length=num_tokens, + image_grid_thw=(t, h * spatial_merge_size, w * spatial_merge_size), + ) + ] + + positions, delta = model.get_mrope_input_positions( + input_tokens=input_tokens, + mm_features=mm_features, + ) + + expected = torch.tensor( + [ + [0, 1, 1, 1, 1, 3, 4, 5], + [0, 1, 1, 2, 2, 3, 4, 5], + [0, 1, 2, 1, 2, 3, 4, 5], + ] + ) + + assert torch.equal(positions, expected) diff --git a/vllm/model_executor/models/paddleocr_vl.py b/vllm/model_executor/models/paddleocr_vl.py index 48a285bc0f0..8eaf94620c1 100644 --- a/vllm/model_executor/models/paddleocr_vl.py +++ b/vllm/model_executor/models/paddleocr_vl.py @@ -15,7 +15,7 @@ # limitations under the License. import math -from collections.abc import Iterable, Mapping, Sequence +from collections.abc import Iterable, Iterator, Mapping, Sequence from functools import partial from typing import Annotated, Literal @@ -1056,121 +1056,83 @@ class PaddleOCRVLForConditionalGeneration(nn.Module, SupportsMultiModal, Support ) -> torch.Tensor | None: return self.language_model.compute_logits(hidden_states) + def iter_mm_grid_thw( + self, mm_features: list[MultiModalFeatureSpec] + ) -> Iterator[tuple[int, int, int, int, float]]: + """ + Iterate over multimodal features and yield grid information. + + Args: + mm_features: List of multimodal feature specifications + + Yields: + Tuple of (offset, grid_t, grid_h, grid_w, t_factor) for each frame/image + """ + spatial_merge_size = self.config.vision_config.spatial_merge_size + tokens_per_second = getattr(self.config.vision_config, "tokens_per_second", 1.0) + for mm_feature in sorted(mm_features, key=lambda f: f.mm_position.offset): + offset = mm_feature.mm_position.offset + if mm_feature.modality == "image": + t, h, w = mm_feature.data["image_grid_thw"].data.tolist() + assert t == 1, f"Image must have 1 frame, got {t}" + yield offset, 1, h // spatial_merge_size, w // spatial_merge_size, 1.0 + elif mm_feature.modality == "video": + t, h, w = mm_feature.data["video_grid_thw"].data.tolist() + second_per_grid_ts = 1.0 + if mm_feature.data.get("second_per_grid_ts", None): + second_per_grid_ts = mm_feature.data[ + "second_per_grid_ts" + ].data.item() + t_factor = second_per_grid_ts * tokens_per_second + yield ( + offset, + t, + h // spatial_merge_size, + w // spatial_merge_size, + t_factor, + ) + else: + raise ValueError(f"Unsupported modality: {mm_feature.modality}") + def get_mrope_input_positions( self, input_tokens: list[int], mm_features: list[MultiModalFeatureSpec], ) -> tuple[torch.Tensor, int]: - kwargs = MultiModalFeatureSpec.gather_kwargs( - mm_features, - {"image_grid_thw", "video_grid_thw", "second_per_grid_ts"}, - ) - image_grid_thw = [item.tolist() for item in kwargs.get("image_grid_thw", [])] - video_grid_thw = [item.tolist() for item in kwargs.get("video_grid_thw", [])] - second_per_grid_ts = kwargs.get("second_per_grid_ts", []) - - hf_config = self.config - image_token_id = hf_config.image_token_id - video_token_id = hf_config.video_token_id - vision_start_token_id = hf_config.vision_start_token_id - spatial_merge_size = hf_config.vision_config.spatial_merge_size - tokens_per_second = getattr(hf_config.vision_config, "tokens_per_second", 1.0) - - input_tokens_tensor = torch.tensor(input_tokens) - vision_start_indices = torch.argwhere( - input_tokens_tensor == vision_start_token_id - ).squeeze(1) - vision_tokens = input_tokens_tensor[vision_start_indices + 1] - image_nums = (vision_tokens == image_token_id).sum() - video_nums = (vision_tokens == video_token_id).sum() llm_pos_ids_list: list = [] - st = 0 - remain_images, remain_videos = image_nums, video_nums - - image_index, video_index = 0, 0 - for _ in range(image_nums + video_nums): - video_second_per_grid_t = 0.0 - if remain_images > 0: - try: - ed_image = input_tokens.index(image_token_id, st) - except ValueError: - ed_image = len(input_tokens) + 1 - else: - ed_image = len(input_tokens) + 1 - if remain_videos > 0: - try: - ed_video = input_tokens.index(video_token_id, st) - except ValueError: - ed_video = len(input_tokens) + 1 - else: - ed_video = len(input_tokens) + 1 - if ed_image < ed_video: - t, h, w = image_grid_thw[image_index] - image_index += 1 - remain_images -= 1 - ed = ed_image - else: - t, h, w = video_grid_thw[video_index] - video_second_per_grid_t = 1.0 - if second_per_grid_ts: - video_second_per_grid_t = second_per_grid_ts[video_index] - video_index += 1 - remain_videos -= 1 - ed = ed_video - - llm_grid_t, llm_grid_h, llm_grid_w = ( - t, - h // spatial_merge_size, - w // spatial_merge_size, - ) - text_len = ed - st + for ( + offset, + llm_grid_t, + llm_grid_h, + llm_grid_w, + t_factor, + ) in self.iter_mm_grid_thw(mm_features): + text_len = offset - st st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 llm_pos_ids_list.append( - torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx + np.broadcast_to(np.arange(text_len), (3, text_len)) + st_idx ) - t_index = ( - ( - torch.arange(llm_grid_t) - .view(-1, 1) - .expand(-1, llm_grid_h * llm_grid_w) - * video_second_per_grid_t - * tokens_per_second - ) - .long() - .flatten() - ) + grid_indices = np.indices((llm_grid_t, llm_grid_h, llm_grid_w)) + if t_factor != 1.0: + grid_indices[0] = (grid_indices[0] * t_factor).astype(np.int64) - h_index = ( - torch.arange(llm_grid_h) - .view(1, -1, 1) - .expand(llm_grid_t, -1, llm_grid_w) - .flatten() - ) - w_index = ( - torch.arange(llm_grid_w) - .view(1, 1, -1) - .expand(llm_grid_t, llm_grid_h, -1) - .flatten() - ) - llm_pos_ids_list.append( - torch.stack([t_index, h_index, w_index]) + text_len + st_idx - ) - st = ed + llm_grid_t * llm_grid_h * llm_grid_w + llm_pos_ids_list.append(grid_indices.reshape(3, -1) + text_len + st_idx) + st = offset + llm_grid_t * llm_grid_h * llm_grid_w if st < len(input_tokens): st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 text_len = len(input_tokens) - st llm_pos_ids_list.append( - torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx + np.broadcast_to(np.arange(text_len), (3, text_len)) + st_idx ) - llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1) + llm_positions = np.concatenate(llm_pos_ids_list, axis=1).reshape(3, -1) mrope_position_delta = (llm_positions.max() + 1 - len(input_tokens)).item() - return llm_positions, mrope_position_delta + return torch.from_numpy(llm_positions), mrope_position_delta def _parse_and_validate_image_input( self, **kwargs: object