From 44198f577fe5cfb4b02297bf6ff7235eb31d742b Mon Sep 17 00:00:00 2001 From: C <3121985758@qq.com> Date: Sat, 3 Oct 2026 09:00:51 +0800 Subject: [PATCH] [Bugfix][Model] Fix M-RoPE offset double-count in Qwen3-Omni (#58890) Signed-off-by: C1ves <3121985758@qq.com> --- tests/model_executor/test_qwen3_omni_mrope.py | 198 ++++++++++++++++++ .../models/qwen3_omni_moe_thinker.py | 12 +- 2 files changed, 202 insertions(+), 8 deletions(-) create mode 100644 tests/model_executor/test_qwen3_omni_mrope.py diff --git a/tests/model_executor/test_qwen3_omni_mrope.py b/tests/model_executor/test_qwen3_omni_mrope.py new file mode 100644 index 0000000000..bb8de63dc5 --- /dev/null +++ b/tests/model_executor/test_qwen3_omni_mrope.py @@ -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) diff --git a/vllm/model_executor/models/qwen3_omni_moe_thinker.py b/vllm/model_executor/models/qwen3_omni_moe_thinker.py index f9338b1189..d3510f3e64 100755 --- a/vllm/model_executor/models/qwen3_omni_moe_thinker.py +++ b/vllm/model_executor/models/qwen3_omni_moe_thinker.py @@ -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