[Bugfix][MiniCPM-o] Restore streaming audio cache and unstall the first duplex response (#6274)

Signed-off-by: BruceLoveDecimal <liuqihao970610@gmail.com>

tests/model_executor/models/minicpmo_4_5/test_streaming_audio_cache.py

@@ -0,0 +1,156 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Regression tests for MiniCPM-o 4.5 streaming audio cache compatibility."""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+from torch import nn
+from transformers.models.whisper.modeling_whisper import WhisperConfig
+
+from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_llm import (
+    WHISPER_ATTENTION_CLASSES,
+    MiniCPMWhisperEncoder,
+    MiniCPMWhisperEncoderLayer,
+    _get_audio_cache_length,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _LegacyWhisperAttention(nn.Module):
+    def __init__(self, **kwargs) -> None:
+        super().__init__()
+        self.seen_cache = None
+
+    def forward(self, hidden_states, past_key_value=None, **kwargs):
+        self.seen_cache = past_key_value
+        return hidden_states, None, past_key_value
+
+
+class _CurrentWhisperAttention(nn.Module):
+    def __init__(self, **kwargs) -> None:
+        super().__init__()
+        self.seen_cache = None
+
+    def forward(self, hidden_states, past_key_values=None, **kwargs):
+        self.seen_cache = past_key_values
+        return hidden_states, None
+
+
+@pytest.mark.parametrize(
+    ("implementation", "attention_cls"),
+    [
+        ("test_legacy_cache", _LegacyWhisperAttention),
+        ("test_current_cache", _CurrentWhisperAttention),
+    ],
+)
+def test_whisper_attention_preserves_streaming_cache(
+    monkeypatch: pytest.MonkeyPatch,
+    implementation: str,
+    attention_cls: type[nn.Module],
+) -> None:
+    monkeypatch.setitem(WHISPER_ATTENTION_CLASSES, implementation, attention_cls)
+    config = WhisperConfig(
+        d_model=4,
+        encoder_attention_heads=1,
+        encoder_ffn_dim=8,
+        encoder_layers=1,
+    )
+    config._attn_implementation = implementation
+    layer = MiniCPMWhisperEncoderLayer(config, layer_idx=0)
+    cache = object()
+
+    outputs = layer(
+        torch.randn(1, 2, config.d_model),
+        attention_mask=None,
+        layer_head_mask=None,
+        past_key_values=cache,
+        use_cache=True,
+    )
+
+    assert layer.self_attn.seen_cache is cache
+    assert outputs[-1] is cache
+
+
+def test_audio_cache_length_uses_current_cache_api() -> None:
+    dynamic_cache = SimpleNamespace(get_seq_length=lambda: 7)
+    encoder_decoder_cache = SimpleNamespace(self_attention_cache=dynamic_cache)
+
+    assert _get_audio_cache_length(encoder_decoder_cache) == 7
+
+
+def test_audio_cache_length_supports_legacy_cache() -> None:
+    key = torch.zeros(1, 1, 5, 2)
+    value = torch.zeros_like(key)
+
+    assert _get_audio_cache_length(((key, value),)) == 5
+
+
+def _tiny_encoder_config(attn_implementation: str) -> WhisperConfig:
+    config = WhisperConfig(
+        num_mel_bins=16,
+        d_model=32,
+        encoder_layers=2,
+        encoder_attention_heads=2,
+        encoder_ffn_dim=64,
+        max_source_positions=100,
+        dropout=0.0,
+        attention_dropout=0.0,
+    )
+    config._attn_implementation = attn_implementation
+    return config
+
+
+@pytest.mark.parametrize("attn_implementation", ["eager", "sdpa"])
+def test_encoder_reuses_cache_across_streaming_chunks(attn_implementation: str) -> None:
+    """Encode two chunks against the installed Whisper attention and cache classes.
+
+    Unlike the stub-based tests above, this drives the real transformers cache
+    API that ``get_audio_embedding_streaming`` relies on, so it fails whenever a
+    removed cache method or renamed attention kwarg breaks streaming reuse.
+    """
+    torch.manual_seed(0)
+    config = _tiny_encoder_config(attn_implementation)
+    encoder = MiniCPMWhisperEncoder(config).eval()
+
+    first_frames, second_frames = 20, 14
+    first_len = (first_frames - 1) // 2 + 1
+    second_len = (second_frames - 1) // 2 + 1
+    first_mel = torch.randn(1, config.num_mel_bins, first_frames)
+    second_mel = torch.randn(1, config.num_mel_bins, second_frames)
+
+    with torch.no_grad():
+        first = encoder(
+            first_mel,
+            attention_mask=torch.zeros(1, 1, first_len, first_len),
+            past_key_values=None,
+            use_cache=True,
+            output_hidden_states=True,
+        )
+        cache = first.past_key_values
+        assert cache is not None
+        assert _get_audio_cache_length(cache) == first_len
+
+        second = encoder(
+            second_mel,
+            attention_mask=torch.zeros(1, 1, second_len, first_len + second_len),
+            past_key_values=cache,
+            use_cache=True,
+            output_hidden_states=True,
+        )
+        assert _get_audio_cache_length(second.past_key_values) == first_len + second_len
+
+        stateless = encoder(
+            second_mel,
+            attention_mask=torch.zeros(1, 1, second_len, second_len),
+            past_key_values=None,
+            use_cache=True,
+            output_hidden_states=True,
+        )
+
+    assert torch.isfinite(second.last_hidden_state).all()
+    assert not torch.allclose(second.last_hidden_state, stateless.last_hidden_state)

vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.py

@@ -58,6 +58,25 @@ class BatchedToken2WavState:
     hift_cache: dict[str, torch.Tensor]
 
 
+def _undecorate_dynamo(module: nn.Module, method: str) -> None:
+    """Restore ``method`` on ``module`` if TorchDynamo wrapped it.
+
+    ``cosyvoice2`` decorates ``UpsampleConformerEncoderV2.forward_chunk`` with
+    ``torch.compile(backend="eager")``. That backend performs no Inductor
+    optimisation, so the wrapper only adds tracing and guard construction, and
+    duplex pays it again on every unseen chunk shape -- seconds inside a live
+    response. Dropping the wrapper leaves the original implementation, which is
+    what the eager backend was executing anyway.
+    """
+    bound = getattr(module, method, None)
+    original = getattr(bound, "_torchdynamo_orig_callable", None) or getattr(bound, "__wrapped__", None)
+    if original is None:
+        return
+    function = getattr(original, "__func__", original)
+    module.__dict__[method] = function.__get__(module, type(module))
+    logger.info("Bypassed TorchDynamo wrapper on %s.%s", type(module).__name__, method)
+
+
 class BatchedToken2Wav(nn.Module):
     """Drive Token2wav's modules with dynamically-sized, request-owned caches.
 
@@ -82,6 +101,7 @@ class BatchedToken2Wav(nn.Module):
         self._trt_stepper = trt_stepper
         self.flow = token2wav.flow
         self.hift = token2wav.hift
+        _undecorate_dynamo(self.flow.encoder, "forward_chunk")
         hift_parameter = next(self.hift.parameters(), None)
         if hift_parameter is not None and hift_parameter.device.type == "cuda":
             # Prime the CUDA state used by HiFT during backend construction.

vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py

@@ -17,6 +17,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import inspect
 import math
 import os
 import warnings
@@ -2505,6 +2506,17 @@ def _in_projection(
     return linear(q, w_q, b_q), linear(k, w_k, b_k), linear(v, w_v, b_v)
 
 
+def _get_audio_cache_length(past_key_values: Any) -> int:
+    cache = getattr(past_key_values, "self_attention_cache", past_key_values)
+    get_seq_length = getattr(cache, "get_seq_length", None)
+    if callable(get_seq_length):
+        try:
+            return int(get_seq_length())
+        except TypeError:
+            return int(get_seq_length(0))
+    return int(cache[0][0].shape[2])
+
+
 # Copied from transformers.models.whisper.modeling_whisper.WhisperEncoderLayer and add use_cache for streaming inference
 class MiniCPMWhisperEncoderLayer(nn.Module):
     def __init__(self, config: WhisperConfig, layer_idx: int = None):
@@ -2517,6 +2529,10 @@ class MiniCPMWhisperEncoderLayer(nn.Module):
             config=config,
             layer_idx=layer_idx,
         )
+        attention_parameters = inspect.signature(self.self_attn.forward).parameters
+        self._past_key_values_kwarg = (
+            "past_key_values" if "past_key_values" in attention_parameters else "past_key_value"
+        )
         self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim)
         self.dropout = config.dropout
         self.activation_fn = ACT2FN[config.activation_function]
