obscura/vllm
2
0
Fork 0
mirror of https://github.com/vllm-project/vllm.git synced 2026-08-14 01:38:11 +00:00
Code Issues Packages Projects Releases Wiki Activity

[Model] Use mm_features to compute mrope positions for PaddleOCR-VL (#39888)

Browse Source 
Signed-off-by: grYe99 <[email protected]>
Co-authored-by: grYe99 <[email protected]>
This commit is contained in:
grYe99
2026-04-16 06:14:00 -07:00
committed by GitHub
co-authored by grYe99
parent edc3648966
commit 4269b79409
2 changed files with 266 additions and 94 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/model_executor/test_paddleocr_vl_mrope.py
+210
View File
@@ -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)
vllm/model_executor/models/paddleocr_vl.py
+56 -94
View File
@@ -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