[Perf] release unused Code2Wav stream caches (#6100)

Signed-off-by: natureofnature <wzliu@connect.hku.hk>

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()

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.