[Bugfix][Model] Fix Qwen3-Omni interleaved M-RoPE boundary positions (#59842)

Signed-off-by: Tianyu Guo <guoty@inferact.ai>
Co-authored-by: OpenAI Codex <codex@openai.com>
This commit is contained in:
Tianyu Guo
2026-10-06 03:35:53 +00:00
committed by GitHub
co-authored by OpenAI Codex
parent 48eb44b719
commit 8c417a2dae
2 changed files with 35 additions and 16 deletions
+32 -14
View File
@@ -145,27 +145,45 @@ def test_mrope_positions_video_at_prompt_end():
assert positions.shape == (3, len(input_tokens))
def test_mrope_positions_audio_in_video_at_prompt_end():
@pytest.mark.parametrize("num_videos", [1, 2])
@pytest.mark.parametrize("trailing_tokens", [0, 3])
def test_mrope_positions_audio_in_video_boundaries(num_videos, trailing_tokens):
"""Interleaved BOS/EOS positions and decode delta match Qwen3-Omni HF."""
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_tokens = (
[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),
]
input_tokens = [TEXT_TOKEN_ID] * 2
mm_features = []
expected = [torch.tensor([[0, 1]]).expand(3, -1)]
video_positions = torch.tensor(
[
[2, 3, 4, 4, 4, 4, 4, 5, 6, 7],
[2, 3, 4, 4, 5, 5, 4, 5, 6, 7],
[2, 3, 4, 5, 4, 5, 4, 5, 6, 7],
]
)
for index in range(num_videos):
offset = len(input_tokens) + 1
mm_features.extend(
[
make_video_feature(offset, (1, 2, 2), use_audio_in_video=True),
make_audio_feature(offset, audio_feature_length=16),
]
)
input_tokens.extend(video_tokens)
expected.append(video_positions + index * 6)
positions, _ = model.get_mrope_input_positions(input_tokens, mm_features)
input_tokens.extend([TEXT_TOKEN_ID] * trailing_tokens)
next_position = 2 + num_videos * 6
expected.append((torch.arange(trailing_tokens) + next_position).expand(3, -1))
positions, delta = model.get_mrope_input_positions(input_tokens, mm_features)
assert positions.shape == (3, len(input_tokens))
torch.testing.assert_close(positions, torch.cat(expected, dim=1))
assert delta == -4 * num_videos
def test_mrope_positions_image_then_text():
@@ -2311,8 +2311,9 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
llm_pos_ids_list.append(eos_pos)
st = offset + video_len + 1
else:
audio_bos_pos = np.broadcast_to(np.array([st_idx - 1]), (3, 1))
audio_bos_pos = np.broadcast_to(np.array([st_idx]), (3, 1))
llm_pos_ids_list.append(audio_bos_pos)
st_idx += 1
pos_ids, _ = self._compute_interleaved_positions(st_idx, data)
llm_pos_ids_list.append(pos_ids)
@@ -2320,7 +2321,7 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
eos_pos = np.broadcast_to(np.array([st_idx]), (3, 1))
llm_pos_ids_list.append(eos_pos)
llm_pos_ids_list.append(eos_pos)
llm_pos_ids_list.append(eos_pos + 1)
video_len = grid_t * grid_h * grid_w
audio_len = self._compute_audio_token_count(