[Bugfix][Model] Fix M-RoPE offset double-count in Qwen3-Omni (#58890)

Signed-off-by: C1ves <3121985758@qq.com>
This commit is contained in:
C
2026-10-03 09:00:51 +08:00
committed by GitHub
parent 8c3b720a4d
commit 44198f577f
2 changed files with 202 additions and 8 deletions
@@ -0,0 +1,198 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import dataclasses
from dataclasses import dataclass
import pytest
import torch
from vllm.model_executor.models.qwen3_omni_moe_thinker import (
Qwen3OmniMoeThinkerForConditionalGeneration,
)
from vllm.multimodal.inputs import (
MultiModalFeatureSpec,
MultiModalFieldElem,
MultiModalKwargsItem,
PlaceholderRange,
)
@pytest.fixture(autouse=True, scope="module")
def _force_cpu_default_device():
# get_mrope_input_positions returns CPU tensors (via torch.from_numpy).
# Ensure the default device is CPU so the rest of the test tensors match.
original = torch.get_default_device()
torch.set_default_device("cpu")
yield
torch.set_default_device(original)
TEXT_TOKEN_ID = 1
VISION_START_TOKEN_ID = 100
VISION_END_TOKEN_ID = 101
IMAGE_TOKEN_ID = 200
VIDEO_TOKEN_ID = 201
AUDIO_TOKEN_ID = 202
AUDIO_START_TOKEN_ID = 300
AUDIO_END_TOKEN_ID = 301
@dataclass
class DummyVisionConfig:
spatial_merge_size: int = 1
@dataclass
class DummyConfig:
position_id_per_seconds: int = 25
vision_config: DummyVisionConfig = dataclasses.field(
default_factory=DummyVisionConfig
)
def make_model() -> Qwen3OmniMoeThinkerForConditionalGeneration:
model = Qwen3OmniMoeThinkerForConditionalGeneration.__new__(
Qwen3OmniMoeThinkerForConditionalGeneration
)
model.config = DummyConfig()
return model
def _elem(data):
return MultiModalFieldElem(
data=torch.tensor(data),
field=None, # HACK.
)
def make_image_feature(offset: int, grid_thw: tuple[int, int, int]):
return MultiModalFeatureSpec(
data=MultiModalKwargsItem({"image_grid_thw": _elem(grid_thw)}),
modality="image",
identifier="DUMMY",
mm_position=PlaceholderRange(
offset=offset, length=int(torch.tensor(grid_thw).prod())
),
)
def make_video_feature(
offset: int, grid_thw: tuple[int, int, int], use_audio_in_video: bool = False
):
data = {"video_grid_thw": _elem(grid_thw)}
if use_audio_in_video:
data["use_audio_in_video"] = _elem(True)
return MultiModalFeatureSpec(
data=MultiModalKwargsItem(data),
modality="video",
identifier="DUMMY",
mm_position=PlaceholderRange(
offset=offset, length=int(torch.tensor(grid_thw).prod())
),
)
def make_audio_feature(offset: int, audio_feature_length: int):
return MultiModalFeatureSpec(
data=MultiModalKwargsItem(
{"audio_feature_lengths": _elem(audio_feature_length)}
),
modality="audio",
identifier="DUMMY",
mm_position=PlaceholderRange(offset=offset, length=0),
)
def test_mrope_positions_image_at_prompt_end():
"""A prompt ending with <|vision_end|> must still produce seq_len positions."""
model = make_model()
# [<|vision_start|>, <|image_pad|> x 4, <|vision_end|>]
input_tokens = (
[VISION_START_TOKEN_ID] + [IMAGE_TOKEN_ID] * 4 + [VISION_END_TOKEN_ID]
)
mm_features = [make_image_feature(offset=1, grid_thw=(1, 2, 2))]
positions, _ = model.get_mrope_input_positions(input_tokens, mm_features)
assert positions.shape == (3, len(input_tokens))
def test_mrope_positions_audio_at_prompt_end():
model = make_model()
# audio_feature_length=16 -> 2 audio tokens.
# [<|audio_start|>, <|audio_pad|> x 2, <|audio_end|>]
input_tokens = [AUDIO_START_TOKEN_ID] + [AUDIO_TOKEN_ID] * 2 + [AUDIO_END_TOKEN_ID]
mm_features = [make_audio_feature(offset=1, audio_feature_length=16)]
positions, _ = model.get_mrope_input_positions(input_tokens, mm_features)
assert positions.shape == (3, len(input_tokens))
def test_mrope_positions_video_at_prompt_end():
model = make_model()
# [T, <|vision_start|>, <|video_pad|> x 8, <|vision_end|>]
input_tokens = (
[TEXT_TOKEN_ID]
+ [VISION_START_TOKEN_ID]
+ [VIDEO_TOKEN_ID] * 8
+ [VISION_END_TOKEN_ID]
)
mm_features = [make_video_feature(offset=2, grid_thw=(2, 2, 2))]
positions, _ = model.get_mrope_input_positions(input_tokens, mm_features)
assert positions.shape == (3, len(input_tokens))
def test_mrope_positions_audio_in_video_at_prompt_end():
model = make_model()
# video grid (1, 2, 2) -> 4 video tokens; audio_feature_length=16 -> 2
# audio tokens, interleaved after the video tokens.
# [T, T, <|vision_start|>, <|audio_start|>, V x 4, A x 2,
# <|audio_end|>, <|vision_end|>]
input_tokens = (
[TEXT_TOKEN_ID] * 2
+ [VISION_START_TOKEN_ID, AUDIO_START_TOKEN_ID]
+ [VIDEO_TOKEN_ID] * 4
+ [AUDIO_TOKEN_ID] * 2
+ [AUDIO_END_TOKEN_ID, VISION_END_TOKEN_ID]
)
mm_features = [
make_video_feature(offset=3, grid_thw=(1, 2, 2), use_audio_in_video=True),
make_audio_feature(offset=3, audio_feature_length=16),
]
positions, _ = model.get_mrope_input_positions(input_tokens, mm_features)
assert positions.shape == (3, len(input_tokens))
def test_mrope_positions_image_then_text():
"""Media followed by text: positions must match the expected layout exactly.
Before the fix this case passed the length guard but every position from
the first <|image_pad|> onward was silently shifted by one slot.
"""
model = make_model()
# [T, T, <|vision_start|>, <|image_pad|> x 4, <|vision_end|>, T, T, T]
input_tokens = (
[TEXT_TOKEN_ID] * 2
+ [VISION_START_TOKEN_ID]
+ [IMAGE_TOKEN_ID] * 4
+ [VISION_END_TOKEN_ID]
+ [TEXT_TOKEN_ID] * 3
)
mm_features = [make_image_feature(offset=3, grid_thw=(1, 2, 2))]
positions, delta = model.get_mrope_input_positions(input_tokens, mm_features)
expected = torch.tensor(
[
[0, 1, 2, 3, 3, 3, 3, 5, 6, 7, 8],
[0, 1, 2, 3, 3, 4, 4, 5, 6, 7, 8],
[0, 1, 2, 3, 4, 3, 4, 5, 6, 7, 8],
]
)
torch.testing.assert_close(positions, expected)
assert delta == 9 - len(input_tokens)
@@ -2260,10 +2260,6 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
)
st_idx += text_len
bos_pos = np.broadcast_to(np.array([st_idx]), (3, 1))
llm_pos_ids_list.append(bos_pos)
st_idx += 1
if modality == "audio":
audio_tokens = self._compute_audio_token_count(
data["audio_feature_length"]
@@ -2276,7 +2272,7 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
eos_pos = np.broadcast_to(np.array([st_idx]), (3, 1))
llm_pos_ids_list.append(eos_pos)
st = offset + 1 + audio_tokens + 1
st = offset + audio_tokens + 1
elif modality == "image":
grid_t = data["grid_t"]
@@ -2294,7 +2290,7 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
eos_pos = np.broadcast_to(np.array([st_idx]), (3, 1))
llm_pos_ids_list.append(eos_pos)
st = offset + 1 + image_len + 1
st = offset + image_len + 1
elif modality == "video":
grid_t = data["grid_t"]
@@ -2313,7 +2309,7 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
eos_pos = np.broadcast_to(np.array([st_idx]), (3, 1))
llm_pos_ids_list.append(eos_pos)
st = offset + 1 + video_len + 1
st = offset + video_len + 1
else:
audio_bos_pos = np.broadcast_to(np.array([st_idx - 1]), (3, 1))
llm_pos_ids_list.append(audio_bos_pos)
@@ -2330,7 +2326,7 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
audio_len = self._compute_audio_token_count(
data["audio_feature_length"]
)
st = offset + 2 + video_len + audio_len + 2
st = offset + 1 + video_len + audio_len + 2
if st < seq_len:
st_idx = int(llm_pos_ids_list[-1].max()) + 1 if llm_pos_ids_list else 0