diff --git a/tests/model_executor/models/minicpmo_4_5/test_audio_chunk_mask.py b/tests/model_executor/models/minicpmo_4_5/test_audio_chunk_mask.py new file mode 100644 index 000000000..7769b38be --- /dev/null +++ b/tests/model_executor/models/minicpmo_4_5/test_audio_chunk_mask.py @@ -0,0 +1,67 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Regression tests for MiniCPM-o 4.5 audio chunk attention masks.""" + +from __future__ import annotations + +import pytest +import torch + +from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_llm import ( + MiniCPMO45OmniLLMForConditionalGeneration, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _reference_chunk_mask( + size: int, + chunk_size: int, + num_left_chunks: int, + num_lookhead: int, +) -> torch.Tensor: + """Preserve the original row-by-row implementation as a test oracle.""" + mask = torch.zeros(size, size, dtype=torch.bool) + for row in range(size): + if num_left_chunks < 0: + start = 0 + else: + start = max( + (row // chunk_size - num_left_chunks) * chunk_size, + 0, + ) + end = min( + (row // chunk_size + 1) * chunk_size + num_lookhead, + size, + ) + mask[row, start:end] = True + return mask + + +@pytest.mark.parametrize("size", [0, 1, 49, 50, 51, 1500]) +@pytest.mark.parametrize("chunk_size", [1, 17, 50]) +@pytest.mark.parametrize("num_left_chunks", [-1, 0, 1, 3]) +@pytest.mark.parametrize("num_lookhead", [0, 1, 7]) +def test_subsequent_chunk_mask_matches_row_by_row_reference( + size: int, + chunk_size: int, + num_left_chunks: int, + num_lookhead: int, +) -> None: + actual = MiniCPMO45OmniLLMForConditionalGeneration.subsequent_chunk_mask( + None, + size=size, + chunk_size=chunk_size, + num_left_chunks=num_left_chunks, + num_lookhead=num_lookhead, + ) + expected = _reference_chunk_mask( + size, + chunk_size, + num_left_chunks, + num_lookhead, + ) + + assert actual.dtype == torch.bool + assert actual.shape == (size, size) + assert torch.equal(actual, expected) diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py index 4671e3b96..ce00f21ed 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py @@ -4255,15 +4255,21 @@ class MiniCPMO45OmniLLMForConditionalGeneration(nn.Module, SupportsMultiModal, S torch.Tensor: mask """ - ret = torch.zeros(size, size, device=device, dtype=torch.bool) - for i in range(size): - if num_left_chunks < 0: - start = 0 - else: - start = max((i // chunk_size - num_left_chunks) * chunk_size, 0) - ending = min((i // chunk_size + 1) * chunk_size + num_lookhead, size) - ret[i, start:ending] = True - return ret + indices = torch.arange(size, device=device) + chunk_index = indices // chunk_size + chunk_end = torch.clamp( + (chunk_index + 1) * chunk_size + num_lookhead, + max=size, + ) + + mask = indices.unsqueeze(0) < chunk_end.unsqueeze(1) + if num_left_chunks >= 0: + chunk_start = torch.clamp( + (chunk_index - num_left_chunks) * chunk_size, + min=0, + ) + mask.logical_and_(indices.unsqueeze(0) >= chunk_start.unsqueeze(1)) + return mask def _get_feat_extract_output_lengths(self, input_lengths: torch.LongTensor): input_lengths_after_cnn = (input_lengths - 1) // 2 + 1