Optimize MiniCPM-o 4.5 Whisper chunk-attention mask construction (#5382)

Signed-off-by: a-kaa <siruhe666@gmail.com>
Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>

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)

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