From 70995834040f488d63ed5c916fcee4e94268a902 Mon Sep 17 00:00:00 2001 From: BruceLoveDecimal <39156883+BruceLoveDecimal@users.noreply.github.com> Date: Wed, 19 Aug 2026 15:26:03 +0800 Subject: [PATCH] [Bugfix][MiniCPM-o] Restore streaming audio cache and unstall the first duplex response (#6274) Signed-off-by: BruceLoveDecimal --- .../test_streaming_audio_cache.py | 156 ++++++++++++++++++ .../models/minicpmo_4_5/batched_token2wav.py | 20 +++ .../minicpmo_4_5/minicpmo_4_5_omni_llm.py | 40 +++-- 3 files changed, 205 insertions(+), 11 deletions(-) create mode 100644 tests/model_executor/models/minicpmo_4_5/test_streaming_audio_cache.py diff --git a/tests/model_executor/models/minicpmo_4_5/test_streaming_audio_cache.py b/tests/model_executor/models/minicpmo_4_5/test_streaming_audio_cache.py new file mode 100644 index 000000000..e867e0cab --- /dev/null +++ b/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) diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.py b/vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.py index 41223e864..fe5c8a57e 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.py +++ b/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. 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 a182848cf..490d8e0cf 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 @@ -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),