diff --git a/tests/model_executor/models/voxcpm2/test_talker_state_eviction.py b/tests/model_executor/models/voxcpm2/test_talker_state_eviction.py index 22d1d2463..28da0b1bc 100644 --- a/tests/model_executor/models/voxcpm2/test_talker_state_eviction.py +++ b/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])) diff --git a/tests/worker/test_gpu_ar_model_runner.py b/tests/worker/test_gpu_ar_model_runner.py index c36600be8..8c6accad5 100644 --- a/tests/worker/test_gpu_ar_model_runner.py +++ b/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) diff --git a/vllm_omni/deploy/voxcpm2.yaml b/vllm_omni/deploy/voxcpm2.yaml index 71ef14824..4e3f7f3f6 100644 --- a/vllm_omni/deploy/voxcpm2.yaml +++ b/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 diff --git a/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py b/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py index 8b60457b0..978b56668 100644 --- a/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py +++ b/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py @@ -12,11 +12,12 @@ from __future__ import annotations import copy import dataclasses +import hashlib import logging import math -import os import time from collections.abc import Iterable +from types import MethodType from typing import Any import torch @@ -46,7 +47,7 @@ from .voxcpm2_import_utils import import_voxcpm2_core logger = init_logger(__name__) -_ENABLE_PROFILING = os.environ.get("VOXCPM2_PROFILE", "0") == "1" +_ENABLE_NVTX_PROFILE = False # Lower bound for the _active_states leak-warn threshold. The effective # threshold is max(_ACTIVE_STATE_LEAK_WARN_MIN, 4 * max_batch_size) so small @@ -54,6 +55,105 @@ _ENABLE_PROFILING = os.environ.get("VOXCPM2_PROFILE", "0") == "1" _ACTIVE_STATE_LEAK_WARN_MIN = 512 +@dataclasses.dataclass(frozen=True) +class _VoxCPM2RuntimeConfig: + enable_profiling: bool = False + enable_nvtx_profile: bool = False + enable_loc_dit_layer_nvtx: bool = False + enable_loc_dit_fused_qkv: bool = True + enable_loc_dit_fused_mlp: bool = True + enable_loc_dit_zero_dt_cache: bool = True + enable_loc_dit_fast_rope: bool = False + enable_loc_dit_skip_qkv_contig: bool = True + enable_loc_dit_reduce_overhead_no_cg: bool = False + enable_loc_dit_fullgraph_no_cg: bool = False + cfg_cutoff_ratio: float = 1.0 + decode_graph_capture_policy: str = "all" + enable_vae_cuda_graph: bool = False + enable_cfm_cuda_graph: bool = False + enable_cfm_prealloc_output: bool = False + enable_batched_cfm: bool = True + deterministic_cfm_noise: bool = False + deterministic_cfm_seed: int = 20260601 + audio_emit_every: int = 1 + vae_decode_every: int = 3 + enable_delayed_audio_copy: bool = False + delayed_audio_copy_use_events: bool = False + coalesce_audio_d2h: bool = True + enable_batched_vae_decode: bool = True + enable_batched_fsq_fusion: bool = False + batched_fsq_fusion_max_batch: int = 32 + + @classmethod + def from_vllm_config(cls, vllm_config: VllmConfig) -> _VoxCPM2RuntimeConfig: + model_config = vllm_config.model_config + raw = getattr(model_config, "voxcpm2_runtime_config", None) + if raw is None: + raw = getattr(getattr(model_config, "hf_config", None), "voxcpm2_runtime_config", None) + if raw is None: + return cls() + if hasattr(raw, "to_dict"): + raw = raw.to_dict() + elif not isinstance(raw, dict) and hasattr(raw, "__dict__"): + raw = vars(raw) + if not isinstance(raw, dict): + logger.warning("Ignoring invalid voxcpm2_runtime_config=%r; expected a dict.", raw) + return cls() + + fields = {field.name: field for field in dataclasses.fields(cls)} + values: dict[str, Any] = {} + for key, value in raw.items(): + if key not in fields: + logger.warning("Ignoring unknown VoxCPM2 runtime config key: %s", key) + continue + default = getattr(cls(), key) + values[key] = cls._coerce_value(key, value, default) + return cls(**values)._normalized() + + def _normalized(self) -> _VoxCPM2RuntimeConfig: + cfg = self + if cfg.enable_batched_cfm and cfg.enable_cfm_cuda_graph: + logger.warning( + "VoxCPM2 batched CFM and CFM CUDA Graph are mutually exclusive; " + "disabling CFM CUDA Graph and keeping batched CFM." + ) + cfg = dataclasses.replace(cfg, enable_cfm_cuda_graph=False) + if cfg.enable_cfm_prealloc_output and not cfg.enable_cfm_cuda_graph: + logger.warning("VoxCPM2 CFM preallocated output requires CFM CUDA Graph; disabling it.") + cfg = dataclasses.replace(cfg, enable_cfm_prealloc_output=False) + if cfg.delayed_audio_copy_use_events and not cfg.enable_delayed_audio_copy: + logger.warning("VoxCPM2 delayed audio copy events require delayed audio copy; disabling event polling.") + cfg = dataclasses.replace(cfg, delayed_audio_copy_use_events=False) + if cfg.enable_batched_vae_decode and ( + cfg.enable_delayed_audio_copy or cfg.audio_emit_every > 1 or cfg.enable_vae_cuda_graph + ): + logger.warning( + "VoxCPM2 batched VAE decode is incompatible with delayed audio copy, " + "audio_emit_every > 1, and VAE CUDA Graph; disabling batched VAE decode." + ) + cfg = dataclasses.replace(cfg, enable_batched_vae_decode=False) + return cfg + + @staticmethod + def _coerce_value(key: str, value: Any, default: Any) -> Any: + if isinstance(default, bool): + if isinstance(value, str): + return value.strip().lower() in ("1", "true", "yes", "on") + return bool(value) + if isinstance(default, int) and not isinstance(default, bool): + value = int(value) + if key in {"audio_emit_every", "vae_decode_every", "batched_fsq_fusion_max_batch"}: + return max(1, value) + return value + if isinstance(default, float): + if key == "cfg_cutoff_ratio": + return min(1.0, max(0.0, float(value))) + return float(value) + if isinstance(default, str): + return str(value) + return value + + def is_cjk_char(c: str) -> bool: """Check if a character is a CJK ideograph.""" cp = ord(c) @@ -198,7 +298,12 @@ class _RequestState: prev_feat_embed: torch.Tensor | None = None curr_prefix_feat_cond: torch.Tensor | None = None last_audio_patch_gpu: torch.Tensor | None = None + cfm_output_gpu: torch.Tensor | None = None + cfm_noise_step: int = 0 precomputed_stop_logits: torch.Tensor | None = None + pending_audio_chunks_gpu: list[torch.Tensor] = dataclasses.field(default_factory=list) + pending_audio_copies: list[_PendingAudioCopy] = dataclasses.field(default_factory=list) + pending_vae_latents_gpu: list[torch.Tensor] = dataclasses.field(default_factory=list) # Rolling tail of previously-decoded latents used as VAE receptive-field context. # Shape (n_pad_frames, feat_dim) on GPU. None before first decode. decode_pad: torch.Tensor | None = None @@ -209,6 +314,7 @@ class _RequestState: prompt_cache: dict | None = None prefill_masks: tuple | None = None is_stopping: bool = False + precomputed_is_stopping: bool | None = None @dataclasses.dataclass @@ -219,6 +325,31 @@ class _CapturedGraph: output: torch.Tensor +@dataclasses.dataclass +class _CapturedVAEGraph: + graph: torch.cuda.CUDAGraph + input_feat: torch.Tensor + output: torch.Tensor + + +@dataclasses.dataclass +class _CapturedCFMGraph: + graph: torch.cuda.CUDAGraph + mu: torch.Tensor + cond: torch.Tensor + noise: torch.Tensor + output: torch.Tensor + buffers: _CFMBufferManager + + +@dataclasses.dataclass +class _PendingAudioCopy: + host: torch.Tensor + event: torch.cuda.Event | None = None + source: torch.Tensor | None = None + async_copy: bool = False + + # =================================================================== # Profiling timer # =================================================================== @@ -283,6 +414,250 @@ class _PerfTimer: self._pairs.clear() +class _NvtxRange: + __slots__ = ("_enabled", "_entered") + + def __init__(self, name: str): + self._enabled = _ENABLE_NVTX_PROFILE and torch.cuda.is_available() + self._entered = False + if self._enabled: + torch.cuda.nvtx.range_push(name) + self._entered = True + + def __enter__(self) -> None: + return None + + def __exit__(self, exc_type: Any, exc: Any, tb: Any) -> None: + if self._entered: + torch.cuda.nvtx.range_pop() + + +def _install_locdit_fused_qkv( + estimator: nn.Module, + *, + enable_fast_rope: bool, + skip_qkv_contig: bool, +) -> int: + """Patch native VoxCPM LocDiT attention modules to use one QKV projection. + + Native LocDiT uses separate q/k/v Linear layers before PyTorch SDPA. This + keeps the attention backend unchanged and only replaces the projection + front-end with a pre-concatenated weight, mirroring the AR-side wrapper. + """ + decoder = getattr(estimator, "decoder", None) + layers = getattr(decoder, "layers", None) + if layers is None: + return 0 + patched = 0 + + def _rotate_half(x: torch.Tensor) -> torch.Tensor: + x1, x2 = x.chunk(2, dim=-1) + return torch.cat((-x2, x1), dim=-1) + + def _apply_rotary_pos_emb_fast_dtype( + q: torch.Tensor, + k: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + cos = cos.to(q.dtype) + sin = sin.to(q.dtype) + return (q * cos) + (_rotate_half(q) * sin), (k * cos) + (_rotate_half(k) * sin) + + for layer in layers: + attn = getattr(layer, "self_attn", None) + if attn is None or getattr(attn, "_voxcpm2_fused_qkv", False): + continue + apply_rotary = type(attn).forward.__globals__.get("apply_rotary_pos_emb") + if apply_rotary is None: + continue + weight = torch.cat([attn.q_proj.weight, attn.k_proj.weight, attn.v_proj.weight], dim=0).detach() + attn.register_buffer("_voxcpm2_fused_qkv_weight", weight, persistent=False) + + def _forward( + self, + hidden_states: torch.Tensor, + position_emb: tuple[torch.Tensor, torch.Tensor] | None, + is_causal: bool, + _apply_rotary=apply_rotary, + ): + bsz, q_len, _ = hidden_states.size() + qkv = nn.functional.linear(hidden_states, self._voxcpm2_fused_qkv_weight) + q_size = self.num_heads * self.head_dim + kv_size = self.num_key_value_heads * self.head_dim + query_states, key_states, value_states = qkv.split([q_size, kv_size, kv_size], dim=-1) + query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) + key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + if position_emb is not None: + cos, sin = position_emb + if enable_fast_rope and query_states.dtype != torch.float32: + query_states, key_states = _apply_rotary_pos_emb_fast_dtype(query_states, key_states, cos, sin) + else: + query_states, key_states = _apply_rotary(query_states, key_states, cos, sin) + if not skip_qkv_contig: + query_states = query_states.contiguous() + key_states = key_states.contiguous() + value_states = value_states.contiguous() + attn_output = torch.nn.functional.scaled_dot_product_attention( + query_states, + key_states, + value_states, + is_causal=is_causal, + enable_gqa=True, + ) + attn_output = attn_output.transpose(1, 2).contiguous() + attn_output = attn_output.reshape(bsz, q_len, self.num_heads * self.head_dim) + attn_output = self.o_proj(attn_output) + return attn_output, (key_states, value_states) + + attn.forward = MethodType(_forward, attn) + attn._voxcpm2_fused_qkv = True + patched += 1 + return patched + + +def _install_locdit_fused_mlp(estimator: nn.Module) -> int: + """Patch LocDiT MLP gate/up projections into one Linear call.""" + decoder = getattr(estimator, "decoder", None) + layers = getattr(decoder, "layers", None) + if layers is None: + return 0 + patched = 0 + for layer in layers: + mlp = getattr(layer, "mlp", None) + if mlp is None or getattr(mlp, "_voxcpm2_fused_mlp", False): + continue + gate_proj = getattr(mlp, "gate_proj", None) + up_proj = getattr(mlp, "up_proj", None) + down_proj = getattr(mlp, "down_proj", None) + if gate_proj is None or up_proj is None or down_proj is None: + continue + gate_bias = getattr(gate_proj, "bias", None) + up_bias = getattr(up_proj, "bias", None) + if (gate_bias is None) != (up_bias is None): + continue + weight = torch.cat([gate_proj.weight, up_proj.weight], dim=0).detach() + mlp.register_buffer("_voxcpm2_fused_gate_up_weight", weight, persistent=False) + if gate_bias is not None and up_bias is not None: + bias = torch.cat([gate_bias, up_bias], dim=0).detach() + mlp.register_buffer("_voxcpm2_fused_gate_up_bias", bias, persistent=False) + else: + mlp._voxcpm2_fused_gate_up_bias = None + + def _forward(self, x: torch.Tensor): + gate_up = nn.functional.linear( + x, + self._voxcpm2_fused_gate_up_weight, + self._voxcpm2_fused_gate_up_bias, + ) + gate, up = gate_up.chunk(2, dim=-1) + return self.down_proj(nn.functional.silu(gate) * up) + + mlp.forward = MethodType(_forward, mlp) + mlp._voxcpm2_fused_mlp = True + patched += 1 + return patched + + +def _install_locdit_zero_dt_cache(estimator: nn.Module) -> bool: + """Cache the constant delta-time embedding used when CFM mean_mode is off.""" + if getattr(estimator, "_voxcpm2_zero_dt_cache", False): + return False + time_embeddings = getattr(estimator, "time_embeddings", None) + delta_time_mlp = getattr(estimator, "delta_time_mlp", None) + out_proj = getattr(estimator, "out_proj", None) + decoder = getattr(estimator, "decoder", None) + if time_embeddings is None or delta_time_mlp is None or out_proj is None or decoder is None: + return False + + device = out_proj.weight.device + dtype = out_proj.weight.dtype + with torch.no_grad(): + zero_dt = torch.zeros(1, device=device, dtype=dtype) + zero_dt_emb = delta_time_mlp(time_embeddings(zero_dt).to(dtype)).detach() + estimator.register_buffer("_voxcpm2_zero_dt_emb", zero_dt_emb, persistent=False) + + def _forward( + self, + x: torch.Tensor, + mu: torch.Tensor, + t: torch.Tensor, + cond: torch.Tensor, + dt: torch.Tensor, + ): + x = self.in_proj(x.transpose(1, 2).contiguous()) + cond = self.cond_proj(cond.transpose(1, 2).contiguous()) + prefix = cond.size(1) + + t = self.time_embeddings(t).to(x.dtype) + t = self.time_mlp(t) + t = t + self._voxcpm2_zero_dt_emb.expand(t.shape[0], -1) + + mu = mu.view(x.size(0), -1, x.size(-1)) + x = torch.cat([mu, t.unsqueeze(1), cond, x], dim=1) + + hidden, _ = self.decoder(x, is_causal=False) + hidden = hidden[:, prefix + mu.size(1) + 1 :, :] + hidden = self.out_proj(hidden) + + return hidden.transpose(1, 2).contiguous() + + estimator.forward = MethodType(_forward, estimator) + estimator._voxcpm2_zero_dt_cache = True + return True + + +def _install_locdit_layer_nvtx(estimator: nn.Module) -> int: + decoder = getattr(estimator, "decoder", None) + layers = getattr(decoder, "layers", None) + if layers is None: + return 0 + patched = 0 + for idx, layer in enumerate(layers): + if getattr(layer, "_voxcpm2_layer_nvtx", False): + continue + + def _forward( + self, + hidden_states: torch.Tensor, + position_emb: tuple[torch.Tensor, torch.Tensor] | None, + is_causal: bool, + _idx: int = idx, + ): + residual = hidden_states + with _NvtxRange(f"voxcpm2.cfm.estimator.layer{_idx}.norm1"): + hidden_states = self.input_layernorm(hidden_states) + with _NvtxRange(f"voxcpm2.cfm.estimator.layer{_idx}.attn"): + hidden_states, present_key_value = self.self_attn( + hidden_states=hidden_states, + position_emb=position_emb, + is_causal=is_causal, + ) + with _NvtxRange(f"voxcpm2.cfm.estimator.layer{_idx}.resid1"): + if self.use_mup: + hidden_states = residual + hidden_states * (self.scale_depth / math.sqrt(self.num_hidden_layers)) + else: + hidden_states = residual + hidden_states + + residual = hidden_states + with _NvtxRange(f"voxcpm2.cfm.estimator.layer{_idx}.norm2"): + hidden_states = self.post_attention_layernorm(hidden_states) + with _NvtxRange(f"voxcpm2.cfm.estimator.layer{_idx}.mlp"): + hidden_states = self.mlp(hidden_states) + with _NvtxRange(f"voxcpm2.cfm.estimator.layer{_idx}.resid2"): + if self.use_mup: + hidden_states = residual + hidden_states * (self.scale_depth / math.sqrt(self.num_hidden_layers)) + else: + hidden_states = residual + hidden_states + return hidden_states, present_key_value + + layer.forward = MethodType(_forward, layer) + layer._voxcpm2_layer_nvtx = True + patched += 1 + return patched + + # =================================================================== # CFM pre-allocated buffers + optimized Euler solver # =================================================================== @@ -306,6 +681,8 @@ class _CFMBufferManager: self.dt_in = torch.zeros(n, device=device, dtype=dtype) self.cond_in = torch.zeros(n, feat_dim, patch_size, device=device, dtype=dtype) self.noise = torch.zeros(max_batch_size, feat_dim, patch_size, device=device, dtype=dtype) + self.x_work = torch.zeros_like(self.noise) + self._zeroed_mu_negative_batch: int | None = None self._sway_coef = sway_sampling_coef self._device = device self._dtype = dtype @@ -330,80 +707,123 @@ def _optimized_solve_euler( use_cfg_zero_star: bool = True, cfg_cutoff_ratio: float = 1.0, perf: _PerfTimer | None = None, +) -> torch.Tensor: + b = mu.size(0) + buffers.noise[:b].normal_() + return _optimized_solve_euler_with_noise( + cfm_module, + mu, + patch_size, + cond, + buffers.noise[:b], + n_timesteps, + cfg_value, + buffers, + use_cfg_zero_star=use_cfg_zero_star, + cfg_cutoff_ratio=cfg_cutoff_ratio, + perf=perf, + ) + + +def _optimized_solve_euler_with_noise( + cfm_module: nn.Module, + mu: torch.Tensor, + patch_size: int, + cond: torch.Tensor, + noise: torch.Tensor, + n_timesteps: int, + cfg_value: float, + buffers: _CFMBufferManager, + use_cfg_zero_star: bool = True, + cfg_cutoff_ratio: float = 1.0, + perf: _PerfTimer | None = None, ) -> torch.Tensor: estimator = cfm_module.estimator mean_mode = getattr(cfm_module, "mean_mode", False) b = mu.size(0) - buffers.noise[:b].normal_() - x = buffers.noise[:b].clone() + with _NvtxRange("voxcpm2.cfm.init_copy"): + buffers.x_work[:b].copy_(noise[:b]) + x = buffers.x_work[:b] t_span = buffers.get_t_span(n_timesteps) t, dt = t_span[0], t_span[0] - t_span[1] zero_init_steps = max(1, int(len(t_span) * 0.04)) cfg_cutoff_step = max(zero_init_steps + 1, int(len(t_span) * cfg_cutoff_ratio)) + if use_cfg_zero_star and zero_init_steps > 0: + t = t_span[zero_init_steps] + if zero_init_steps < len(t_span) - 1: + dt = t - t_span[zero_init_steps + 1] + start_step = zero_init_steps + 1 + else: + start_step = 1 - for step in range(1, len(t_span)): - if use_cfg_zero_star and step <= zero_init_steps: - dphi_dt = torch.zeros_like(x) - elif step <= cfg_cutoff_step: - buffers.x_in[:b].copy_(x) - buffers.x_in[b : 2 * b].copy_(x) - buffers.mu_in[:b].copy_(mu) - buffers.mu_in[b : 2 * b].zero_() - # Broadcast the 0-dim GPU scalar directly instead of - # ``.fill_(t.item())`` — ``.item()`` forces a GPU->CPU sync. - buffers.t_in[: 2 * b].copy_(t) - if mean_mode: - buffers.dt_in[: 2 * b].copy_(dt) - else: - buffers.dt_in.zero_() - buffers.cond_in[:b].copy_(cond[:b]) + with _NvtxRange("voxcpm2.cfm.static_inputs"): + buffers.mu_in[:b].copy_(mu) + buffers.cond_in[:b].copy_(cond[:b]) + if cfg_cutoff_step >= start_step: + if buffers._zeroed_mu_negative_batch != b: + buffers.mu_in[b : 2 * b].zero_() + buffers._zeroed_mu_negative_batch = b buffers.cond_in[b : 2 * b].copy_(cond[:b]) + for step in range(start_step, len(t_span)): + if step <= cfg_cutoff_step: + with _NvtxRange("voxcpm2.cfm.step_cfg_inputs"): + buffers.x_in[:b].copy_(x) + buffers.x_in[b : 2 * b].copy_(x) + buffers.t_in[: 2 * b].copy_(t) + if mean_mode: + buffers.dt_in[: 2 * b].copy_(dt) + if perf: perf.start(" cfm.estimator_cfg") - raw_out = estimator( - buffers.x_in[: 2 * b], - buffers.mu_in[: 2 * b], - buffers.t_in[: 2 * b], - buffers.cond_in[: 2 * b], - buffers.dt_in[: 2 * b], - ) + with _NvtxRange("voxcpm2.cfm.estimator_cfg"): + raw_out = estimator( + buffers.x_in[: 2 * b], + buffers.mu_in[: 2 * b], + buffers.t_in[: 2 * b], + buffers.cond_in[: 2 * b], + buffers.dt_in[: 2 * b], + ) if perf: perf.stop(" cfm.estimator_cfg") - dphi_dt, cfg_dphi_dt = raw_out[:b], raw_out[b : 2 * b] - if use_cfg_zero_star: - pos = dphi_dt.reshape(b, -1) - neg = cfg_dphi_dt.reshape(b, -1) - st = torch.sum(pos * neg, 1, keepdim=True) / (torch.sum(neg**2, 1, keepdim=True) + 1e-8) - st = st.view(b, *([1] * (len(dphi_dt.shape) - 1))) - else: - st = 1.0 - dphi_dt = cfg_dphi_dt * st + cfg_value * (dphi_dt - cfg_dphi_dt * st) + with _NvtxRange("voxcpm2.cfm.cfg_post"): + dphi_dt, cfg_dphi_dt = raw_out[:b], raw_out[b : 2 * b] + if use_cfg_zero_star: + pos = dphi_dt.reshape(b, -1) + neg = cfg_dphi_dt.reshape(b, -1) + st = torch.sum(pos * neg, 1, keepdim=True) / (torch.sum(neg**2, 1, keepdim=True) + 1e-8) + st = st.view(b, *([1] * (len(dphi_dt.shape) - 1))) + dphi_dt = cfg_dphi_dt * st + cfg_value * (dphi_dt - cfg_dphi_dt * st) + else: + st = 1.0 + dphi_dt = cfg_dphi_dt * st + cfg_value * (dphi_dt - cfg_dphi_dt * st) else: - buffers.x_in[:b].copy_(x) - buffers.mu_in[:b].copy_(mu) - # Broadcast the 0-dim GPU scalar; ``.fill_(t.item())`` would sync. - buffers.t_in[:b].copy_(t) - if mean_mode: - buffers.dt_in[:b].copy_(dt) - else: - buffers.dt_in[:b].zero_() - buffers.cond_in[:b].copy_(cond[:b]) + with _NvtxRange("voxcpm2.cfm.step_nocfg_inputs"): + buffers.x_in[:b].copy_(x) + buffers.t_in[:b].copy_(t) + if mean_mode: + buffers.dt_in[:b].copy_(dt) if perf: perf.start(" cfm.estimator_nocfg") - dphi_dt = estimator( - buffers.x_in[:b], buffers.mu_in[:b], buffers.t_in[:b], buffers.cond_in[:b], buffers.dt_in[:b] - ) + with _NvtxRange("voxcpm2.cfm.estimator_nocfg"): + dphi_dt = estimator( + buffers.x_in[:b], + buffers.mu_in[:b], + buffers.t_in[:b], + buffers.cond_in[:b], + buffers.dt_in[:b], + ) if perf: perf.stop(" cfm.estimator_nocfg") - x = x - dt * dphi_dt - t = t - dt - if step < len(t_span) - 1: - dt = t - t_span[step + 1] + with _NvtxRange("voxcpm2.cfm.euler_update"): + x.sub_(dt * dphi_dt) + t = t - dt + if step < len(t_span) - 1: + dt = t - t_span[step + 1] return x @@ -417,6 +837,9 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): super().__init__() self.vllm_config = vllm_config self.config = vllm_config.model_config.hf_config + self._runtime_config = _VoxCPM2RuntimeConfig.from_vllm_config(vllm_config) + global _ENABLE_NVTX_PROFILE + _ENABLE_NVTX_PROFILE = self._runtime_config.enable_nvtx_profile self.have_multimodal_outputs = True self.has_preprocess = True @@ -460,7 +883,7 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): self._inference_timesteps = 10 self._cfg_value = 2.0 - self._cfg_cutoff_ratio = 1.0 + self._cfg_cutoff_ratio = self._runtime_config.cfg_cutoff_ratio # Number of trailing latent frames to keep as VAE receptive-field context # for sliding-window streaming decode. 12 matches the nanovllm reference # implementation and covers the longest VAE decoder receptive field. @@ -474,11 +897,31 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): self._speaker_cache = get_speaker_cache() self._load_custom_voice_profiles() - self._perf = _PerfTimer(enabled=_ENABLE_PROFILING) + self._enable_profiling = self._runtime_config.enable_profiling + self._perf = _PerfTimer(enabled=self._enable_profiling) self._cfm_buffers: _CFMBufferManager | None = None self._enable_cuda_graph = True self._scaffold_graphs: dict[int, _CapturedGraph] = {} self._residual_graphs: dict[int, _CapturedGraph] = {} + self._decode_graph_capture_policy = self._runtime_config.decode_graph_capture_policy + self._vae_graphs: dict[tuple[int, int], _CapturedVAEGraph] = {} + self._enable_vae_cuda_graph = self._runtime_config.enable_vae_cuda_graph + self._cfm_graphs: dict[int, _CapturedCFMGraph] = {} + self._enable_cfm_cuda_graph = self._runtime_config.enable_cfm_cuda_graph + self._enable_cfm_prealloc_output = self._runtime_config.enable_cfm_prealloc_output + self._enable_batched_cfm = self._runtime_config.enable_batched_cfm + self._deterministic_cfm_noise = self._runtime_config.deterministic_cfm_noise + self._deterministic_cfm_seed = self._runtime_config.deterministic_cfm_seed + self._vae_decode_sr_cond: torch.Tensor | None = None + self._audio_emit_every = self._runtime_config.audio_emit_every + self._vae_decode_every = self._runtime_config.vae_decode_every + self._enable_delayed_audio_copy = self._runtime_config.enable_delayed_audio_copy + self._delayed_audio_copy_use_events = self._runtime_config.delayed_audio_copy_use_events + self._coalesce_audio_d2h = self._runtime_config.coalesce_audio_d2h + self._enable_batched_vae_decode = self._runtime_config.enable_batched_vae_decode + self._enable_batched_fsq_fusion = self._runtime_config.enable_batched_fsq_fusion + self._batched_fsq_fusion_max_batch = self._runtime_config.batched_fsq_fusion_max_batch + self._audio_copy_stream: torch.cuda.Stream | None = None self._max_cached_graphs = self._max_batch_size self._cuda_graph_warmup_steps = 0 self._cuda_graph_warmup_threshold = 3 @@ -490,6 +933,7 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): self._pending_requests: list[tuple[str, bool, torch.Tensor | None, int]] = [] self._results_queue: list[tuple[str, torch.Tensor | None]] = [] self._audio_queue: list[tuple[str, Any]] = [] + self._last_audio_output_req_ids: list[str] = [] self._deferred_cleanup_ids: set[str] = set() self._active_state_warn_threshold = max(_ACTIVE_STATE_LEAK_WARN_MIN, 4 * self._max_batch_size) # one-shot by design: fires at most once per process to avoid log spam. @@ -568,7 +1012,12 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): return self._get_or_create_state(request_id) def _cleanup_request(self, request_id: str) -> None: - self._active_states.pop(request_id, None) + state = self._active_states.pop(request_id, None) + if state is not None: + state.pending_audio_chunks_gpu.clear() + state.pending_audio_copies.clear() + state.pending_vae_latents_gpu.clear() + state.cfm_output_gpu = None if self._current_request_id == request_id: self._current_request_id = None @@ -664,13 +1113,87 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): return targets: list[str] = [] + cfg = self._runtime_config - try: - tts.feat_decoder.estimator = torch.compile(estimator, mode="reduce-overhead", fullgraph=False) - tts.feat_decoder.estimator._compiled = True - targets.append("LocDiT") - except Exception as e: - logger.warning("torch.compile LocDiT failed: %s", e) + if cfg.enable_loc_dit_fused_qkv: + patched = _install_locdit_fused_qkv( + estimator, + enable_fast_rope=cfg.enable_loc_dit_fast_rope, + skip_qkv_contig=cfg.enable_loc_dit_skip_qkv_contig, + ) + if patched: + targets.append(f"LocDiT fused-qkv attention ({patched} layers)") + if cfg.enable_loc_dit_fused_mlp: + patched = _install_locdit_fused_mlp(estimator) + if patched: + targets.append(f"LocDiT fused gate-up MLP ({patched} layers)") + if cfg.enable_loc_dit_zero_dt_cache and not getattr(tts.feat_decoder, "mean_mode", False): + if _install_locdit_zero_dt_cache(estimator): + targets.append("LocDiT zero-dt embedding cache") + + if cfg.enable_loc_dit_layer_nvtx: + patched = _install_locdit_layer_nvtx(estimator) + if patched: + targets.append(f"LocDiT layer NVTX ({patched} layers, compile skipped)") + estimator._compiled = True + elif self._enable_cfm_cuda_graph: + try: + if cfg.enable_loc_dit_reduce_overhead_no_cg: + import torch._inductor.config as inductor_config + + old_cudagraphs = inductor_config.triton.cudagraphs + old_cudagraph_trees = inductor_config.triton.cudagraph_trees + inductor_config.triton.cudagraphs = False + inductor_config.triton.cudagraph_trees = False + try: + tts.feat_decoder.estimator = torch.compile( + estimator, + mode="reduce-overhead", + fullgraph=False, + ) + finally: + inductor_config.triton.cudagraphs = old_cudagraphs + inductor_config.triton.cudagraph_trees = old_cudagraph_trees + targets.append("LocDiT (reduce-overhead no-inductor-cudagraph + CUDA Graph)") + else: + tts.feat_decoder.estimator = torch.compile( + estimator, + fullgraph=cfg.enable_loc_dit_fullgraph_no_cg, + options={ + "triton.cudagraphs": False, + "triton.cudagraph_trees": False, + }, + ) + graph_mode = "fullgraph" if cfg.enable_loc_dit_fullgraph_no_cg else "partitioned" + targets.append(f"LocDiT ({graph_mode} compile no-inductor-cudagraph + CUDA Graph)") + tts.feat_decoder.estimator._compiled = True + except Exception as e: + logger.warning("torch.compile LocDiT reduce-overhead for CUDA Graph failed: %s", e) + try: + tts.feat_decoder.estimator = torch.compile( + estimator, + fullgraph=cfg.enable_loc_dit_fullgraph_no_cg, + options={ + "triton.cudagraphs": False, + "triton.cudagraph_trees": False, + }, + ) + tts.feat_decoder.estimator._compiled = True + graph_mode = "fullgraph" if cfg.enable_loc_dit_fullgraph_no_cg else "partitioned" + targets.append(f"LocDiT ({graph_mode} compile no-inductor-cudagraph + CUDA Graph)") + except Exception as inner_e: + logger.warning("torch.compile LocDiT for CUDA Graph failed: %s", inner_e) + else: + try: + tts.feat_decoder.estimator = torch.compile( + estimator, + mode="reduce-overhead", + fullgraph=False, + ) + tts.feat_decoder.estimator._compiled = True + targets.append("LocDiT") + except Exception as e: + logger.warning("torch.compile LocDiT failed: %s", e) try: if not hasattr(tts.feat_encoder, "_compiled"): @@ -680,7 +1203,7 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): except Exception as e: logger.warning("torch.compile feat_encoder failed: %s", e) - if self._compile_vae: + if self._compile_vae and not self._enable_vae_cuda_graph: try: if not hasattr(tts.audio_vae, "_compiled"): tts.audio_vae.decode = torch.compile(tts.audio_vae.decode, mode="reduce-overhead", fullgraph=False) @@ -795,6 +1318,8 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): inputs_embeds: torch.Tensor, positions: torch.Tensor, batch_size: int, + *, + clone_output: bool = True, ) -> torch.Tensor: """Copy fresh inputs into static buffers, then replay. @@ -805,7 +1330,208 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): g.input_embeds[:batch_size].copy_(inputs_embeds[:batch_size]) g.positions[:batch_size].copy_(positions[:batch_size]) g.graph.replay() - return g.output[:batch_size].clone() + output = g.output[:batch_size] + return output.clone() if clone_output else output + + def _should_use_decode_graph(self, batch_size: int) -> bool: + if batch_size > self._max_cached_graphs: + return False + policy = self._decode_graph_capture_policy + if policy == "all": + return True + if policy == "power2": + return batch_size > 0 and (batch_size & (batch_size - 1)) == 0 + if policy == "power2_or_cached": + return ( + batch_size in self._scaffold_graphs + or batch_size in self._residual_graphs + or (batch_size > 0 and (batch_size & (batch_size - 1)) == 0) + ) + return True + + def _capture_vae_graph(self, feat: torch.Tensor) -> _CapturedVAEGraph: + batch_size = feat.shape[0] + num_frames = feat.shape[-1] + input_feat = torch.zeros_like(feat) + sr_cond = self._get_vae_decode_sr_cond(feat.device) + + with torch.no_grad(): + for _ in range(3): + output = self.tts.audio_vae.decode(input_feat, sr_cond=sr_cond) + + g = _CapturedVAEGraph( + graph=torch.cuda.CUDAGraph(), + input_feat=input_feat, + output=output, + ) + with torch.cuda.graph(g.graph, pool=current_platform.get_global_graph_pool()): + g.output = self.tts.audio_vae.decode(g.input_feat, sr_cond=sr_cond) + + logger.info("CUDA Graph captured for AudioVAE decode (batch_size=%d, frames=%d)", batch_size, num_frames) + return g + + def _get_vae_decode_sr_cond(self, device: torch.device) -> torch.Tensor: + sr_cond = self._vae_decode_sr_cond + if sr_cond is None or sr_cond.device != device: + sr_cond = torch.tensor( + [self.tts.audio_vae.out_sample_rate], + device=device, + dtype=torch.int32, + ) + self._vae_decode_sr_cond = sr_cond + return sr_cond + + def _run_vae_decode(self, feat: torch.Tensor) -> torch.Tensor: + if not feat.is_cuda: + return self.tts.audio_vae.decode(feat) + + sr_cond = self._get_vae_decode_sr_cond(feat.device) + if not self._enable_vae_cuda_graph: + return self.tts.audio_vae.decode(feat, sr_cond=sr_cond) + + graph_key = (feat.shape[0], feat.shape[-1]) + graph = self._vae_graphs.get(graph_key) + if graph is None: + graph = self._capture_vae_graph(feat) + self._vae_graphs[graph_key] = graph + + graph.input_feat.copy_(feat) + graph.graph.replay() + return graph.output + + def _capture_cfm_graph( + self, + mu: torch.Tensor | None, + cond: torch.Tensor, + *, + batch_size: int | None = None, + dit_hidden: int | None = None, + ) -> _CapturedCFMGraph: + if mu is not None: + batch_size = mu.shape[0] + dit_hidden = mu.shape[-1] + device = mu.device + dtype = mu.dtype + else: + assert batch_size is not None + assert dit_hidden is not None + device = cond.device + dtype = cond.dtype + graph_buffers = _CFMBufferManager( + device=device, + dtype=dtype, + feat_dim=self._feat_dim, + patch_size=self._patch_size, + dit_hidden_size=dit_hidden, + max_batch_size=batch_size, + ) + static_mu = torch.zeros(batch_size, dit_hidden, device=device, dtype=dtype) + static_cond = torch.zeros_like(cond) + static_noise = torch.zeros(batch_size, self._feat_dim, self._patch_size, device=device, dtype=dtype) + + with torch.no_grad(): + for _ in range(3): + output = _optimized_solve_euler_with_noise( + self.tts.feat_decoder, + static_mu, + self._patch_size, + static_cond, + static_noise, + self._inference_timesteps, + self._cfg_value, + graph_buffers, + cfg_cutoff_ratio=self._cfg_cutoff_ratio, + ) + + graph = _CapturedCFMGraph( + graph=torch.cuda.CUDAGraph(), + mu=static_mu, + cond=static_cond, + noise=static_noise, + output=output, + buffers=graph_buffers, + ) + with torch.cuda.graph(graph.graph, pool=current_platform.get_global_graph_pool()): + graph.output = _optimized_solve_euler_with_noise( + self.tts.feat_decoder, + graph.mu, + self._patch_size, + graph.cond, + graph.noise, + self._inference_timesteps, + self._cfg_value, + graph.buffers, + cfg_cutoff_ratio=self._cfg_cutoff_ratio, + ) + + logger.info("CUDA Graph captured for VoxCPM2 CFM solver (batch_size=%d)", batch_size) + return graph + + def _get_cfm_cuda_graph( + self, + mu: torch.Tensor | None, + cond: torch.Tensor, + *, + batch_size: int | None = None, + dit_hidden: int | None = None, + ) -> _CapturedCFMGraph: + if mu is not None: + graph_batch_size = mu.shape[0] + else: + assert batch_size is not None + graph_batch_size = batch_size + graph = self._cfm_graphs.get(graph_batch_size) + if graph is None: + graph = self._capture_cfm_graph(mu, cond, batch_size=batch_size, dit_hidden=dit_hidden) + self._cfm_graphs[graph_batch_size] = graph + return graph + + def _run_cfm_cuda_graph(self, mu: torch.Tensor, cond: torch.Tensor) -> torch.Tensor: + graph = self._get_cfm_cuda_graph(mu, cond) + with _NvtxRange("voxcpm2.cfm.graph_copy_mu"): + graph.mu.copy_(mu) + return self._run_cfm_cuda_graph_from_static_mu(graph, cond) + + def _run_cfm_cuda_graph_from_static_mu(self, graph: _CapturedCFMGraph, cond: torch.Tensor) -> torch.Tensor: + with _NvtxRange("voxcpm2.cfm.graph_copy_cond"): + graph.cond.copy_(cond) + with _NvtxRange("voxcpm2.cfm.graph_noise"): + graph.noise.normal_() + with _NvtxRange("voxcpm2.cfm.graph_replay"): + graph.graph.replay() + with _NvtxRange("voxcpm2.cfm.graph_output_clone"): + return graph.output.clone() + + def _run_cfm_cuda_graph_to_state_buffer( + self, + state: _RequestState, + mu: torch.Tensor, + cond: torch.Tensor, + ) -> torch.Tensor: + graph = self._get_cfm_cuda_graph(mu, cond) + with _NvtxRange("voxcpm2.cfm.graph_copy_mu"): + graph.mu.copy_(mu) + with _NvtxRange("voxcpm2.cfm.graph_copy_cond"): + graph.cond.copy_(cond) + with _NvtxRange("voxcpm2.cfm.graph_noise"): + if self._deterministic_cfm_noise: + self._fill_deterministic_cfm_noise(state, graph.noise) + else: + graph.noise.normal_() + with _NvtxRange("voxcpm2.cfm.graph_replay"): + graph.graph.replay() + + output_shape = graph.output.shape + if ( + state.cfm_output_gpu is None + or state.cfm_output_gpu.shape != output_shape + or state.cfm_output_gpu.device != graph.output.device + or state.cfm_output_gpu.dtype != graph.output.dtype + ): + state.cfm_output_gpu = torch.empty(output_shape, device=graph.output.device, dtype=graph.output.dtype) + with _NvtxRange("voxcpm2.cfm.graph_output_copy"): + state.cfm_output_gpu.copy_(graph.output) + return state.cfm_output_gpu.transpose(1, 2) # -------------------- vllm hooks -------------------- @@ -823,6 +1549,7 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): self._perf.start("forward_total") dev = input_ids.device + self._last_audio_output_req_ids = [req_id for req_id, _, _, _ in self._pending_requests] num_reqs = len(self._pending_requests) num_decode = sum(1 for _, is_p, _, n in self._pending_requests if not is_p and n == 1) is_all_decode = num_decode == num_reqs and num_reqs > 0 @@ -837,6 +1564,11 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): ) if can_use_graph and is_all_decode and num_reqs <= self._max_cached_graphs: + use_scaffold_graph = self._should_use_decode_graph(num_reqs) + else: + use_scaffold_graph = False + + if use_scaffold_graph: self._perf.start("scaffold_fwd") if num_reqs not in self._scaffold_graphs: self._scaffold_graphs[num_reqs] = self._capture_graph(self.model, num_reqs, "scaffold") @@ -853,11 +1585,30 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): if isinstance(scaffold_hidden, tuple): scaffold_hidden = scaffold_hidden[0] - # Phase 1: per-request FSQ + residual input token_offset = 0 residual_inputs: list[torch.Tensor] = [] residual_positions: list[torch.Tensor] = [] - req_metas: list[tuple] = [] + req_metas: list[tuple[_RequestState, bool, dict]] = [] + pending_decode_fsq: list[tuple[_RequestState, torch.Tensor, torch.Tensor]] = [] + + def flush_decode_fsq_batch() -> None: + if not pending_decode_fsq: + return + if len(pending_decode_fsq) == 1: + state, req_hidden, req_pos = pending_decode_fsq[0] + res_input, meta = self._prepare_residual_decode(state, req_hidden, dev) + residual_inputs.append(res_input) + residual_positions.append(req_pos) + req_metas.append((state, False, meta)) + else: + states = [state for state, _, _ in pending_decode_fsq] + hidden_list = [req_hidden for _, req_hidden, _ in pending_decode_fsq] + batched = self._prepare_residual_decode_batch(states, hidden_list, dev) + for (state, _, req_pos), (res_input, meta) in zip(pending_decode_fsq, batched): + residual_inputs.append(res_input) + residual_positions.append(req_pos) + req_metas.append((state, False, meta)) + pending_decode_fsq.clear() for req_id, is_prefill, _req_embeds, n in self._pending_requests: state = self._switch_to_request(req_id) @@ -865,10 +1616,25 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): req_pos = positions[token_offset : token_offset + n] if is_prefill: + flush_decode_fsq_batch() res_input, meta = self._prepare_residual_prefill(state, req_hidden, dev) elif state.prefill_completed: + if ( + self._enable_batched_fsq_fusion + and n == 1 + and state.prev_feat_embed is not None + and req_hidden.ndim == 2 + and req_hidden.shape[0] == 1 + ): + pending_decode_fsq.append((state, req_hidden, req_pos)) + token_offset += n + if len(pending_decode_fsq) >= self._batched_fsq_fusion_max_batch: + flush_decode_fsq_batch() + continue + flush_decode_fsq_batch() res_input, meta = self._prepare_residual_decode(state, req_hidden, dev) else: + flush_decode_fsq_batch() token_offset += n self._results_queue.append((req_id, None)) self._audio_queue.append((req_id, None)) @@ -878,6 +1644,7 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): residual_positions.append(req_pos) req_metas.append((state, is_prefill, meta)) token_offset += n + flush_decode_fsq_batch() # Phase 2: batch residual_model forward if residual_inputs: @@ -890,36 +1657,70 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): and is_all_decode and graph_ready and residual_batch_size == num_reqs # 1 token per request - and residual_batch_size <= self._max_cached_graphs + and self._should_use_decode_graph(residual_batch_size) ) self._perf.start("residual_fwd") - if use_residual_graph: - if residual_batch_size not in self._residual_graphs: - self._residual_graphs[residual_batch_size] = self._capture_graph( - self.residual_model, residual_batch_size, "residual", is_residual=True + with _NvtxRange("voxcpm2.residual_fwd"): + if use_residual_graph: + if residual_batch_size not in self._residual_graphs: + self._residual_graphs[residual_batch_size] = self._capture_graph( + self.residual_model, residual_batch_size, "residual", is_residual=True + ) + batch_out = self._replay_graph( + self._residual_graphs[residual_batch_size], + batch_in, + batch_pos, + residual_batch_size, ) - batch_out = self._replay_graph( - self._residual_graphs[residual_batch_size], batch_in, batch_pos, residual_batch_size - ) - else: - batch_out = self.residual_model(batch_pos, batch_in) + else: + batch_out = self.residual_model(batch_pos, batch_in) self._perf.stop("residual_fwd") - # Phase 3: per-request LocDiT + update - offset = 0 - for idx, (state, is_prefill, meta) in enumerate(req_metas): - n = residual_inputs[idx].shape[0] - res_out = batch_out[offset : offset + n] - offset += n + can_batch_decode_tail = ( + bool(req_metas) + and all(not is_prefill for _, is_prefill, _ in req_metas) + and all(x.shape[0] == 1 for x in residual_inputs) + ) + if can_batch_decode_tail: + self._finish_decode_batch(req_metas, batch_out) + self._precompute_stop_flags_for_audio_collect([state for state, _, _ in req_metas]) + ready_audio_by_req = self._drain_ready_audio_copies_for_states([state for state, _, _ in req_metas]) + audio_by_req = self._collect_audio_batch( + [state for state, _, _ in req_metas], + initial_delayed_chunks_by_req=ready_audio_by_req, + ) + for state, _, _ in req_metas: + self._results_queue.append((state.request_id, state.precomputed_stop_logits)) + self._audio_queue.append((state.request_id, audio_by_req.get(state.request_id))) + else: + offset = 0 + decoded_states: list[_RequestState] = [] + for idx, (state, is_prefill, meta) in enumerate(req_metas): + n = residual_inputs[idx].shape[0] + res_out = batch_out[offset : offset + n] + offset += n - if is_prefill: - self._finish_prefill(state, meta, res_out, dev) - else: - self._finish_decode(state, meta, res_out, dev) + if is_prefill: + self._finish_prefill(state, meta, res_out, dev) + else: + self._finish_decode(state, meta, res_out, dev) + decoded_states.append(state) - self._results_queue.append((state.request_id, state.precomputed_stop_logits)) - self._audio_queue.append((state.request_id, self._collect_audio(state))) + collect_states = [state for state, _, _ in req_metas] + self._precompute_stop_flags_for_audio_collect(collect_states) + ready_audio_by_req = self._drain_ready_audio_copies_for_states(collect_states) + audio_by_req = self._collect_audio_batch( + collect_states, + initial_delayed_chunks_by_req={ + state.request_id: ready_audio_by_req.get(state.request_id) + for state, is_prefill, _ in req_metas + if not is_prefill + }, + ) + for state, is_prefill, _ in req_metas: + self._results_queue.append((state.request_id, state.precomputed_stop_logits)) + self._audio_queue.append((state.request_id, audio_by_req.get(state.request_id))) self._pending_requests.clear() self._flush_deferred_cleanup() @@ -973,26 +1774,135 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): res_input = tts.fusion_concat_proj(torch.cat([lm_h, prev], dim=-1)) return res_input, {"new_lm_hidden": lm_h} + def _prepare_residual_decode_batch( + self, + states: list[_RequestState], + base_lm_outs: list[torch.Tensor], + dev: Any, + ) -> list[tuple[torch.Tensor, dict]]: + tts = self.tts + + for state in states: + state.decode_step_count += 1 + if state.decode_step_count >= self._max_decode_steps: + logger.warning( + "MAX_DECODE_STEPS for %s (%d), forcing stop", + state.request_id, + state.decode_step_count, + ) + state.is_stopping = True + + with _NvtxRange("voxcpm2.fsq_fusion_batch"): + hidden_batch = torch.cat( + [h.reshape(1, -1) if h.ndim == 1 else h for h in base_lm_outs], + dim=0, + ) + lm_h_batch = tts.fsq_layer(hidden_batch) + if lm_h_batch.ndim == 1: + lm_h_batch = lm_h_batch.unsqueeze(0) + prev_batch = torch.cat( + [state.prev_feat_embed.to(dev, dtype=self._side_dtype).reshape(1, -1) for state in states], + dim=0, + ) + res_batch = tts.fusion_concat_proj(torch.cat([lm_h_batch, prev_batch], dim=-1)) + + return [ + ( + res_batch[i : i + 1], + {"new_lm_hidden": lm_h_batch[i : i + 1]}, + ) + for i in range(len(states)) + ] + def _run_cfm(self, dit_h: torch.Tensor, cond: torch.Tensor) -> torch.Tensor: - if self._cfm_buffers is not None: - return _optimized_solve_euler( - self.tts.feat_decoder, - dit_h, - self._patch_size, - cond, - self._inference_timesteps, - self._cfg_value, - self._cfm_buffers, - cfg_cutoff_ratio=self._cfg_cutoff_ratio, - perf=self._perf, + with _NvtxRange("voxcpm2.cfm"): + if self._cfm_buffers is not None: + if self._enable_cfm_cuda_graph and dit_h.is_cuda: + return self._run_cfm_cuda_graph(dit_h, cond).transpose(1, 2) + return _optimized_solve_euler( + self.tts.feat_decoder, + dit_h, + self._patch_size, + cond, + self._inference_timesteps, + self._cfg_value, + self._cfm_buffers, + cfg_cutoff_ratio=self._cfg_cutoff_ratio, + perf=self._perf, + ).transpose(1, 2) + return self.tts.feat_decoder( + mu=dit_h, + patch_size=self._patch_size, + cond=cond, + n_timesteps=self._inference_timesteps, + cfg_value=self._cfg_value, ).transpose(1, 2) - return self.tts.feat_decoder( - mu=dit_h, - patch_size=self._patch_size, - cond=cond, - n_timesteps=self._inference_timesteps, - cfg_value=self._cfg_value, - ).transpose(1, 2) + + def _run_cfm_for_state(self, state: _RequestState, dit_h: torch.Tensor, cond: torch.Tensor) -> torch.Tensor: + with _NvtxRange("voxcpm2.cfm"): + if self._cfm_buffers is not None: + if self._enable_cfm_cuda_graph and dit_h.is_cuda: + if self._deterministic_cfm_noise and not self._enable_cfm_prealloc_output: + graph = self._get_cfm_cuda_graph(dit_h, cond) + with _NvtxRange("voxcpm2.cfm.graph_copy_mu"): + graph.mu.copy_(dit_h) + with _NvtxRange("voxcpm2.cfm.graph_copy_cond"): + graph.cond.copy_(cond) + with _NvtxRange("voxcpm2.cfm.graph_noise"): + self._fill_deterministic_cfm_noise(state, graph.noise) + with _NvtxRange("voxcpm2.cfm.graph_replay"): + graph.graph.replay() + with _NvtxRange("voxcpm2.cfm.graph_output_clone"): + return graph.output.clone().transpose(1, 2) + if not self._enable_cfm_prealloc_output: + return self._run_cfm_cuda_graph(dit_h, cond).transpose(1, 2) + return self._run_cfm_cuda_graph_to_state_buffer(state, dit_h, cond) + if self._deterministic_cfm_noise: + noise = self._cfm_buffers.noise[: dit_h.shape[0]] + self._fill_deterministic_cfm_noise(state, noise) + return _optimized_solve_euler_with_noise( + self.tts.feat_decoder, + dit_h, + self._patch_size, + cond, + noise, + self._inference_timesteps, + self._cfg_value, + self._cfm_buffers, + cfg_cutoff_ratio=self._cfg_cutoff_ratio, + perf=self._perf, + ).transpose(1, 2) + return _optimized_solve_euler( + self.tts.feat_decoder, + dit_h, + self._patch_size, + cond, + self._inference_timesteps, + self._cfg_value, + self._cfm_buffers, + cfg_cutoff_ratio=self._cfg_cutoff_ratio, + perf=self._perf, + ).transpose(1, 2) + return self.tts.feat_decoder( + mu=dit_h, + patch_size=self._patch_size, + cond=cond, + n_timesteps=self._inference_timesteps, + cfg_value=self._cfg_value, + ).transpose(1, 2) + + def _fill_deterministic_cfm_noise(self, state: _RequestState, out: torch.Tensor) -> None: + """Fill CFM noise deterministically for benchmark replay only.""" + request_key = state.request_id.split("_", 1)[0] + if not request_key.isdigit(): + request_key = state.request_id + key = f"{self._deterministic_cfm_seed}:{request_key}:{state.cfm_noise_step}".encode() + digest = hashlib.blake2b(key, digest_size=8).digest() + seed = int.from_bytes(digest, "little") & 0x7FFF_FFFF_FFFF_FFFF + gen = torch.Generator(device=out.device) + gen.manual_seed(seed) + out.normal_(generator=gen) + state.cfm_noise_step += 1 def _finish_prefill(self, state: _RequestState, meta: dict, res_out: torch.Tensor, dev: Any): tts = self.tts @@ -1000,22 +1910,19 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): prefix_feat_cond = meta["prefix_feat_cond"] residual_hidden = res_out[-1:, :] - state.precomputed_stop_logits = tts.stop_head(tts.stop_actn(tts.stop_proj(lm_hidden))).detach() + stop_logits = tts.stop_head(tts.stop_actn(tts.stop_proj(lm_hidden))).detach() dit_h = torch.cat([tts.lm_to_dit_proj(lm_hidden), tts.res_to_dit_proj(residual_hidden)], dim=-1) self._setup_cfm_buffers() if self._enable_torch_compile: self._setup_torch_compile() - pred_feat = self._run_cfm(dit_h, prefix_feat_cond.transpose(1, 2).contiguous()) + pred_feat = self._run_cfm_for_state(state, dit_h, prefix_feat_cond.transpose(1, 2).contiguous()) with torch.no_grad(): curr_embed = tts.enc_to_lm_proj(tts.feat_encoder(pred_feat.unsqueeze(1))).squeeze(1) - state.curr_embed_for_next = curr_embed.detach() - state.prev_feat_embed = curr_embed.detach() - state.curr_prefix_feat_cond = pred_feat[0].detach() - state.last_audio_patch_gpu = pred_feat.detach() + self._commit_decode_state(state, stop_logits, curr_embed, pred_feat) state.decode_step_count = 0 state.request_start_time = time.perf_counter() state.prefill_completed = True @@ -1035,27 +1942,243 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): dit_proj = getattr(self, "_compiled_dit_proj", None) or self._dit_proj_fn stop_fn = getattr(self, "_compiled_stop_fn", None) or self._stop_fn - dit_h = dit_proj(lm_h, res_h) pfc = state.curr_prefix_feat_cond.to(self._side_dtype) if pfc.ndim == 2: pfc = pfc.unsqueeze(0) + cond = pfc.transpose(1, 2).contiguous() - pred_feat = self._run_cfm(dit_h, pfc.transpose(1, 2).contiguous()) + dit_h = dit_proj(lm_h, res_h) + pred_feat = self._run_cfm_for_state(state, dit_h, cond) next_embed = tts.enc_to_lm_proj(tts.feat_encoder(pred_feat.unsqueeze(1))).squeeze(1) + stop_logits = stop_fn(lm_h).detach() + self._commit_decode_state(state, stop_logits, next_embed, pred_feat) - state.precomputed_stop_logits = stop_fn(lm_h).detach() + self._perf.stop("decode_step") + if self._enable_profiling and state.decode_step_count % 20 == 0: + logger.info("Step %d[%s]:\n%s", state.decode_step_count, state.request_id, self._perf.breakdown()) + + def _finish_decode_batch(self, req_metas: list[tuple], batch_out: torch.Tensor): + self._perf.start("decode_step") + tts = self.tts + + dit_proj = getattr(self, "_compiled_dit_proj", None) or self._dit_proj_fn + stop_fn = getattr(self, "_compiled_stop_fn", None) or self._stop_fn + + if self._enable_batched_cfm and not self._enable_cfm_cuda_graph: + states = [state for state, _, _ in req_metas] + lm_h = torch.cat([meta["new_lm_hidden"] for _, _, meta in req_metas], dim=0) + pfc = torch.cat( + [ + state.curr_prefix_feat_cond.to(self._side_dtype).unsqueeze(0) + if state.curr_prefix_feat_cond.ndim == 2 + else state.curr_prefix_feat_cond.to(self._side_dtype) + for state in states + ], + dim=0, + ) + with _NvtxRange("voxcpm2.dit_proj"): + dit_h = dit_proj(lm_h, batch_out) + cond = pfc.transpose(1, 2).contiguous() + if self._deterministic_cfm_noise and self._cfm_buffers is not None: + noise = self._cfm_buffers.noise[: dit_h.size(0)] + for i, state in enumerate(states): + self._fill_deterministic_cfm_noise(state, noise[i : i + 1]) + pred_feat = _optimized_solve_euler_with_noise( + self.tts.feat_decoder, + dit_h, + self._patch_size, + cond, + noise, + self._inference_timesteps, + self._cfg_value, + self._cfm_buffers, + cfg_cutoff_ratio=self._cfg_cutoff_ratio, + perf=self._perf, + ).transpose(1, 2) + else: + pred_feat = self._run_cfm(dit_h, cond) + with _NvtxRange("voxcpm2.feat_encoder_feedback"): + next_embed = tts.enc_to_lm_proj(tts.feat_encoder(pred_feat.unsqueeze(1))).squeeze(1) + with _NvtxRange("voxcpm2.stop_fn"): + stop_logits = stop_fn(lm_h).detach() + + for i, state in enumerate(states): + self._commit_decode_state( + state, + stop_logits[i : i + 1], + next_embed[i : i + 1], + pred_feat[i : i + 1], + ) + + self._perf.stop("decode_step") + return + + for i, (state, _, meta) in enumerate(req_metas): + lm_h = meta["new_lm_hidden"] + res_h = batch_out[i : i + 1] + pfc = state.curr_prefix_feat_cond.to(self._side_dtype) + if pfc.ndim == 2: + pfc = pfc.unsqueeze(0) + cond = pfc.transpose(1, 2).contiguous() + with _NvtxRange("voxcpm2.dit_proj"): + dit_h = dit_proj(lm_h, res_h) + pred_feat = self._run_cfm_for_state(state, dit_h, cond) + with _NvtxRange("voxcpm2.feat_encoder_feedback"): + next_embed = tts.enc_to_lm_proj(tts.feat_encoder(pred_feat.unsqueeze(1))).squeeze(1) + with _NvtxRange("voxcpm2.stop_fn"): + stop_logits = stop_fn(lm_h).detach() + self._commit_decode_state(state, stop_logits, next_embed, pred_feat) + + self._perf.stop("decode_step") + + def _commit_decode_state( + self, + state: _RequestState, + stop_logits: torch.Tensor, + next_embed: torch.Tensor, + pred_feat: torch.Tensor, + ) -> None: + state.precomputed_stop_logits = stop_logits + state.precomputed_is_stopping = None state.curr_embed_for_next = next_embed.detach() state.prev_feat_embed = next_embed.detach() state.curr_prefix_feat_cond = pred_feat[0].detach() state.last_audio_patch_gpu = pred_feat.detach() - self._perf.stop("decode_step") - if _ENABLE_PROFILING and state.decode_step_count % 20 == 0: - logger.info("Step %d[%s]:\n%s", state.decode_step_count, state.request_id, self._perf.breakdown()) - # -------------------- audio collection -------------------- - def _collect_audio(self, state: _RequestState) -> torch.Tensor | None: + def _uses_sparse_audio_outputs(self) -> bool: + return ( + self._audio_emit_every > 1 or self._enable_delayed_audio_copy or getattr(self, "_vae_decode_every", 1) > 1 + ) + + def _vae_output_storage_may_be_reused(self) -> bool: + return self._enable_vae_cuda_graph or bool(getattr(self.tts.audio_vae, "_compiled", False)) + + def _enqueue_delayed_audio_copy(self, state: _RequestState, audio: torch.Tensor) -> None: + src = audio.detach().contiguous() + if not src.is_cuda: + state.pending_audio_copies.append(_PendingAudioCopy(host=src.cpu().contiguous())) + return + + # Compiled/graph VAE output storage can be reused on the next token. + if self._vae_output_storage_may_be_reused(): + src = src.clone() + host = torch.empty(src.shape, dtype=src.dtype, device="cpu", pin_memory=True) + if self._audio_copy_stream is None: + self._audio_copy_stream = torch.cuda.Stream(device=src.device) + copy_stream = self._audio_copy_stream + copy_stream.wait_stream(torch.cuda.current_stream(src.device)) + with torch.cuda.stream(copy_stream): + host.copy_(src, non_blocking=True) + src.record_stream(copy_stream) + event = None + if self._delayed_audio_copy_use_events: + event = torch.cuda.Event() + event.record(copy_stream) + state.pending_audio_copies.append(_PendingAudioCopy(host=host, event=event, source=src, async_copy=True)) + + def _pending_audio_copy_ready(self, copy: _PendingAudioCopy) -> bool: + if copy.event is not None: + return copy.event.query() + if not copy.async_copy: + return True + return self._audio_copy_stream is not None and self._audio_copy_stream.query() + + def _drain_ready_audio_copies_for_states(self, states: list[_RequestState]) -> dict[str, list[torch.Tensor]]: + if ( + not self._enable_delayed_audio_copy + or self._audio_emit_every != 1 + or self._delayed_audio_copy_use_events + or self._audio_copy_stream is None + or not any(state.pending_audio_copies for state in states) + or not self._audio_copy_stream.query() + ): + return {} + + ready_by_req: dict[str, list[torch.Tensor]] = {} + for state in states: + ready = self._drain_pending_audio_copies(state, force=False) + if ready: + ready_by_req[state.request_id] = ready + return ready_by_req + + def _drain_pending_audio_copies(self, state: _RequestState, *, force: bool) -> list[torch.Tensor]: + ready: list[torch.Tensor] = [] + if ( + force + and not self._delayed_audio_copy_use_events + and state.pending_audio_copies + and any(copy.async_copy for copy in state.pending_audio_copies) + and self._audio_copy_stream is not None + and not self._audio_copy_stream.query() + ): + self._audio_copy_stream.synchronize() + while state.pending_audio_copies: + pending = state.pending_audio_copies[0] + if pending.event is not None and force: + pending.event.synchronize() + elif not self._pending_audio_copy_ready(pending): + break + state.pending_audio_copies.pop(0) + ready.append(pending.host.float()) + return ready + + @staticmethod + def _merge_audio_chunks(chunks: list[torch.Tensor]) -> torch.Tensor | None: + if not chunks: + return None + if len(chunks) == 1: + return chunks[0] + return torch.cat([chunk.reshape(-1) for chunk in chunks], dim=0) + + def _precompute_stop_flags_for_audio_collect(self, states: list[_RequestState]) -> None: + if ( + self._audio_emit_every == 1 + and not self._enable_delayed_audio_copy + and getattr(self, "_vae_decode_every", 1) == 1 + ): + return + + pending: list[tuple[_RequestState, torch.Tensor]] = [] + for state in states: + if state.is_stopping or state.precomputed_is_stopping is not None: + continue + stop_logits = state.precomputed_stop_logits + if stop_logits is None or not stop_logits.is_cuda: + continue + pending.append((state, stop_logits)) + if not pending: + return + + stacked = torch.stack([stop_logits[0] for _, stop_logits in pending], dim=0) + stop_mask = stacked[:, 1] > stacked[:, 0] + stop_mask_cpu = stop_mask.cpu() + for i, (state, _) in enumerate(pending): + is_stopping = bool(stop_mask_cpu[i]) + state.precomputed_is_stopping = is_stopping + if is_stopping: + state.is_stopping = True + + @staticmethod + def _should_stop_from_cached_logits(state: _RequestState) -> bool: + if state.is_stopping: + return True + cached = state.precomputed_is_stopping + if cached is not None: + return cached + stop_logits = state.precomputed_stop_logits + if stop_logits is None: + return False + is_stopping = bool(stop_logits[0, 1] > stop_logits[0, 0]) + state.precomputed_is_stopping = is_stopping + if is_stopping: + state.is_stopping = True + return is_stopping + + def _collect_audio( + self, state: _RequestState, initial_delayed_chunks: list[torch.Tensor] | None = None + ) -> torch.Tensor | None: """Per-step sliding-window VAE decode (nanovllm pattern). Each decode step feeds ``[decode_pad, new_patch]`` through the VAE @@ -1067,13 +2190,27 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): Returns the delta audio chunk (not cumulative) so the output processor can stream each chunk to the client independently. """ + delayed_chunks: list[torch.Tensor] = list(initial_delayed_chunks or []) + if self._enable_delayed_audio_copy and self._audio_emit_every == 1: + if initial_delayed_chunks is None: + delayed_chunks.extend(self._drain_pending_audio_copies(state, force=False)) + patch = state.last_audio_patch_gpu if patch is None: - return None + return self._merge_audio_chunks(delayed_chunks) state.last_audio_patch_gpu = None # patch shape: (patch_size, feat_dim) or (1, patch_size, feat_dim) new_latent = patch.reshape(-1, self._feat_dim).to(torch.float32) + vae_decode_every = getattr(self, "_vae_decode_every", 1) + if vae_decode_every > 1: + is_stopping = self._should_stop_from_cached_logits(state) + state.pending_vae_latents_gpu.append(new_latent.detach()) + if not is_stopping and len(state.pending_vae_latents_gpu) < vae_decode_every: + return self._merge_audio_chunks(delayed_chunks) + new_latent = torch.cat(state.pending_vae_latents_gpu, dim=0) + state.pending_vae_latents_gpu.clear() + n_new = new_latent.shape[0] # = patch_size (typically 4) self._perf.start("vae_decode") @@ -1087,21 +2224,143 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): pad_frames = 0 # VAE decode: (1, feat_dim, T_frames) -> (1, 1, T_samples) - feat = vae_input.unsqueeze(0).transpose(1, 2).contiguous() - with torch.no_grad(): - audio = self.tts.audio_vae.decode(feat.to(self._device)).reshape(-1) + with _NvtxRange("voxcpm2.vae_prepare"): + feat = vae_input.unsqueeze(0).transpose(1, 2).contiguous() + vae_frames = vae_input.shape[0] + with _NvtxRange("voxcpm2.vae_decode"): + with torch.no_grad(): + audio = self._run_vae_decode(feat.to(self._device)).reshape(-1) # Slice out only the new audio (after the pad region). # Each latent frame maps to decoder_chunk_size audio samples. - dcs = int(getattr(self.tts.audio_vae, "decode_chunk_size", audio.numel() // vae_input.shape[0])) - new_audio = audio[pad_frames * dcs : (pad_frames + n_new) * dcs].detach().cpu().float() + with _NvtxRange("voxcpm2.vae_slice"): + dcs = int(getattr(self.tts.audio_vae, "decode_chunk_size", audio.numel() // vae_frames)) + new_audio = audio[pad_frames * dcs : (pad_frames + n_new) * dcs] # Roll the pad buffer: keep last N latent frames as context for next step. - all_latents = vae_input # [pad + new] - state.decode_pad = all_latents[-self._n_decode_pad_frames :].detach() + with _NvtxRange("voxcpm2.vae_pad_update"): + all_latents = vae_input # [pad + new] + state.decode_pad = all_latents[-self._n_decode_pad_frames :].detach() self._perf.stop("vae_decode") - return new_audio + if self._enable_delayed_audio_copy and self._audio_emit_every == 1: + is_stopping = self._should_stop_from_cached_logits(state) + self._enqueue_delayed_audio_copy(state, new_audio) + if is_stopping: + delayed_chunks.extend(self._drain_pending_audio_copies(state, force=True)) + return self._merge_audio_chunks(delayed_chunks) + + if self._audio_emit_every > 1: + is_stopping = self._should_stop_from_cached_logits(state) + audio_chunk = new_audio.detach() + if self._vae_output_storage_may_be_reused(): + audio_chunk = audio_chunk.clone() + state.pending_audio_chunks_gpu.append(audio_chunk) + if not is_stopping and len(state.pending_audio_chunks_gpu) < self._audio_emit_every: + return None + if len(state.pending_audio_chunks_gpu) == 1: + merged_audio = state.pending_audio_chunks_gpu[0] + else: + merged_audio = torch.cat([chunk.reshape(-1) for chunk in state.pending_audio_chunks_gpu], dim=0) + state.pending_audio_chunks_gpu.clear() + return merged_audio.detach().cpu().float() + if self._coalesce_audio_d2h: + audio_chunk = new_audio.detach() + if self._vae_output_storage_may_be_reused(): + audio_chunk = audio_chunk.clone() + return audio_chunk + return new_audio.detach().cpu().float() + + def _can_collect_audio_batch( + self, + states: list[_RequestState], + initial_delayed_chunks_by_req: dict[str, list[torch.Tensor] | None] | None, + ) -> bool: + if not getattr(self, "_enable_batched_vae_decode", False): + return False + if self._enable_delayed_audio_copy or self._audio_emit_every > 1 or self._enable_vae_cuda_graph: + return False + if initial_delayed_chunks_by_req and any(initial_delayed_chunks_by_req.values()): + return False + return len(states) > 1 + + def _collect_audio_batch( + self, + states: list[_RequestState], + initial_delayed_chunks_by_req: dict[str, list[torch.Tensor] | None] | None = None, + ) -> dict[str, torch.Tensor | None]: + if not self._can_collect_audio_batch(states, initial_delayed_chunks_by_req): + return { + state.request_id: self._collect_audio( + state, + initial_delayed_chunks=None + if initial_delayed_chunks_by_req is None + else initial_delayed_chunks_by_req.get(state.request_id), + ) + for state in states + } + + outputs: dict[str, torch.Tensor | None] = {state.request_id: None for state in states} + pending_by_shape: dict[ + tuple[torch.device, torch.dtype, int, int], + list[tuple[_RequestState, torch.Tensor, torch.Tensor, int, int]], + ] = {} + vae_decode_every = getattr(self, "_vae_decode_every", 1) + + for state in states: + patch = state.last_audio_patch_gpu + if patch is None: + continue + state.last_audio_patch_gpu = None + + new_latent = patch.reshape(-1, self._feat_dim).to(torch.float32) + if vae_decode_every > 1: + is_stopping = self._should_stop_from_cached_logits(state) + state.pending_vae_latents_gpu.append(new_latent.detach()) + if not is_stopping and len(state.pending_vae_latents_gpu) < vae_decode_every: + outputs[state.request_id] = None + continue + new_latent = torch.cat(state.pending_vae_latents_gpu, dim=0) + state.pending_vae_latents_gpu.clear() + + n_new = new_latent.shape[0] + if state.decode_pad is not None: + vae_input = torch.cat([state.decode_pad, new_latent], dim=0) + pad_frames = state.decode_pad.shape[0] + else: + vae_input = new_latent + pad_frames = 0 + + with _NvtxRange("voxcpm2.vae_prepare"): + feat = vae_input.unsqueeze(0).transpose(1, 2).contiguous() + key = (feat.device, feat.dtype, feat.shape[1], feat.shape[2]) + pending_by_shape.setdefault(key, []).append((state, feat, vae_input, pad_frames, n_new)) + + for group in pending_by_shape.values(): + with _NvtxRange("voxcpm2.vae_decode"): + with torch.no_grad(): + feat_batch = torch.cat([feat for _, feat, _, _, _ in group], dim=0).to(self._device) + audio_batch = self._run_vae_decode(feat_batch) + + for i, (state, _feat, vae_input, pad_frames, n_new) in enumerate(group): + audio = audio_batch[i].reshape(-1) + vae_frames = vae_input.shape[0] + with _NvtxRange("voxcpm2.vae_slice"): + dcs = int(getattr(self.tts.audio_vae, "decode_chunk_size", audio.numel() // vae_frames)) + new_audio = audio[pad_frames * dcs : (pad_frames + n_new) * dcs] + + with _NvtxRange("voxcpm2.vae_pad_update"): + state.decode_pad = vae_input[-self._n_decode_pad_frames :].detach() + + if self._coalesce_audio_d2h: + audio_chunk = new_audio.detach() + if self._vae_output_storage_may_be_reused(): + audio_chunk = audio_chunk.clone() + outputs[state.request_id] = audio_chunk + else: + outputs[state.request_id] = new_audio.detach().cpu().float() + + return outputs # -------------------- compute_logits -------------------- @@ -1128,12 +2387,15 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): logits[i, 0] = 0.0 logits[i, 1] = 1.0 state.precomputed_stop_logits = None + state.precomputed_is_stopping = None else: logits[i, 0] = stop_logits[0, 0] logits[i, 1] = stop_logits[0, 1] if state is not None: - state.is_stopping = bool(stop_logits[0, 1] > stop_logits[0, 0]) + if state.precomputed_is_stopping is not None: + state.is_stopping = state.precomputed_is_stopping state.precomputed_stop_logits = None + state.precomputed_is_stopping = None elif state and state.prefill_completed: logits[i, 1] = 1.0 else: @@ -1151,10 +2413,43 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): mm: dict[str, Any] = {} if self._audio_queue: - audio_by_req = {rid: audio for rid, audio in self._audio_queue} - order = [r for r, _ in self._audio_queue] - mm["model_outputs"] = [audio_by_req.get(r) for r in order] - mm["sr"] = [torch.tensor(self._sample_rate, dtype=torch.int32) for _ in order] + audio_by_req: dict[str, torch.Tensor] = {} + for req_id, audio in self._audio_queue: + if audio is None: + continue + if req_id in audio_by_req: + audio_by_req[req_id] = torch.cat([audio_by_req[req_id].reshape(-1), audio.reshape(-1)], dim=0) + else: + audio_by_req[req_id] = audio + if audio_by_req: + sr = torch.tensor(self._sample_rate, dtype=torch.int32) + if self._uses_sparse_audio_outputs(): + ready_req_ids = list(audio_by_req) + chunks = [audio_by_req[req_id].reshape(-1) for req_id in ready_req_ids] + if self._coalesce_audio_d2h and any(chunk.is_cuda for chunk in chunks): + sizes = [int(chunk.numel()) for chunk in chunks] + merged = torch.cat(chunks, dim=0) if len(chunks) > 1 else chunks[0] + merged_cpu = merged.detach().cpu().contiguous() + mm["model_outputs"] = list(merged_cpu.split(sizes)) + else: + mm["model_outputs"] = chunks + mm["sr"] = [sr for _ in ready_req_ids] + mm["meta"] = {"req_id": ready_req_ids, "sparse_audio": ["1"]} + elif self._coalesce_audio_d2h and any(audio.is_cuda for audio in audio_by_req.values()): + ready_req_ids = list(audio_by_req) + chunks = [audio_by_req[req_id].reshape(-1) for req_id in ready_req_ids] + sizes = [int(chunk.numel()) for chunk in chunks] + merged = torch.cat(chunks, dim=0) if len(chunks) > 1 else chunks[0] + merged_cpu = merged.detach().cpu().float() + mm["model_outputs"] = list(merged_cpu.split(sizes)) + mm["sr"] = [sr for _ in ready_req_ids] + else: + mm["model_outputs"] = list(audio_by_req.values()) + mm["sr"] = [sr for _ in audio_by_req] + elif self._uses_sparse_audio_outputs(): + mm["model_outputs"] = [] + mm["sr"] = [] + mm["meta"] = {"req_id": [], "sparse_audio": ["1"]} self._audio_queue.clear() return OmniOutput(text_hidden_states=model_outputs, multimodal_outputs=mm) @@ -1213,7 +2508,14 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): state.prefill_completed = False state.decode_step_count = 0 state.precomputed_stop_logits = None + state.precomputed_is_stopping = None state.last_audio_patch_gpu = None + if not hasattr(state, "pending_audio_chunks_gpu"): + state.pending_audio_chunks_gpu = [] + if not hasattr(state, "pending_audio_copies"): + state.pending_audio_copies = [] + state.pending_audio_chunks_gpu.clear() + state.pending_audio_copies.clear() state.curr_embed_for_next = None state.prev_feat_embed = None state.curr_prefix_feat_cond = None @@ -1307,7 +2609,7 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): def postprocess(self, hidden_states: torch.Tensor, **info: Any) -> dict[str, Any]: req_id = info.get("request_id", self._current_request_id or "default") - if _ENABLE_PROFILING: + if self._enable_profiling: state = self._active_states.get(req_id) if state and state.decode_step_count > 0: logger.info( diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index b267b8948..c59b3b4c7 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/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():