[Perf][VoxCPM2] Optimize VoxCPM2 high-concurrency decode throughput. (#3882)

Signed-off-by: Sy03 <1370724210@qq.com>
Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: Gao Han <hgaoaf@connect.ust.hk>

tests/model_executor/models/voxcpm2/test_talker_state_eviction.py

@@ -5,6 +5,7 @@
 from __future__ import annotations
 
 import functools
+from types import SimpleNamespace
 
 import pytest
 
@@ -17,13 +18,14 @@ def _voxcpm2_talker_mod():
     from vllm_omni.model_executor.models.voxcpm2.voxcpm2_talker import (
         VoxCPM2TalkerForConditionalGeneration,
         _RequestState,
+        _VoxCPM2RuntimeConfig,
     )
 
-    return VoxCPM2TalkerForConditionalGeneration, _RequestState
+    return VoxCPM2TalkerForConditionalGeneration, _RequestState, _VoxCPM2RuntimeConfig
 
 
 def _make_bare_talker():
-    VoxCPM2TalkerForConditionalGeneration, _ = _voxcpm2_talker_mod()
+    VoxCPM2TalkerForConditionalGeneration, _, _ = _voxcpm2_talker_mod()
     talker = VoxCPM2TalkerForConditionalGeneration.__new__(VoxCPM2TalkerForConditionalGeneration)
     talker._active_states = {}
     talker._current_request_id = None
@@ -34,11 +36,27 @@ def _make_bare_talker():
     talker._max_batch_size = 4
     talker._active_state_warn_threshold = 512
     talker._active_state_warned = False
+    talker._enable_profiling = False
+    talker._audio_emit_every = 1
+    talker._vae_decode_every = 1
+    talker._enable_delayed_audio_copy = False
+    talker._delayed_audio_copy_use_events = False
+    talker._coalesce_audio_d2h = False
+    talker._enable_batched_vae_decode = False
+    talker._audio_copy_stream = None
+    talker._enable_vae_cuda_graph = False
+    talker._enable_cfm_cuda_graph = False
+    talker._enable_cfm_prealloc_output = False
+    talker._enable_batched_cfm = True
+    talker._deterministic_cfm_noise = False
+    talker._cfm_buffers = None
+    talker._last_audio_output_req_ids = []
+    talker._batched_fsq_fusion_max_batch = 32
     return talker
 
 
 def _seed_cached_decode(talker, req_id: str):
-    _, _RequestState = _voxcpm2_talker_mod()
+    _, _RequestState, _ = _voxcpm2_talker_mod()
     state = _RequestState(request_id=req_id)
     state.prefill_completed = True
     state.decode_step_count = 5
@@ -47,6 +65,23 @@ def _seed_cached_decode(talker, req_id: str):
 
 
 class TestStateEvictionContract:
+    def test_runtime_config_normalizes_mutually_exclusive_paths(self) -> None:
+        _, _, RuntimeConfig = _voxcpm2_talker_mod()
+
+        cfg = RuntimeConfig(
+            enable_batched_cfm=True,
+            enable_cfm_cuda_graph=True,
+            enable_cfm_prealloc_output=True,
+            enable_batched_vae_decode=True,
+            enable_delayed_audio_copy=True,
+        )._normalized()
+
+        assert cfg.enable_batched_cfm is True
+        assert cfg.enable_cfm_cuda_graph is False
+        assert cfg.enable_cfm_prealloc_output is False
+        assert cfg.enable_batched_vae_decode is False
+        assert cfg.enable_delayed_audio_copy is True
+
     def test_pending_requests_is_not_used_for_eviction(self) -> None:
         talker = _make_bare_talker()
 
@@ -61,6 +96,61 @@ class TestStateEvictionContract:
             assert rid in talker._active_states
             assert talker._active_states[rid].prefill_completed is True
 
+    def test_compute_logits_does_not_sync_stop_bool_on_default_path(self) -> None:
+        _, RState, _ = _voxcpm2_talker_mod()
+        talker = _make_bare_talker()
+        talker.config = SimpleNamespace(vocab_size=4)
+        state = RState(request_id="req", precomputed_stop_logits=torch.tensor([[0.0, 1.0]]))
+        talker._active_states["req"] = state
+        talker._results_queue = [("req", state.precomputed_stop_logits)]
+
+        logits = talker.compute_logits(torch.zeros(1, 1))
+
+        assert logits[0, 0] == 0.0
+        assert logits[0, 1] == 1.0
+        assert state.is_stopping is False
+        assert state.precomputed_stop_logits is None
+        assert state.precomputed_is_stopping is None
+
+    def test_compute_logits_reuses_cached_stop_bool_from_sparse_audio_path(self) -> None:
+        _, RState, _ = _voxcpm2_talker_mod()
+        talker = _make_bare_talker()
+        talker.config = SimpleNamespace(vocab_size=4)
+        state = RState(
+            request_id="req",
+            precomputed_stop_logits=torch.tensor([[0.0, 1.0]]),
+            precomputed_is_stopping=True,
+        )
+        talker._active_states["req"] = state
+        talker._results_queue = [("req", state.precomputed_stop_logits)]
+
+        logits = talker.compute_logits(torch.zeros(1, 1))
+
+        assert logits[0, 0] == 0.0
+        assert logits[0, 1] == 1.0
+        assert state.is_stopping is True
+        assert state.precomputed_stop_logits is None
+        assert state.precomputed_is_stopping is None
+
+    def test_batched_stop_precompute_sets_audio_collect_cache(self) -> None:
+        if not torch.cuda.is_available():
+            pytest.skip("CUDA is required to exercise batched stop D2H precompute")
+
+        _, RState, _ = _voxcpm2_talker_mod()
+        talker = _make_bare_talker()
+        talker._audio_emit_every = 2
+        states = [
+            RState(request_id="stop", precomputed_stop_logits=torch.tensor([[0.0, 1.0]], device="cuda")),
+            RState(request_id="keep", precomputed_stop_logits=torch.tensor([[1.0, 0.0]], device="cuda")),
+        ]
+
+        talker._precompute_stop_flags_for_audio_collect(states)
+
+        assert states[0].precomputed_is_stopping is True
+        assert states[0].is_stopping is True
+        assert states[1].precomputed_is_stopping is False
+        assert states[1].is_stopping is False
+
     def test_on_requests_finished_defers_cleanup(self) -> None:
         talker = _make_bare_talker()
         _seed_cached_decode(talker, "req-A")
@@ -119,7 +209,7 @@ class TestLeakWarnGuard:
         talker = _make_bare_talker()
         talker._active_state_warn_threshold = 3
 
-        _, RState = _voxcpm2_talker_mod()
+        _, RState, _ = _voxcpm2_talker_mod()
         for i in range(4):
             talker._active_states[f"seed-{i}"] = RState(request_id=f"seed-{i}")
 
@@ -129,3 +219,548 @@ class TestLeakWarnGuard:
         leak_warnings = [m for m in calls if "cleanup path leak" in m]
         assert len(leak_warnings) == 1
         assert talker._active_state_warned is True
+
+
+class _NoopPerf:
+    def start(self, name: str) -> None:
+        pass
+
+    def stop(self, name: str) -> None:
+        pass
+
+
+class _FakeTTS:
+    feat_decoder = SimpleNamespace(estimator=SimpleNamespace(_compiled=False))
+
+    def fsq_layer(self, x: torch.Tensor) -> torch.Tensor:
+        return x + 10
+
+    def fusion_concat_proj(self, x: torch.Tensor) -> torch.Tensor:
+        half = x.shape[-1] // 2
+        return x[:, :half] * 2 + x[:, half:] * 3
+
+    def feat_encoder(self, x: torch.Tensor) -> torch.Tensor:
+        return x
+
+    def enc_to_lm_proj(self, x: torch.Tensor) -> torch.Tensor:
+        return x.squeeze(1).mean(dim=1).repeat(1, 4)
+
+
+class TestDecodeBatchContract:
+    def test_forward_skip_pending_prefill_keeps_audio_queue_aligned(self) -> None:
+        class FakeScaffold:
+            def __call__(self, input_ids, positions, intermediate_tensors, inputs_embeds):
+                return torch.zeros(input_ids.shape[0], 4)
+
+        talker = _make_bare_talker()
+        talker._perf = _NoopPerf()
+        talker._tts = _FakeTTS()
+        talker.model = FakeScaffold()
+        talker._enable_cuda_graph = False
+        talker._cuda_graph_warmup_steps = 0
+        talker._cuda_graph_warmup_threshold = 3
+        talker._max_cached_graphs = 4
+        talker._scaffold_graphs = {}
+        talker._residual_graphs = {}
+        talker._pending_requests = [("req", False, None, 1)]
+
+        out = talker.forward(
+            torch.tensor([1]),
+            torch.tensor([0]),
+            inputs_embeds=torch.zeros(1, 4),
+        )
+
+        torch.testing.assert_close(out, torch.zeros(1, 4))
+        assert talker._results_queue == [("req", None)]
+        assert talker._audio_queue == [("req", None)]
+
+    def _make_decode_talker(self) -> object:
+        talker = _make_bare_talker()
+        talker._perf = _NoopPerf()
+        talker._tts = _FakeTTS()
+        talker._side_dtype = torch.float32
+        talker._patch_size = 2
+        talker._feat_dim = 3
+
+        def dit_proj(lm_h: torch.Tensor, res_h: torch.Tensor) -> torch.Tensor:
+            return lm_h + 2 * res_h
+
+        def stop_fn(lm_h: torch.Tensor) -> torch.Tensor:
+            return torch.stack([lm_h.sum(dim=1), -lm_h.sum(dim=1)], dim=1)
+
+        def run_cfm(dit_h: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:
+            return dit_h[:, :1].view(-1, 1, 1) + cond.transpose(1, 2) * 0.1
+
+        talker._dit_proj_fn = dit_proj
+        talker._stop_fn = stop_fn
+        talker._run_cfm = run_cfm
+        talker._run_cfm_for_state = lambda _state, dit_h, cond: talker._run_cfm(dit_h, cond)
+        return talker
+
+    def _make_decode_state(self, request_id: str, offset: float):
+        _, RState, _ = _voxcpm2_talker_mod()
+        state = RState(request_id=request_id)
+        state.curr_prefix_feat_cond = torch.arange(6, dtype=torch.float32).reshape(2, 3) + offset
+        return state
+
+    def test_batched_fsq_fusion_matches_per_request_decode_prepare(self) -> None:
+        seq_talker = self._make_decode_talker()
+        batch_talker = self._make_decode_talker()
+        seq_talker._max_decode_steps = 2000
+        batch_talker._max_decode_steps = 2000
+
+        seq_states = [self._make_decode_state("req-0", 0.0), self._make_decode_state("req-1", 10.0)]
+        batch_states = [self._make_decode_state("req-0", 0.0), self._make_decode_state("req-1", 10.0)]
+        for idx, (seq_state, batch_state) in enumerate(zip(seq_states, batch_states, strict=True)):
+            prev = torch.full((1, 4), float(idx + 1))
+            seq_state.prev_feat_embed = prev.clone()
+            batch_state.prev_feat_embed = prev.clone()
+
+        hiddens = [
+            torch.tensor([[1.0, 2.0, 3.0, 4.0]]),
+            torch.tensor([[5.0, 6.0, 7.0, 8.0]]),
+        ]
+
+        expected = [
+            seq_talker._prepare_residual_decode(state, hidden, torch.device("cpu"))
+            for state, hidden in zip(seq_states, hiddens, strict=True)
+        ]
+        actual = batch_talker._prepare_residual_decode_batch(batch_states, hiddens, torch.device("cpu"))
+
+        for (expected_res, expected_meta), (actual_res, actual_meta) in zip(expected, actual, strict=True):
+            torch.testing.assert_close(actual_res, expected_res)
+            torch.testing.assert_close(actual_meta["new_lm_hidden"], expected_meta["new_lm_hidden"])
+        assert [state.decode_step_count for state in batch_states] == [1, 1]
+
+    def test_batch_decode_matches_per_request_state_updates(self) -> None:
+        seq_talker = self._make_decode_talker()
+        batch_talker = self._make_decode_talker()
+
+        seq_states = [self._make_decode_state("req-0", 0.0), self._make_decode_state("req-1", 10.0)]
+        batch_states = [self._make_decode_state("req-0", 0.0), self._make_decode_state("req-1", 10.0)]
+        metas = [
+            {"new_lm_hidden": torch.tensor([[1.0, 2.0, 3.0, 4.0]])},
+            {"new_lm_hidden": torch.tensor([[5.0, 6.0, 7.0, 8.0]])},
+        ]
+        batch_out = torch.tensor([[0.5, 1.0, 1.5, 2.0], [2.5, 3.0, 3.5, 4.0]])
+
+        for state, meta, res_out in zip(seq_states, metas, batch_out, strict=True):
+            seq_talker._finish_decode(state, meta, res_out.unsqueeze(0), torch.device("cpu"))
+
+        batch_talker._finish_decode_batch(
+            [(state, False, meta) for state, meta in zip(batch_states, metas, strict=True)],
+            batch_out,
+        )
+
+        for expected, actual in zip(seq_states, batch_states, strict=True):
+            torch.testing.assert_close(actual.precomputed_stop_logits, expected.precomputed_stop_logits)
+            torch.testing.assert_close(actual.curr_embed_for_next, expected.curr_embed_for_next)
+            torch.testing.assert_close(actual.prev_feat_embed, expected.prev_feat_embed)
+            torch.testing.assert_close(actual.curr_prefix_feat_cond, expected.curr_prefix_feat_cond)
+            torch.testing.assert_close(actual.last_audio_patch_gpu, expected.last_audio_patch_gpu)
+
+    def test_batch_decode_preserves_per_request_cfm_rng_order(self) -> None:
+        talker = self._make_decode_talker()
+        talker._enable_batched_cfm = False
+        states = [self._make_decode_state("req-0", 0.0), self._make_decode_state("req-1", 10.0)]
+        metas = [
+            {"new_lm_hidden": torch.tensor([[1.0, 2.0, 3.0, 4.0]])},
+            {"new_lm_hidden": torch.tensor([[5.0, 6.0, 7.0, 8.0]])},
+        ]
+        batch_out = torch.tensor([[0.5, 1.0, 1.5, 2.0], [2.5, 3.0, 3.5, 4.0]])
+
+        captured_noise: list[torch.Tensor] = []
+        torch.manual_seed(1234)
+        expected_noise = [torch.empty(1, 3, 2).normal_().detach().clone() for _ in states]
+
+        def run_cfm(
+            dit_h: torch.Tensor,
+            cond: torch.Tensor,
+        ) -> torch.Tensor:
+            captured_noise.append(torch.empty(1, 3, 2).normal_().detach().clone())
+            return dit_h[:, :1].view(-1, 1, 1) + cond.transpose(1, 2) * 0.1
+
+        talker._run_cfm = run_cfm
+
+        torch.manual_seed(1234)
+        talker._finish_decode_batch(
+            [(state, False, meta) for state, meta in zip(states, metas, strict=True)],
+            batch_out,
+        )
+
+        assert len(captured_noise) == len(states)
+        for actual, expected in zip(captured_noise, expected_noise, strict=True):
+            torch.testing.assert_close(actual, expected)
+
+
+class _FakeAudioVAE:
+    decode_chunk_size = 2
+
+
+class _FakeAudioTTS:
+    audio_vae = _FakeAudioVAE()
+
+
+class TestAudioEmitCoalescing:
+    def test_audio_emit_every_accumulates_gpu_chunks_without_changing_order(self) -> None:
+        _, RState, _ = _voxcpm2_talker_mod()
+        talker = _make_bare_talker()
+        talker._audio_emit_every = 2
+        talker._feat_dim = 1
+        talker._n_decode_pad_frames = 1
+        talker._device = torch.device("cpu")
+        talker._tts = _FakeAudioTTS()
+        talker._perf = _NoopPerf()
+
+        values = [torch.tensor([10.0, 11.0]), torch.tensor([99.0, 99.0, 20.0, 21.0])]
+
+        def run_vae_decode(feat: torch.Tensor) -> torch.Tensor:
+            return values.pop(0)
+
+        talker._run_vae_decode = run_vae_decode
+        state = RState(request_id="req")
+
+        state.last_audio_patch_gpu = torch.tensor([[1.0]])
+        assert talker._collect_audio(state) is None
+        assert len(state.pending_audio_chunks_gpu) == 1
+
+        state.last_audio_patch_gpu = torch.tensor([[2.0]])
+        audio = talker._collect_audio(state)
+
+        torch.testing.assert_close(audio, torch.tensor([10.0, 11.0, 20.0, 21.0]))
+        assert state.pending_audio_chunks_gpu == []
+
+    def test_audio_emit_every_flushes_on_stop(self) -> None:
+        _, RState, _ = _voxcpm2_talker_mod()
+        talker = _make_bare_talker()
+        talker._audio_emit_every = 8
+        talker._feat_dim = 1
+        talker._n_decode_pad_frames = 1
+        talker._device = torch.device("cpu")
+        talker._tts = _FakeAudioTTS()
+        talker._perf = _NoopPerf()
+        talker._run_vae_decode = lambda feat: torch.tensor([3.0, 4.0])
+        state = RState(request_id="req", precomputed_stop_logits=torch.tensor([[0.0, 1.0]]))
+
+        state.last_audio_patch_gpu = torch.tensor([[1.0]])
+        audio = talker._collect_audio(state)
+
+        torch.testing.assert_close(audio, torch.tensor([3.0, 4.0]))
+        assert state.pending_audio_chunks_gpu == []
+
+    def test_vae_decode_every_accumulates_latents_before_decode(self) -> None:
+        _, RState, _ = _voxcpm2_talker_mod()
+        talker = _make_bare_talker()
+        talker._vae_decode_every = 2
+        talker._feat_dim = 1
+        talker._n_decode_pad_frames = 1
+        talker._device = torch.device("cpu")
+        talker._tts = _FakeAudioTTS()
+        talker._perf = _NoopPerf()
+
+        seen_feats = []
+        values = [torch.tensor([10.0, 11.0, 20.0, 21.0])]
+
+        def run_vae_decode(feat: torch.Tensor) -> torch.Tensor:
+            seen_feats.append(feat.detach().clone())
+            return values.pop(0)
+
+        talker._run_vae_decode = run_vae_decode
+        state = RState(request_id="req")
+
+        state.last_audio_patch_gpu = torch.tensor([[1.0]])
+        assert talker._collect_audio(state) is None
+        assert len(state.pending_vae_latents_gpu) == 1
+        assert seen_feats == []
+
+        state.last_audio_patch_gpu = torch.tensor([[2.0]])
+        audio = talker._collect_audio(state)
+
+        torch.testing.assert_close(audio, torch.tensor([10.0, 11.0, 20.0, 21.0]))
+        assert state.pending_vae_latents_gpu == []
+        torch.testing.assert_close(seen_feats[0].reshape(-1), torch.tensor([1.0, 2.0]))
+
+    def test_vae_decode_every_flushes_pending_latent_on_stop(self) -> None:
+        _, RState, _ = _voxcpm2_talker_mod()
+        talker = _make_bare_talker()
+        talker._vae_decode_every = 4
+        talker._feat_dim = 1
+        talker._n_decode_pad_frames = 1
+        talker._device = torch.device("cpu")
+        talker._tts = _FakeAudioTTS()
+        talker._perf = _NoopPerf()
+        talker._run_vae_decode = lambda feat: torch.tensor([7.0, 8.0])
+        state = RState(request_id="req", precomputed_stop_logits=torch.tensor([[0.0, 1.0]]))
+
+        state.last_audio_patch_gpu = torch.tensor([[1.0]])
+        audio = talker._collect_audio(state)
+
+        torch.testing.assert_close(audio, torch.tensor([7.0, 8.0]))
+        assert state.pending_vae_latents_gpu == []
+        assert state.is_stopping is True
+
+    def test_batched_vae_decode_collects_ready_requests_in_order(self) -> None:
+        _, RState, _ = _voxcpm2_talker_mod()
+        talker = _make_bare_talker()
+        talker._enable_batched_vae_decode = True
+        talker._coalesce_audio_d2h = True
+        talker._vae_decode_every = 2
+        talker._feat_dim = 1
+        talker._n_decode_pad_frames = 1
+        talker._device = torch.device("cpu")
+        talker._tts = _FakeAudioTTS()
+        talker._perf = _NoopPerf()
+
+        seen_feats: list[torch.Tensor] = []
+
+        def run_vae_decode(feat: torch.Tensor) -> torch.Tensor:
+            seen_feats.append(feat.detach().clone())
+            # One output row per request. decode_chunk_size=2, and each ready
+            # request has two new latent frames, so the full decoded length is 4.
+            return torch.tensor(
+                [
+                    [[10.0, 11.0, 12.0, 13.0]],
+                    [[20.0, 21.0, 22.0, 23.0]],
+                ]
+            )
+
+        talker._run_vae_decode = run_vae_decode
+
+        ready0 = RState(request_id="ready-0")
+        ready0.pending_vae_latents_gpu.append(torch.tensor([[1.0]]))
+        ready0.last_audio_patch_gpu = torch.tensor([[2.0]])
+
+        pending = RState(request_id="pending")
+        pending.last_audio_patch_gpu = torch.tensor([[99.0]])
+
+        ready1 = RState(request_id="ready-1")
+        ready1.pending_vae_latents_gpu.append(torch.tensor([[3.0]]))
+        ready1.last_audio_patch_gpu = torch.tensor([[4.0]])
+
+        outputs = talker._collect_audio_batch([ready0, pending, ready1])
+
+        assert list(outputs) == ["ready-0", "pending", "ready-1"]
+        torch.testing.assert_close(outputs["ready-0"], torch.tensor([10.0, 11.0, 12.0, 13.0]))
+        assert outputs["pending"] is None
+        torch.testing.assert_close(outputs["ready-1"], torch.tensor([20.0, 21.0, 22.0, 23.0]))
+        assert len(seen_feats) == 1
+        torch.testing.assert_close(seen_feats[0].reshape(2, -1), torch.tensor([[1.0, 2.0], [3.0, 4.0]]))
+        assert pending.pending_vae_latents_gpu
+        assert ready0.pending_vae_latents_gpu == []
+        assert ready1.pending_vae_latents_gpu == []
+
+    def test_audio_emit_every_keeps_private_chunk_when_vae_output_storage_is_reused(self) -> None:
+        _, RState, _ = _voxcpm2_talker_mod()
+        talker = _make_bare_talker()
+        talker._audio_emit_every = 2
+        talker._enable_vae_cuda_graph = True
+        talker._feat_dim = 1
+        talker._n_decode_pad_frames = 1
+        talker._device = torch.device("cpu")
+        talker._tts = _FakeAudioTTS()
+        talker._perf = _NoopPerf()
+
+        shared_output = torch.empty(4)
+        values = [torch.tensor([10.0, 11.0, -1.0, -1.0]), torch.tensor([99.0, 99.0, 20.0, 21.0])]
+
+        def run_vae_decode(feat: torch.Tensor) -> torch.Tensor:
+            shared_output.copy_(values.pop(0))
+            return shared_output
+
+        talker._run_vae_decode = run_vae_decode
+        state = RState(request_id="req")
+
+        state.last_audio_patch_gpu = torch.tensor([[1.0]])
+        assert talker._collect_audio(state) is None
+
+        state.last_audio_patch_gpu = torch.tensor([[2.0]])
+        audio = talker._collect_audio(state)
+
+        torch.testing.assert_close(audio, torch.tensor([10.0, 11.0, 20.0, 21.0]))
+
+    def test_audio_emit_every_keeps_private_chunk_when_compiled_vae_reuses_storage(self) -> None:
+        _, RState, _ = _voxcpm2_talker_mod()
+        talker = _make_bare_talker()
+        talker._audio_emit_every = 2
+        talker._feat_dim = 1
+        talker._n_decode_pad_frames = 1
+        talker._device = torch.device("cpu")
+        talker._tts = _FakeAudioTTS()
+        talker._tts.audio_vae._compiled = True
+        talker._perf = _NoopPerf()
+
+        shared_output = torch.empty(4)
+        values = [torch.tensor([10.0, 11.0, -1.0, -1.0]), torch.tensor([99.0, 99.0, 20.0, 21.0])]
+
+        def run_vae_decode(feat: torch.Tensor) -> torch.Tensor:
+            shared_output.copy_(values.pop(0))
+            return shared_output
+
+        talker._run_vae_decode = run_vae_decode
+        state = RState(request_id="req")
+
+        state.last_audio_patch_gpu = torch.tensor([[1.0]])
+        assert talker._collect_audio(state) is None
+
+        state.last_audio_patch_gpu = torch.tensor([[2.0]])
+        audio = talker._collect_audio(state)
+
+        torch.testing.assert_close(audio, torch.tensor([10.0, 11.0, 20.0, 21.0]))
+
+    def test_delayed_audio_copy_emits_previous_chunk_and_flushes_on_stop(self) -> None:
+        _, RState, _ = _voxcpm2_talker_mod()
+        talker = _make_bare_talker()
+        talker._enable_delayed_audio_copy = True
+        talker._audio_emit_every = 1
+        talker._feat_dim = 1
+        talker._n_decode_pad_frames = 1
+        talker._device = torch.device("cpu")
+        talker._tts = _FakeAudioTTS()
+        talker._perf = _NoopPerf()
+
+        values = [
+            torch.tensor([10.0, 11.0]),
+            torch.tensor([99.0, 99.0, 20.0, 21.0]),
+            torch.tensor([88.0, 88.0, 30.0, 31.0]),
+        ]
+
+        def run_vae_decode(feat: torch.Tensor) -> torch.Tensor:
+            return values.pop(0)
+
+        talker._run_vae_decode = run_vae_decode
+        state = RState(request_id="req")
+
+        state.last_audio_patch_gpu = torch.tensor([[1.0]])
+        assert talker._collect_audio(state) is None
+        assert len(state.pending_audio_copies) == 1
+
+        state.last_audio_patch_gpu = torch.tensor([[2.0]])
+        audio = talker._collect_audio(state)
+        torch.testing.assert_close(audio, torch.tensor([10.0, 11.0]))
+        assert len(state.pending_audio_copies) == 1
+
+        state.precomputed_stop_logits = torch.tensor([[0.0, 1.0]])
+        state.last_audio_patch_gpu = torch.tensor([[3.0]])
+        audio = talker._collect_audio(state)
+        torch.testing.assert_close(audio, torch.tensor([20.0, 21.0, 30.0, 31.0]))
+        assert state.pending_audio_copies == []
+        assert state.is_stopping is True
+
+    def test_delayed_audio_copy_flushes_on_forced_stop(self) -> None:
+        _, RState, _ = _voxcpm2_talker_mod()
+        talker = _make_bare_talker()
+        talker._enable_delayed_audio_copy = True
+        talker._audio_emit_every = 1
+        talker._feat_dim = 1
+        talker._n_decode_pad_frames = 1
+        talker._device = torch.device("cpu")
+        talker._tts = _FakeAudioTTS()
+        talker._perf = _NoopPerf()
+
+        values = [
+            torch.tensor([10.0, 11.0]),
+            torch.tensor([99.0, 99.0, 20.0, 21.0]),
+        ]
+
+        def run_vae_decode(feat: torch.Tensor) -> torch.Tensor:
+            return values.pop(0)
+
+        talker._run_vae_decode = run_vae_decode
+        state = RState(request_id="req")
+
+        state.last_audio_patch_gpu = torch.tensor([[1.0]])
+        assert talker._collect_audio(state) is None
+
+        state.is_stopping = True
+        state.last_audio_patch_gpu = torch.tensor([[2.0]])
+        audio = talker._collect_audio(state)
+
+        torch.testing.assert_close(audio, torch.tensor([10.0, 11.0, 20.0, 21.0]))
+        assert state.pending_audio_copies == []
+
+    def test_make_omni_output_uses_sparse_req_ids_for_deferred_chunks(self) -> None:
+        talker = _make_bare_talker()
+        talker._audio_emit_every = 4
+        talker._sample_rate = 16000
+        talker._last_audio_output_req_ids = ["req-0", "req-1", "req-2"]
+        talker._audio_queue = [
+            ("req-0", None),
+            ("req-1", torch.tensor([1.0, 2.0])),
+            ("req-2", None),
+        ]
+
+        out = talker.make_omni_output(torch.zeros(3, 1))
+
+        audio = out.multimodal_outputs["model_outputs"]
+        sr = out.multimodal_outputs["sr"]
+        assert out.multimodal_outputs["meta"]["req_id"] == ["req-1"]
+        assert out.multimodal_outputs["meta"]["sparse_audio"] == ["1"]
+        assert len(audio) == 1
+        torch.testing.assert_close(audio[0], torch.tensor([1.0, 2.0]))
+        assert int(sr[0]) == 16000
+
+    def test_make_omni_output_emits_empty_sparse_marker_for_deferred_chunks(self) -> None:
+        talker = _make_bare_talker()
+        talker._audio_emit_every = 4
+        talker._sample_rate = 16000
+        talker._audio_queue = [("req-0", None), ("req-1", None)]
+
+        out = talker.make_omni_output(torch.zeros(2, 1))
+
+        assert out.multimodal_outputs["model_outputs"] == []
+        assert out.multimodal_outputs["sr"] == []
+        assert out.multimodal_outputs["meta"]["req_id"] == []
+        assert out.multimodal_outputs["meta"]["sparse_audio"] == ["1"]
+
+    def test_make_omni_output_uses_sparse_req_ids_for_delayed_audio_copy(self) -> None:
+        talker = _make_bare_talker()
+        talker._enable_delayed_audio_copy = True
+        talker._sample_rate = 16000
+        talker._audio_queue = [
+            ("req-0", None),
+            ("req-1", torch.tensor([1.0, 2.0])),
+            ("req-2", None),
+        ]
+
+        out = talker.make_omni_output(torch.zeros(3, 1))
+
+        assert out.multimodal_outputs["meta"]["req_id"] == ["req-1"]
+        assert out.multimodal_outputs["meta"]["sparse_audio"] == ["1"]
+        assert len(out.multimodal_outputs["model_outputs"]) == 1
+        torch.testing.assert_close(out.multimodal_outputs["model_outputs"][0], torch.tensor([1.0, 2.0]))
+
+    def test_make_omni_output_uses_sparse_req_ids_for_deferred_vae_decode(self) -> None:
+        talker = _make_bare_talker()
+        talker._vae_decode_every = 2
+        talker._sample_rate = 16000
+        talker._audio_queue = [
+            ("req-0", None),
+            ("req-1", torch.tensor([1.0, 2.0])),
+            ("req-2", None),
+        ]
+
+        out = talker.make_omni_output(torch.zeros(3, 1))
+
+        assert out.multimodal_outputs["meta"]["req_id"] == ["req-1"]
+        assert out.multimodal_outputs["meta"]["sparse_audio"] == ["1"]
+        assert len(out.multimodal_outputs["model_outputs"]) == 1
+        torch.testing.assert_close(out.multimodal_outputs["model_outputs"][0], torch.tensor([1.0, 2.0]))
+
+    def test_make_omni_output_coalesces_gpu_audio_to_cpu(self) -> None:
+        if not torch.cuda.is_available():
+            pytest.skip("CUDA is required to exercise D2H coalescing")
+        talker = _make_bare_talker()
+        talker._coalesce_audio_d2h = True
+        talker._sample_rate = 16000
+        talker._audio_queue = [
+            ("req-0", torch.tensor([1.0, 2.0], device="cuda")),
+            ("req-1", torch.tensor([3.0], device="cuda")),
+        ]
+
+        out = talker.make_omni_output(torch.zeros(2, 1))
+
+        audio = out.multimodal_outputs["model_outputs"]
+        assert len(audio) == 2
+        assert all(not chunk.is_cuda for chunk in audio)
+        torch.testing.assert_close(audio[0], torch.tensor([1.0, 2.0]))
+        torch.testing.assert_close(audio[1], torch.tensor([3.0]))

tests/worker/test_gpu_ar_model_runner.py

@@ -44,6 +44,14 @@ def test_resolve_pooler_payload_req_ids_downstream_stage_uses_filtered_requests(
     assert payload_req_ids == ["r2"]
 
 
+def test_sparse_mm_req_ids_requires_sparse_audio_marker():
+    assert GPUARModelRunner._sparse_mm_req_ids({"meta": {"req_id": ["r1"]}}) is None
+    assert GPUARModelRunner._sparse_mm_req_ids({"meta.req_id": ["r1"]}) is None
+
+    assert GPUARModelRunner._sparse_mm_req_ids({"meta": {"req_id": ["r1"], "sparse_audio": ["1"]}}) == ["r1"]
+    assert GPUARModelRunner._sparse_mm_req_ids({"meta.req_id": ["r1"], "meta.sparse_audio": ["1"]}) == ["r1"]
+
+
 @pytest.mark.parametrize("query_start_loc_attr", ["method", "tensor_attr"])
 def test_sample_tokens_tail_only_prefix_cache_uses_staged_cpu_hidden_states(monkeypatch, query_start_loc_attr):
     runner = object.__new__(GPUARModelRunner)

vllm_omni/deploy/voxcpm2.yaml

@@ -1,19 +1,13 @@
 # VoxCPM2 deploy: single-stage AR pipeline with per-request state batching.
 # Verified on 1x H20 141GB.
 #
-# Fields omitted from a stage fall back to StageDeployConfig defaults (see
-# vllm_omni/config/stage_config.py). Notable choices:
-#   * enforce_eager=true because the talker's KV-cache save/restore loop is
-#     not cudagraph-compatible (captured separately via voxcpm2_talker's own
-#     _CapturedGraph path).
-#   * max_num_seqs=4 matches the legacy max_batch_size=4 contract: concurrent
-#     requests are supported via per-request StaticKVCache save/restore.
+# VoxCPM2-local runtime knobs are passed through hf_overrides.
 async_chunk: false
 dtype: bfloat16
 
 stages:
   - stage_id: 0
-    max_num_seqs: 4
+    max_num_seqs: 32
     gpu_memory_utilization: 0.9
     enforce_eager: true
     trust_remote_code: true
@@ -22,6 +16,18 @@ stages:
     max_num_batched_tokens: 4096
     max_model_len: 4096
     devices: "0"
+    engine_extras:
+      hf_overrides:
+        voxcpm2_runtime_config:
+          audio_emit_every: 1
+          vae_decode_every: 3
+          enable_batched_cfm: true
+          enable_batched_vae_decode: true
+          coalesce_audio_d2h: true
+          enable_loc_dit_fused_qkv: true
+          enable_loc_dit_fused_mlp: true
+          enable_loc_dit_zero_dt_cache: true
+          enable_loc_dit_skip_qkv_contig: true
     default_sampling_params:
       temperature: 0.0
       top_p: 1.0

vllm_omni/worker/gpu_ar_model_runner.py

@@ -200,6 +200,33 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
             downstream_req_ids = req_ids_output_copy
         return engine_output_type, downstream_req_ids
 
+    @staticmethod
+    def _sparse_mm_req_ids(multimodal_outputs: Any) -> list[str] | None:
+        if not isinstance(multimodal_outputs, dict):
+            return None
+        meta = multimodal_outputs.get("meta")
+        req_ids = None
+        sparse_audio = False
+        if isinstance(meta, dict):
+            req_ids = meta.get("req_id")
+            sparse_audio = GPUARModelRunner._is_sparse_audio_marker(meta.get("sparse_audio"))
+        if req_ids is None:
+            req_ids = multimodal_outputs.get("meta.req_id")
+            sparse_audio = GPUARModelRunner._is_sparse_audio_marker(multimodal_outputs.get("meta.sparse_audio"))
+        if not sparse_audio:
+            return None
+        if not isinstance(req_ids, list):
+            return None
+        return [rid for rid in req_ids if isinstance(rid, str)]
+
+    @staticmethod
+    def _is_sparse_audio_marker(value: Any) -> bool:
+        if isinstance(value, list):
+            return any(str(item).lower() in ("1", "true", "yes", "on") for item in value)
+        if isinstance(value, str):
+            return value.lower() in ("1", "true", "yes", "on")
+        return bool(value)
+
     def capture_model(self) -> int:
         result = super().capture_model()
         self._capture_talker_mtp_graphs()
@@ -978,10 +1005,16 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
         self.kv_connector_output = None
 
         engine_output_type, downstream_req_ids = self._resolve_pooler_payload_req_ids(req_ids_output_copy)
+        sparse_mm_req_ids = self._sparse_mm_req_ids(multimodal_outputs)
+        sparse_mm_index = {rid: i for i, rid in enumerate(sparse_mm_req_ids or [])}
+        if engine_output_type == "audio" and sparse_mm_req_ids is not None:
+            sparse_req_id_set = set(sparse_mm_req_ids)
+            downstream_req_ids = [rid for rid in req_ids_output_copy if rid in sparse_req_id_set]
         needs_pooler_payload = len(downstream_req_ids) > 0
         downstream_req_id_set = set(downstream_req_ids)
         hidden_states_cpu = None
         req_hidden_states_cpu: dict[str, torch.Tensor] | None = None
+        audio_sparse_output = engine_output_type == "audio" and sparse_mm_req_ids is not None
         needs_scheduled_hidden_payload = needs_pooler_payload and (
             self.omni_prefix_cache is None or not self._model_needs_full_prefix_hidden_states()
         )
@@ -994,7 +1027,9 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
                 int(scheduler_output.total_num_scheduled_tokens),
                 int(hidden_states.shape[0]),
             )
-            if len(downstream_req_ids) == len(req_ids_output_copy):
+            if audio_sparse_output:
+                pass
+            elif len(downstream_req_ids) == len(req_ids_output_copy):
                 hidden_states_cpu = hidden_states[:num_valid_tokens].detach().to("cpu").contiguous()
             else:
                 req_hidden_states_cpu = {}
@@ -1024,8 +1059,6 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
         pooler_output: list[dict[str, object]] | None = None
         if needs_pooler_payload:
             mm_cpu = None
-            # Prior to applying the post-processing func, extract
-            # the prefix cached hidden states and multimodal states.
             if self.omni_prefix_cache is not None:
                 (
                     combined_hidden_states,
@@ -1036,7 +1069,6 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
                     multimodal_outputs,
                     scheduler_output.num_scheduled_tokens,
                 )
-            # Otherwise we don't have the mm CPU data yet, so we still need to build it
             if self.omni_prefix_cache is None or combined_multimodal_outputs is None:
                 mm_cpu = build_mm_cpu(flatten_payload(multimodal_outputs) if multimodal_outputs else multimodal_outputs)
 
@@ -1059,7 +1091,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
                     req_hidden_states_cpu[rid] = hidden_states[start:end].detach().to("cpu").contiguous()
 
             pooler_output = []
-            for out_idx, rid in enumerate(req_ids_output_copy):
+            for rid in req_ids_output_copy:
                 if rid not in downstream_req_id_set:
                     pooler_output.append({})
                     continue
@@ -1067,31 +1099,24 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
                 start = int(query_start_loc_cpu[idx])
                 sched = int(num_scheduled_tokens_np[idx])
                 end = start + sched
-                # If prefix cache is enabled, we have already split everything
-                # by request and converted the states to CPU tensors
-                if req_hidden_states_cpu is not None and combined_hidden_states is None:
-                    req_hidden_states = req_hidden_states_cpu[rid]
-                else:
-                    req_hidden_states = self._resolve_req_hidden_states(
-                        hidden_states_cpu,
-                        combined_hidden_states,
-                        rid,
-                        start,
-                        end,
-                    )
-                payload: dict[str, object] = {"hidden": req_hidden_states}
+                payload: dict[str, object] = {}
+                if not audio_sparse_output:
+                    if req_hidden_states_cpu is not None and combined_hidden_states is None:
+                        req_hidden_states = req_hidden_states_cpu[rid]
+                    else:
+                        req_hidden_states = self._resolve_req_hidden_states(
+                            hidden_states_cpu,
+                            combined_hidden_states,
+                            rid,
+                            start,
+                            end,
+                        )
+                    payload["hidden"] = req_hidden_states
 
                 mm_payload: dict[str, object] = {}
                 if combined_multimodal_outputs or mm_cpu:
                     if combined_multimodal_outputs:
-                        # Prefix cache enabled; all items have already been processed
-                        # and split apart for each request as needed, and all tensors
-                        # have already been detached to the CPU.  Lists are kept as
-                        # passthrough data for consistent behavior in postprocess.
-                        # Recurse into nested dicts so list-valued sub-keys (e.g.
-                        # embed.tts_bos = [tensor]) are unwrapped to bare tensors
-                        # at the leaves; downstream flatten_payload then yields a
-                        # wire-clean dict[str, torch.Tensor].
+
                         def _unwrap_lists(v):
                             if isinstance(v, list):
                                 return v[idx] if idx < len(v) else v[0]
@@ -1103,8 +1128,26 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
                             mm_payload[mm_key] = _unwrap_lists(combined_multimodal_outputs[mm_key][rid])
 
                     else:
-                        # Prefix cache disabled; we still need to process the data
                         for mm_key, mm_val in mm_cpu.items():
+                            if mm_key in {"meta.req_id", "meta.sparse_audio"}:
+                                continue
+                            if audio_sparse_output and isinstance(mm_val, list):
+                                sparse_idx = sparse_mm_index.get(rid)
+                                if sparse_idx is None:
+                                    continue
+                                if sparse_idx >= len(mm_val):
+                                    logger.warning(
+                                        "Sparse multimodal payload mismatch for request %s: index %d >= %d.",
+                                        rid,
+                                        sparse_idx,
+                                        len(mm_val),
+                                    )
+                                    continue
+                                sparse_val = mm_val[sparse_idx]
+                                mm_payload[mm_key] = (
+                                    sparse_val.clone() if isinstance(sparse_val, torch.Tensor) else sparse_val
+                                )
+                                continue
                             mm_payload[mm_key] = to_payload_element(
                                 element=mm_val,
                                 idx=idx,
@@ -1114,8 +1157,6 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
                                 seq_len=seq_len,
                             )
                     payload.update(mm_payload)
-                # Flatten nested dicts to dotted keys so pooling_output
-                # stays dict[str, torch.Tensor] for msgspec serialization.
                 pooler_output.append(flatten_payload(payload))
 
         if pooler_output and self._should_accumulate_full_payload_output():