From 0a033592245706f6f7cb6b763c6262ca6400e8d5 Mon Sep 17 00:00:00 2001 From: NATURE Date: Sat, 1 Aug 2026 11:01:06 +0800 Subject: [PATCH] [Bugfix][MiniCPM-o] Stabilize single-GPU audio generation (#5637) Signed-off-by: natureofnature --- tests/config/test_config_factory.py | 26 +++++++++++++++++++ tests/helpers/minicpmo_4_5_duplex.py | 13 ++++++++++ vllm_omni/deploy/minicpmo_4_5.yaml | 7 +++++ .../minicpmo_4_5_3gpu_stage1_replicas.yaml | 7 +++++ .../minicpmo_4_5_4gpu_stage1_replicas.yaml | 7 +++++ .../models/minicpmo_4_5/batched_token2wav.py | 21 +++++++++++++++ 6 files changed, 81 insertions(+) diff --git a/tests/config/test_config_factory.py b/tests/config/test_config_factory.py index a811868c1..33dcd0551 100644 --- a/tests/config/test_config_factory.py +++ b/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(): diff --git a/tests/helpers/minicpmo_4_5_duplex.py b/tests/helpers/minicpmo_4_5_duplex.py index dd32c0dcf..e19903e84 100644 --- a/tests/helpers/minicpmo_4_5_duplex.py +++ b/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( diff --git a/vllm_omni/deploy/minicpmo_4_5.yaml b/vllm_omni/deploy/minicpmo_4_5.yaml index abe8b5ad2..b3e211175 100644 --- a/vllm_omni/deploy/minicpmo_4_5.yaml +++ b/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 diff --git a/vllm_omni/deploy/minicpmo_4_5_3gpu_stage1_replicas.yaml b/vllm_omni/deploy/minicpmo_4_5_3gpu_stage1_replicas.yaml index 17cd014e9..7a0513167 100644 --- a/vllm_omni/deploy/minicpmo_4_5_3gpu_stage1_replicas.yaml +++ b/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 diff --git a/vllm_omni/deploy/minicpmo_4_5_4gpu_stage1_replicas.yaml b/vllm_omni/deploy/minicpmo_4_5_4gpu_stage1_replicas.yaml index e4fe45922..da28456fe 100644 --- a/vllm_omni/deploy/minicpmo_4_5_4gpu_stage1_replicas.yaml +++ b/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 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 f0e07e50b..f2d96be0d 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,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)