diff --git a/tests/model_executor/models/minicpmo_4_5/test_code2wav_batching.py b/tests/model_executor/models/minicpmo_4_5/test_code2wav_batching.py index 7b5bb1044..38d9d9093 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_code2wav_batching.py +++ b/tests/model_executor/models/minicpmo_4_5/test_code2wav_batching.py @@ -50,6 +50,8 @@ class _FakeEstimator(nn.Module): self.blocks = [_FakeBlock()] self.cfg_batches: list[int] = [] self.speaker_order: list[list[float]] = [] + self.register_buffer("att_cache_buffer", torch.ones(1), persistent=False) + self.register_buffer("cnn_cache_buffer", torch.ones(1), persistent=False) def t_embedder(self, time): return time[:, None] @@ -79,6 +81,8 @@ class _FakeDecoder(nn.Module): self.estimator = _FakeEstimator() self.inference_cfg_rate = 0.7 self.register_buffer("rand_noise", torch.zeros(1, 1, 100), persistent=False) + self.register_buffer("att_cache_buffer", torch.ones(1), persistent=False) + self.register_buffer("cnn_cache_buffer", torch.ones(1), persistent=False) class _FakeFlow(nn.Module): @@ -158,6 +162,24 @@ def _model(): return model, token2wav +def test_adapter_releases_unused_upstream_streaming_buffers(): + token2wav = _FakeToken2Wav() + decoder = token2wav.flow.decoder + modules = (decoder, decoder.estimator) + + assert all( + getattr(module, name) is not None for module in modules for name in ("att_cache_buffer", "cnn_cache_buffer") + ) + + BatchedToken2Wav(token2wav) + + assert all( + name in module._buffers and getattr(module, name) is None + for module in modules + for name in ("att_cache_buffer", "cnn_cache_buffer") + ) + + def test_code2wav_resolves_hf_model_id_for_assets(mocker, tmp_path): resolved_root = tmp_path / "snapshot" resolved_root.mkdir() 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 f2d96be0d..15e0c8f20 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 @@ -65,6 +65,14 @@ class BatchedToken2Wav(nn.Module): self._token2wav = token2wav self.flow = token2wav.flow self.hift = token2wav.hift + # The upstream streaming path preallocates fixed-size CFM and DiT + # caches. This adapter never calls that path and supplies dynamically + # sized request-owned buffers to ``blocks_forward_chunk`` instead. + decoder = self.flow.decoder + for module in (decoder, decoder.estimator): + for buffer_name in ("att_cache_buffer", "cnn_cache_buffer"): + if buffer_name in module._buffers: + setattr(module, buffer_name, None) 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.