[Bugfix][MiniCPM-o] Stabilize single-GPU audio generation (#5637)

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

tests/config/test_config_factory.py

@@ -2214,6 +2214,32 @@ class TestPlatformOverrides:
         assert rocm.stages[0].enforce_eager is None
         assert rocm.stages[1].enforce_eager is True
 
+    def test_minicpmo_4_5_cuda_caps_talker_kv_cache(self):
+        pipeline = resolve_pipeline_config("minicpmo_4_5")
+        assert isinstance(pipeline, PipelineConfig)
+        deploy_path = Path(get_deploy_config_path("minicpmo_4_5.yaml"))
+
+        cuda = _apply_platform_overrides(load_deploy_config(deploy_path), platform="cuda")
+        cuda_stages = merge_pipeline_deploy(pipeline, cuda)
+        assert cuda_stages[1].yaml_engine_args["kv_cache_memory_bytes"] == 2 * 1024**3
+
+        # The CUDA memory budget must not leak into the existing NPU profile.
+        npu = _apply_platform_overrides(load_deploy_config(deploy_path), platform="npu")
+        npu_stages = merge_pipeline_deploy(pipeline, npu)
+        assert "kv_cache_memory_bytes" not in npu_stages[1].yaml_engine_args
+
+        # Multi-GPU replica profiles place Talker on dedicated devices and
+        # preserve their previous automatic KV-cache sizing.
+        for filename in (
+            "minicpmo_4_5_3gpu_stage1_replicas.yaml",
+            "minicpmo_4_5_4gpu_stage1_replicas.yaml",
+        ):
+            replica = _apply_platform_overrides(
+                load_deploy_config(Path(get_deploy_config_path(filename))), platform="cuda"
+            )
+            replica_stages = merge_pipeline_deploy(pipeline, replica)
+            assert replica_stages[1].yaml_engine_args["kv_cache_memory_bytes"] is None
+
     def test_npu_overrides(self):
         deploy_path = Path(get_deploy_config_path("qwen3_omni_moe.yaml"))
         if not deploy_path.exists():

tests/helpers/minicpmo_4_5_duplex.py

@@ -24,6 +24,19 @@ DEPLOY_CONFIG = modify_stage_config(
             1: {"kv_cache_memory_bytes": 512 * 1024 * 1024},
             2: {"kv_cache_memory_bytes": 256 * 1024 * 1024},
         },
+        # Platform overrides are applied after ordinary stage settings. Keep
+        # the constrained duplex test's Talker budget ahead of the base CUDA
+        # profile's single-GPU default.
+        "platforms": {
+            "cuda": {
+                "stages": [
+                    {
+                        "stage_id": 1,
+                        "kv_cache_memory_bytes": 512 * 1024 * 1024,
+                    }
+                ]
+            }
+        },
     },
 )
 CORE_DEPLOY_CONFIG = modify_stage_config(

vllm_omni/deploy/minicpmo_4_5.yaml

@@ -75,6 +75,13 @@ stages:
       detokenize: true
 
 platforms:
+  cuda:
+    stages:
+      - stage_id: 1
+        # Four Talker streams need much less KV than automatic sizing.
+        # Leave single-GPU headroom for concurrent Stage 2 diffusion batches.
+        kv_cache_memory_bytes: 2147483648  # 2 GiB
+
   npu:
     stages:
       - stage_id: 0

vllm_omni/deploy/minicpmo_4_5_3gpu_stage1_replicas.yaml

@@ -15,3 +15,10 @@ stages:
     devices: "1,2"
     num_replicas: 2
     async_scheduling: false
+
+platforms:
+  cuda:
+    stages:
+      - stage_id: 1
+        # Stage 1 has dedicated GPUs in this profile; retain automatic sizing.
+        kv_cache_memory_bytes: null

vllm_omni/deploy/minicpmo_4_5_4gpu_stage1_replicas.yaml

@@ -15,3 +15,10 @@ stages:
     devices: "1,2,3"
     num_replicas: 3
     async_scheduling: false
+
+platforms:
+  cuda:
+    stages:
+      - stage_id: 1
+        # Stage 1 has dedicated GPUs in this profile; retain automatic sizing.
+        kv_cache_memory_bytes: null

vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.py

@@ -65,6 +65,27 @@ class BatchedToken2Wav(nn.Module):
         self._token2wav = token2wav
         self.flow = token2wav.flow
         self.hift = token2wav.hift
+        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.
+            # Otherwise, the first live audio chunk can fail when async stages
+            # share one GPU.
+            device = hift_parameter.device
+            dtype = hift_parameter.dtype
+            mel_channels = int(self.hift.conv_pre.in_channels)
+            with (
+                torch.inference_mode(),
+                torch.random.fork_rng(devices=[device]),
+                _autocast_disabled(device),
+            ):
+                # 50 mel frames match the default first streamed vocoder chunk.
+                speech, source = self.hift(
+                    torch.zeros((1, mel_channels, 50), device=device, dtype=dtype),
+                    torch.zeros((1, 1, 0), device=device, dtype=dtype),
+                )
+            torch.accelerator.synchronize(device)
+            del speech, source
+            torch.accelerator.empty_cache()
         self.float16 = bool(token2wav.float16)
         self.n_timesteps = int(token2wav.n_timesteps)
         self.mel_cache_len = int(token2wav.mel_cache_len)