@@ -2554,16 +2570,18 @@ class MiniCPMWhisperEncoderLayer(nn.Module):
         """
         residual = hidden_states
         hidden_states = self.self_attn_layer_norm(hidden_states)
-        attn_out = self.self_attn(
-            hidden_states=hidden_states,
-            attention_mask=attention_mask,
-            layer_head_mask=layer_head_mask,
-            output_attentions=output_attentions,
-            past_key_value=past_key_values,
-        )
+        attention_kwargs = {
+            "hidden_states": hidden_states,
+            "attention_mask": attention_mask,
+            "layer_head_mask": layer_head_mask,
+            "output_attentions": output_attentions,
+            self._past_key_values_kwarg: past_key_values,
+        }
+        attn_out = self.self_attn(**attention_kwargs)
         hidden_states = attn_out[0]
         attn_weights = attn_out[1] if len(attn_out) > 1 else None
-        past_key_values = attn_out[2] if len(attn_out) > 2 else None
+        if len(attn_out) > 2:
+            past_key_values = attn_out[2]
         hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
         hidden_states = residual + hidden_states
 
@@ -2769,7 +2787,7 @@ class MiniCPMWhisperEncoder(WhisperEncoder):
                 past_key_values = EncoderDecoderCache(past_key_values, DynamicCache())
             else:
                 pass
-            past_key_values_length = past_key_values.self_attention_cache.get_usable_length(inputs_embeds.shape[1])
+            past_key_values_length = _get_audio_cache_length(past_key_values)
             if inputs_embeds.shape[1] + past_key_values_length > embed_pos.shape[0]:
                 logger.warning("seems the audio is longer than 30s. repeating the last part of the audio")
                 embed_pos_front = embed_pos[past_key_values_length:, :]
@@ -4468,7 +4486,7 @@ class MiniCPMO45OmniLLMForConditionalGeneration(nn.Module, SupportsMultiModal, S
             return []
 
         if self.audio_past_key_values is not None:
-            cache_length = self.audio_past_key_values[0][0].shape[2]
+            cache_length = _get_audio_cache_length(self.audio_past_key_values)
             apm_max_len = self.apm.embed_positions.weight.shape[0]
             if cache_length + current_seq_len >= apm_max_len:
                 logger.warning(
@@ -4480,7 +4498,7 @@ class MiniCPMO45OmniLLMForConditionalGeneration(nn.Module, SupportsMultiModal, S
 
         past_len = 0
         if self.audio_past_key_values is not None:
-            past_len = self.audio_past_key_values[0][0].shape[2]
+            past_len = _get_audio_cache_length(self.audio_past_key_values)
         total_seq_len = past_len + current_seq_len
         audio_attention_mask = torch.zeros(
             (batch_size, 1, current_seq_len, total_seq_len),