From 9edc1b4667d4b56d12effe1e5b300315520d1ae3 Mon Sep 17 00:00:00 2001 From: geray <48796550+gerayking@users.noreply.github.com> Date: Sat, 12 Sep 2026 20:31:49 +0800 Subject: [PATCH] Optimize CosyVoice3 Stage1 flow batching (#4876) Signed-off-by: gerayking <399geray@gmail.com> Co-authored-by: Hongsheng Liu --- .../examples/online_serving/text_to_speech.md | 9 + tests/config/test_environment_variables.py | 6 +- .../cosyvoice3/test_cosyvoice3_components.py | 252 ++++++++++++++++++ .../test_cosyvoice3_model_helpers.py | 64 +++++ .../config/environment_variable_inventory.py | 2 + .../models/cosyvoice3_audio/cosyvoice3_dit.py | 173 +++++++----- .../models/cosyvoice3/code2wav_core/cfm.py | 176 ++++++------ .../models/cosyvoice3/cosyvoice3.py | 97 +++++-- .../models/cosyvoice3/cosyvoice3_code2wav.py | 213 ++++++++++++--- .../models/cosyvoice3/runtime.py | 31 +++ 10 files changed, 825 insertions(+), 198 deletions(-) create mode 100644 vllm_omni/model_executor/models/cosyvoice3/runtime.py diff --git a/docs/user_guide/examples/online_serving/text_to_speech.md b/docs/user_guide/examples/online_serving/text_to_speech.md index 9a33dc858..2b4f617f4 100644 --- a/docs/user_guide/examples/online_serving/text_to_speech.md +++ b/docs/user_guide/examples/online_serving/text_to_speech.md @@ -120,6 +120,15 @@ vllm serve FunAudioLLM/Fun-CosyVoice3-0.5B-2512 --omni --port 8091 --trust-remot Streaming is on by default via `async_chunk: true` in `vllm_omni/deploy/cosyvoice3.yaml`. Pass `--no-async-chunk` (or `NO_ASYNC_CHUNK=1 ./cosyvoice3/run_server.sh`) for the legacy synchronous path. +Cross-request Stage-1 flow batching is opt-in. Enable it in the server environment when concurrent requests should share a flow-estimator call: + +```bash +export COSYVOICE3_BATCH_FLOW=1 +vllm serve FunAudioLLM/Fun-CosyVoice3-0.5B-2512 --omni --port 8091 --trust-remote-code +``` + +Batching preserves output lengths and streaming cache alignment, but the different GEMM shapes can produce small waveform differences compared with processing each request separately. Leave `COSYVOICE3_BATCH_FLOW` unset (or set it to `0`) when request-independent numerical behavior is required. Set `COSYVOICE3_BATCH_FLOW_DEBUG=1` to log the observed group-size distribution and enable detailed Stage-1 profiler scopes; diagnostics are disabled by default to avoid per-step profiling overhead. + ### CLI client The client defaults to the official upstream zero-shot prompt, so it runs without extra flags: diff --git a/tests/config/test_environment_variables.py b/tests/config/test_environment_variables.py index 7603945a0..cc37cd6b1 100644 --- a/tests/config/test_environment_variables.py +++ b/tests/config/test_environment_variables.py @@ -167,7 +167,7 @@ def test_inventory_matches_reviewed_snapshot_counts(): EnvironmentVariableCategory.PUBLIC_OMNI: 24, EnvironmentVariableCategory.INHERITED_VLLM: 20, EnvironmentVariableCategory.PLATFORM_EXTERNAL: 27, - EnvironmentVariableCategory.MODEL_SPECIFIC: 61, + EnvironmentVariableCategory.MODEL_SPECIFIC: 63, EnvironmentVariableCategory.BENCHMARK_TRANSITIONAL: 20, EnvironmentVariableCategory.INTERNAL: 2, } @@ -178,10 +178,10 @@ def test_inventory_matches_reviewed_snapshot_counts(): if item.category is EnvironmentVariableCategory.MODEL_SPECIFIC ) assert {disposition: disposition_counts[disposition] for disposition in ModelEnvironmentVariableDisposition} == { - ModelEnvironmentVariableDisposition.PROMOTE: 37, + ModelEnvironmentVariableDisposition.PROMOTE: 38, ModelEnvironmentVariableDisposition.REQUEST_SCOPE: 6, ModelEnvironmentVariableDisposition.EXTERNAL: 0, - ModelEnvironmentVariableDisposition.INTERNALIZE: 13, + ModelEnvironmentVariableDisposition.INTERNALIZE: 14, ModelEnvironmentVariableDisposition.DEPRECATE_REMOVE: 5, } diff --git a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_components.py b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_components.py index 4f01022f8..1b1df6367 100644 --- a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_components.py +++ b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_components.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Unit tests for CosyVoice3 components.""" +import types from contextlib import contextmanager from types import SimpleNamespace @@ -9,6 +10,7 @@ import pytest import torch import torch.nn as nn +from tests.helpers.mark import hardware_test from vllm_omni.diffusion.config import set_current_diffusion_config from vllm_omni.diffusion.data import AttentionConfig from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan import ( @@ -29,6 +31,23 @@ def _force_torch_sdpa(): yield +def test_batch_flow_runtime_flags_are_opt_in(monkeypatch): + from vllm_omni.model_executor.models.cosyvoice3.runtime import ( + cosyvoice3_batch_flow_debug, + cosyvoice3_batch_flow_enabled, + ) + + monkeypatch.delenv("COSYVOICE3_BATCH_FLOW", raising=False) + monkeypatch.delenv("COSYVOICE3_BATCH_FLOW_DEBUG", raising=False) + assert not cosyvoice3_batch_flow_enabled() + assert not cosyvoice3_batch_flow_debug() + + monkeypatch.setenv("COSYVOICE3_BATCH_FLOW", "1") + monkeypatch.setenv("COSYVOICE3_BATCH_FLOW_DEBUG", "1") + assert cosyvoice3_batch_flow_enabled() + assert cosyvoice3_batch_flow_debug() + + @pytest.fixture def causal_hift(): return CausalHiFTGenerator( @@ -139,6 +158,39 @@ class TestDiTAttention: assert attention.to_k.out_features == 512 assert attention.to_v.out_features == 512 + @pytest.mark.core_model + @pytest.mark.cpu + def test_forward_passes_2d_mask_without_casting_qkv(self): + """The float32 fallback receives the per-row validity mask.""" + from vllm_omni.diffusion.models.cosyvoice3_audio.cosyvoice3_dit import DiTAttention + + class CaptureAttention(nn.Module): + def __init__(self): + super().__init__() + self.last_metadata = None + self.last_dtype = None + + def forward(self, q, k, v, attn_metadata=None): + self.last_metadata = attn_metadata + self.last_dtype = q.dtype + return torch.zeros_like(q) + + attention = DiTAttention(dim=16, heads=2, dim_head=8, dropout=0.0) + capture = CaptureAttention() + attention.attn = capture + + x = torch.randn(2, 5, 16) + mask = torch.tensor([[True, True, True, False, False], [True, True, True, True, False]]) + + out = attention(x, mask=mask) + + assert out.shape == x.shape + assert capture.last_metadata is not None + assert torch.equal(capture.last_metadata.attn_mask, mask) + assert capture.last_dtype == x.dtype + assert out.dtype == x.dtype + assert torch.allclose(out[~mask], torch.zeros_like(out[~mask])) + class TestDiTBlock: """Tests for DiTBlock.""" @@ -422,6 +474,52 @@ class TestCFM: assert reused_context is context assert reused_stream is stream + @pytest.mark.core_model + @pytest.mark.cpu + def test_causal_conditional_cfm_batches_cfg_estimator(self): + """Batched flow calls should invoke the CFG estimator with 2B rows.""" + from omegaconf import DictConfig + + from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.cfm import CausalConditionalCFM + + class RecordingEstimator(nn.Module): + def __init__(self): + super().__init__() + self.shapes: list[tuple[int, ...]] = [] + + def forward(self, x, mask, mu, t, spks=None, cond=None): + self.shapes.append(tuple(x.shape)) + return torch.zeros_like(x) + + estimator = RecordingEstimator() + cfm_params = DictConfig( + { + "sigma_min": 1e-6, + "solver": "euler", + "t_scheduler": "cosine", + "training_cfg_rate": 0.2, + "inference_cfg_rate": 0.7, + } + ) + cfm = CausalConditionalCFM( + in_channels=80, + cfm_params=cfm_params, + n_spks=1, + spk_emb_dim=80, + estimator=estimator, + ) + + batch, mel_dim, seq_len = 3, 80, 16 + mu = torch.randn(batch, mel_dim, seq_len) + mask = torch.ones(batch, 1, seq_len) + spks = torch.randn(batch, 80) + cond = torch.randn(batch, mel_dim, seq_len) + + out, _ = cfm(mu, mask, n_timesteps=2, spks=spks, cond=cond) + + assert out.shape == mu.shape + assert estimator.shapes == [(2 * batch, mel_dim, seq_len)] * 2 + class TestSDPAFallback: """Test SDPA fallback for float32 inputs.""" @@ -484,3 +582,157 @@ def test_code2wav_forward_finalizes_hift_tail(): assert out.shape == (1, 1, 8) assert model.hift.finalize_calls == [True] assert forward_mel_calls[0]["token_offset_tokens"] == 0 + + +def test_code2wav_streaming_batch_pads_codec_tokens_and_preserves_lengths(): + from vllm_omni.model_executor.models.cosyvoice3.cosyvoice3_code2wav import CosyVoice3Code2Wav + + model = object.__new__(CosyVoice3Code2Wav) + model.flow_model = SimpleNamespace(pre_lookahead_len=1, token_mel_ratio=2) + forward_mel_calls = [] + + def fake_forward_mel(self, **kwargs): + forward_mel_calls.append(kwargs) + return torch.arange(2 * 80 * 8, dtype=torch.float32).reshape(2, 80, 8) + + def fake_stream_hift(self, feat, *, cache_state=None, finalize=False): + return feat, None + + model._forward_mel = types.MethodType(fake_forward_mel, model) + model._stream_hift_from_feat = types.MethodType(fake_stream_hift, model) + items = [ + { + "token": torch.ones(1, 3, dtype=torch.int32), + "prompt_token": torch.ones(1, 4, dtype=torch.int32), + "prompt_feat": torch.ones(1, 8, 80), + "embedding": torch.ones(1, 192), + "finalize": False, + }, + { + "token": torch.ones(1, 5, dtype=torch.int32), + "prompt_token": torch.ones(1, 4, dtype=torch.int32), + "prompt_feat": torch.ones(1, 8, 80), + "embedding": torch.ones(1, 192), + "finalize": False, + }, + ] + + results = model.forward_streaming_batch(items) + + assert len(forward_mel_calls) == 1 + call = forward_mel_calls[0] + assert call["token"].shape == (2, 5) + assert torch.equal(call["token_lens"], torch.tensor([3, 5], dtype=torch.int32)) + assert results[0][0].shape[-1] == 4 + assert results[1][0].shape[-1] == 8 + + +@pytest.mark.core_model +@hardware_test(res={"cuda": "L4"}, num_cards=1) +def test_code2wav_streaming_batch_matches_ragged_flow_numerics(monkeypatch): + """A padded flow batch must match individual flow calls on valid mels.""" + from omegaconf import DictConfig + + from vllm_omni.diffusion.models.cosyvoice3_audio.cosyvoice3_dit import DiT + from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.cfm import ( + CausalConditionalCFM, + CausalMaskedDiffWithDiT, + ) + from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.layers import PreLookaheadLayer + from vllm_omni.model_executor.models.cosyvoice3.cosyvoice3_code2wav import CosyVoice3Code2Wav + + torch.manual_seed(0) + estimator = DiT( + dim=32, + depth=1, + heads=4, + dim_head=8, + dropout=0.0, + ff_mult=2, + mel_dim=80, + mu_dim=80, + spk_dim=80, + out_channels=80, + ) + decoder = CausalConditionalCFM( + in_channels=80, + cfm_params=DictConfig( + { + "sigma_min": 1e-6, + "solver": "euler", + "t_scheduler": "cosine", + "training_cfg_rate": 0.2, + "inference_cfg_rate": 0.7, + } + ), + n_spks=1, + spk_emb_dim=80, + estimator=estimator, + ) + flow_model = ( + CausalMaskedDiffWithDiT( + input_size=80, + output_size=80, + spk_embed_dim=192, + vocab_size=64, + input_frame_rate=25, + only_mask_loss=True, + token_mel_ratio=2, + pre_lookahead_len=1, + pre_lookahead_layer=PreLookaheadLayer(in_channels=80, channels=80, pre_lookahead_len=1), + decoder=decoder, + ) + .to(device="cuda", dtype=torch.bfloat16) + .eval() + ) + + model = object.__new__(CosyVoice3Code2Wav) + nn.Module.__init__(model) + model.flow_model = flow_model + + def return_mel(self, feat, *, cache_state=None, finalize=False): + return feat, None + + model._stream_hift_from_feat = types.MethodType(return_mel, model) + + original_randn = torch.randn + + def length_consistent_randn(*size, **kwargs): + shape = tuple(size[0]) if len(size) == 1 and isinstance(size[0], (tuple, list)) else tuple(size) + if len(shape) == 3 and shape[1] == 80: + device = kwargs.get("device") + dtype = kwargs.get("dtype", torch.float32) + channels = torch.arange(shape[1], device=device, dtype=torch.float32).view(1, -1, 1) + positions = torch.arange(shape[2], device=device, dtype=torch.float32).view(1, 1, -1) + noise = torch.sin(channels * 0.17 + positions * 0.31) + return noise.expand(shape[0], -1, -1).to(dtype=dtype).clone() + return original_randn(*size, **kwargs) + + monkeypatch.setattr(torch, "randn", length_consistent_randn) + common = { + "prompt_token": torch.tensor([[7, 8]], dtype=torch.int32), + "prompt_feat": torch.linspace(-0.5, 0.5, 4 * 80).reshape(1, 4, 80), + "embedding": torch.linspace(-1.0, 1.0, 192).reshape(1, 192), + "finalize": False, + } + items = [ + {**common, "token": torch.tensor([[1, 2, 3]], dtype=torch.int32)}, + {**common, "token": torch.tensor([[1, 2, 3, 4, 5]], dtype=torch.int32)}, + ] + + batched = model.forward_streaming_batch(items, n_timesteps=2) + individual = [ + model.forward_streaming( + token=item["token"], + prompt_token=item["prompt_token"], + prompt_feat=item["prompt_feat"], + embedding=item["embedding"], + n_timesteps=2, + ) + for item in items + ] + + for (batched_mel, _), (individual_mel, _) in zip(batched, individual): + assert batched_mel.shape == individual_mel.shape + rel_mean = (batched_mel - individual_mel).abs().mean() / individual_mel.abs().mean().clamp_min(1e-6) + assert rel_mean.item() < 0.05 diff --git a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_model_helpers.py b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_model_helpers.py index e16927096..1d583de6c 100644 --- a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_model_helpers.py +++ b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_model_helpers.py @@ -42,6 +42,7 @@ class _DummyCode2Wav: self.outputs = list(outputs or []) self.forward_calls: list[dict[str, object]] = [] self.forward_streaming_calls: list[dict[str, object]] = [] + self.forward_streaming_batch_calls: list[list[dict[str, object]]] = [] def forward(self, **kwargs): self.forward_calls.append(kwargs) @@ -65,6 +66,22 @@ class _DummyCode2Wav: } return audio, new_state + def forward_streaming_batch(self, items, *, n_timesteps: int = 10): + self.forward_streaming_batch_calls.append(items) + return [ + self.forward_streaming( + token=item["token"], + prompt_token=item["prompt_token"], + prompt_feat=item["prompt_feat"], + embedding=item["embedding"], + cache_state=item.get("cache_state"), + n_timesteps=n_timesteps, + token_offset_tokens=int(item.get("token_offset_tokens", 0)), + finalize=bool(item.get("finalize", False)), + ) + for item in items + ] + def _make_code2wav_model( *, @@ -373,6 +390,53 @@ def test_forward_clears_streaming_cache_on_terminal_chunk(): assert "rid-stream" not in model._stream_vocoder_cache_by_req +def test_forward_batches_streaming_flow_items(monkeypatch): + monkeypatch.setenv("COSYVOICE3_BATCH_FLOW", "1") + model = _make_code2wav_model() + runtime_info = [ + { + "embed": { + "speech_token": torch.tensor([[1, 2, 3]], dtype=torch.long), + "speech_feat": torch.tensor([[[0.1, 0.2], [0.3, 0.4]]], dtype=torch.float32), + "embedding": torch.tensor([[0.5, 0.6]], dtype=torch.float32), + }, + "meta": { + "req_id": ["rid-a"], + "stream_finished": torch.tensor(False), + "left_context_size": 0, + }, + }, + { + "embed": { + "speech_token": torch.tensor([[1, 2, 3]], dtype=torch.long), + "speech_feat": torch.tensor([[[0.1, 0.2], [0.3, 0.4]]], dtype=torch.float32), + "embedding": torch.tensor([[0.7, 0.8]], dtype=torch.float32), + }, + "meta": { + "req_id": ["rid-b"], + "stream_finished": torch.tensor(False), + "left_context_size": 1, + }, + }, + ] + + out = model.forward( + input_ids=torch.tensor([0, 1, 2, 0, 1, 2], dtype=torch.long), + positions=torch.arange(6, dtype=torch.long), + model_intermediate_buffer=runtime_info, + seq_token_counts=[3, 3], + ) + + assert len(out.multimodal_outputs["audio"]) == 2 + assert len(model.code2wav.forward_streaming_batch_calls) == 1 + batch_items = model.code2wav.forward_streaming_batch_calls[0] + assert [item["index"] for item in batch_items] == [0, 1] + assert torch.equal(batch_items[0]["token"], torch.tensor([[0, 1, 2]])) + assert batch_items[1]["token_offset_tokens"] == 1 + assert "rid-a" in model._stream_vocoder_cache_by_req + assert "rid-b" in model._stream_vocoder_cache_by_req + + def test_sample_uses_ras_rejection_for_recent_repetition(): model = _make_talker_model() metadata = _make_sampling_metadata(output_token_ids=[[1] * 10]) diff --git a/vllm_omni/config/environment_variable_inventory.py b/vllm_omni/config/environment_variable_inventory.py index c1e380a22..4a4966d9c 100644 --- a/vllm_omni/config/environment_variable_inventory.py +++ b/vllm_omni/config/environment_variable_inventory.py @@ -137,6 +137,7 @@ _MODEL_PROMOTE = ( "AR_DIFFUSION_KV_SCRATCH_BLOCKS_PER_BRANCH", "COSMOS3_SOUND_TOKENIZER_CONFIG_PATH", "COSMOS3_SOUND_TOKENIZER_PATH", + "COSYVOICE3_BATCH_FLOW", "COSYVOICE3_ESTIMATOR_ONNX", "COSYVOICE3_TRT", "COSYVOICE3_TRT_CACHE", @@ -188,6 +189,7 @@ _MODEL_REQUEST_SCOPE = ( _MODEL_EXTERNAL: tuple[str, ...] = () _MODEL_INTERNALIZE = ( + "COSYVOICE3_BATCH_FLOW_DEBUG", "DOTS_TTS_BETA_TRACE", "DYNIN_S2U_VENDOR_ROOT", "DZ_PHASE_TIMING", diff --git a/vllm_omni/diffusion/models/cosyvoice3_audio/cosyvoice3_dit.py b/vllm_omni/diffusion/models/cosyvoice3_audio/cosyvoice3_dit.py index 5ae75b256..d0538d79f 100644 --- a/vllm_omni/diffusion/models/cosyvoice3_audio/cosyvoice3_dit.py +++ b/vllm_omni/diffusion/models/cosyvoice3_audio/cosyvoice3_dit.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project # Adapted from https://github.com/FunAudioLLM/CosyVoice/tree/main/cosyvoice/flow/DiT # Refactored to use vllm_omni diffusion infrastructure for optimized attention backends @@ -15,8 +15,10 @@ from torch import nn from vllm.logger import init_logger from x_transformers.x_transformers import RotaryEmbedding, apply_rotary_pos_emb +from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata from vllm_omni.diffusion.attention.layer import Attention as DiffusionAttention from vllm_omni.model_executor.layers.timestep_embedding import DiTTimestepEmbedding +from vllm_omni.model_executor.models.cosyvoice3.runtime import cosyvoice3_batch_flow_profile logger = init_logger(__name__) @@ -52,6 +54,22 @@ def get_pos_embed_indices(start, length, max_pos, scale=1.0): return pos +def _normalize_attention_mask(mask: torch.Tensor | None, batch_size: int, seq_len: int) -> torch.Tensor | None: + if mask is None: + return None + if mask.dim() == 2: + attn_mask = mask + elif mask.dim() == 3 and mask.shape[1] == 1: + attn_mask = mask[:, 0] + elif mask.dim() == 4: + attn_mask = mask[:, 0, -1] + else: + return None + if attn_mask.shape != (batch_size, seq_len): + return None + return attn_mask.bool() + + class FeedForward(nn.Module): """Feed-forward network with GELU activation.""" @@ -114,42 +132,39 @@ class DiTAttention(nn.Module): ) -> torch.Tensor: batch_size, seq_len = x.shape[0], x.shape[1] - # Project to Q, K, V - query = self.to_q(x) - key = self.to_k(x) - value = self.to_v(x) + with cosyvoice3_batch_flow_profile("cosyvoice3_dit_attention_qkv_projection"): + query = self.to_q(x) + key = self.to_k(x) + value = self.to_v(x) - # Apply rotary position embedding - if rope is not None: - freqs, xpos_scale = rope - q_xpos_scale, k_xpos_scale = (xpos_scale, xpos_scale**-1.0) if xpos_scale is not None else (1.0, 1.0) - query = apply_rotary_pos_emb(query, freqs, q_xpos_scale) - key = apply_rotary_pos_emb(key, freqs, k_xpos_scale) + with cosyvoice3_batch_flow_profile("cosyvoice3_dit_attention_rope"): + if rope is not None: + freqs, xpos_scale = rope + q_xpos_scale, k_xpos_scale = (xpos_scale, xpos_scale**-1.0) if xpos_scale is not None else (1.0, 1.0) + query = apply_rotary_pos_emb(query, freqs, q_xpos_scale) + key = apply_rotary_pos_emb(key, freqs, k_xpos_scale) - # Reshape for attention: (batch, seq, heads, head_dim) - query = query.view(batch_size, seq_len, self.heads, self.dim_head) - key = key.view(batch_size, seq_len, self.heads, self.dim_head) - value = value.view(batch_size, seq_len, self.heads, self.dim_head) + with cosyvoice3_batch_flow_profile("cosyvoice3_dit_attention_backend"): + # Reshape for attention: (batch, seq, heads, head_dim) + query = query.view(batch_size, seq_len, self.heads, self.dim_head) + key = key.view(batch_size, seq_len, self.heads, self.dim_head) + value = value.view(batch_size, seq_len, self.heads, self.dim_head) - # Use diffusion attention backend - # The diffusion Attention layer expects (batch, seq, heads, head_dim) - out = self.attn(query, key, value, attn_metadata=None) + # The diffusion Attention layer expects (batch, seq, heads, head_dim) + attn_mask = _normalize_attention_mask(mask, batch_size, seq_len) + attn_metadata = AttentionMetadata(attn_mask=attn_mask) if attn_mask is not None else None + out = self.attn(query, key, value, attn_metadata=attn_metadata) - # Some attention backends return a non-contiguous layout. - out = out.reshape(batch_size, seq_len, self.inner_dim) - out = out.to(query.dtype) - - # Output projection - out = self.to_out(out) + with cosyvoice3_batch_flow_profile("cosyvoice3_dit_attention_output_projection"): + # Some attention backends return a non-contiguous layout. + out = out.reshape(batch_size, seq_len, self.inner_dim) + out = out.to(x.dtype) + out = self.to_out(out) # Apply mask if provided - if mask is not None: - if mask.dim() == 2: - mask = mask.unsqueeze(-1) - elif mask.dim() == 4: - # (batch, heads, seq, seq) -> use last dim - mask = mask[:, 0, -1].unsqueeze(-1) - out = out.masked_fill(~mask.bool(), 0.0) + attn_mask = _normalize_attention_mask(mask, batch_size, seq_len) + if attn_mask is not None: + out = out.masked_fill(~attn_mask.unsqueeze(-1), 0.0) return out @@ -172,18 +187,19 @@ class DiTBlock(nn.Module): self.ff = FeedForward(dim=dim, mult=ff_mult, dropout=dropout, approximate="tanh") def forward(self, x, t, mask=None, rope=None): - # pre-norm & modulation for attention input - norm, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.attn_norm(x, emb=t) + with cosyvoice3_batch_flow_profile("cosyvoice3_dit_adalayernorm_attention"): + norm, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.attn_norm(x, emb=t) - # attention - attn_output = self.attn(x=norm, mask=mask, rope=rope) + with cosyvoice3_batch_flow_profile("cosyvoice3_dit_attention"): + attn_output = self.attn(x=norm, mask=mask, rope=rope) - # process attention output for input x - x = x + gate_msa.unsqueeze(1) * attn_output + with cosyvoice3_batch_flow_profile("cosyvoice3_dit_attention_residual"): + x = x + gate_msa.unsqueeze(1) * attn_output - ff_norm = self.ff_norm(x) * (1 + scale_mlp[:, None]) + shift_mlp[:, None] - ff_output = self.ff(ff_norm) - x = x + gate_mlp.unsqueeze(1) * ff_output + with cosyvoice3_batch_flow_profile("cosyvoice3_dit_layernorm_ffn"): + ff_norm = self.ff_norm(x) * (1 + scale_mlp[:, None]) + shift_mlp[:, None] + ff_output = self.ff(ff_norm) + x = x + gate_mlp.unsqueeze(1) * ff_output return x @@ -205,19 +221,20 @@ class CausalConvPositionEmbedding(nn.Module): ) def forward(self, x: torch.Tensor, mask: torch.Tensor | None = None): - if mask is not None: - mask = mask[..., None] - x = x.masked_fill(~mask, 0.0) + with cosyvoice3_batch_flow_profile("cosyvoice3_dit_causal_conv_position_embedding"): + if mask is not None: + mask = mask[..., None] + x = x.masked_fill(~mask, 0.0) - x = x.permute(0, 2, 1) - x = F.pad(x, (self.kernel_size - 1, 0, 0, 0)) - x = self.conv1(x) - x = F.pad(x, (self.kernel_size - 1, 0, 0, 0)) - x = self.conv2(x) - out = x.permute(0, 2, 1) + x = x.permute(0, 2, 1) + x = F.pad(x, (self.kernel_size - 1, 0, 0, 0)) + x = self.conv1(x) + x = F.pad(x, (self.kernel_size - 1, 0, 0, 0)) + x = self.conv2(x) + out = x.permute(0, 2, 1) - if mask is not None: - out = out.masked_fill(~mask, 0.0) + if mask is not None: + out = out.masked_fill(~mask, 0.0) return out @@ -327,13 +344,15 @@ class InputEmbedding(nn.Module): self.conv_pos_embed = CausalConvPositionEmbedding(dim=out_dim) def forward(self, x, cond, text_embed, spks): - to_cat = [x, cond, text_embed] - if self.spk_dim > 0: - spks = repeat(spks, "b c -> b t c", t=x.shape[1]) - to_cat.append(spks) + with cosyvoice3_batch_flow_profile("cosyvoice3_dit_input_projection"): + to_cat = [x, cond, text_embed] + if self.spk_dim > 0: + spks = repeat(spks, "b c -> b t c", t=x.shape[1]) + to_cat.append(spks) - x = self.proj(torch.cat(to_cat, dim=-1)) - x = self.conv_pos_embed(x) + x + x = self.proj(torch.cat(to_cat, dim=-1)) + with cosyvoice3_batch_flow_profile("cosyvoice3_dit_input_causal_conv"): + x = self.conv_pos_embed(x) + x return x @@ -385,31 +404,37 @@ class DiT(nn.Module): self.num_decoding_left_chunks = num_decoding_left_chunks def forward(self, x, mask, mu, t, spks=None, cond=None): - x = x.transpose(1, 2) - mu = mu.transpose(1, 2) - cond = cond.transpose(1, 2) - spks = spks.unsqueeze(dim=1) - batch, seq_len = x.shape[0], x.shape[1] - if t.ndim == 0: - t = t.repeat(batch) + with cosyvoice3_batch_flow_profile("cosyvoice3_dit_prepare_inputs"): + x = x.transpose(1, 2) + mu = mu.transpose(1, 2) + cond = cond.transpose(1, 2) + spks = spks.unsqueeze(dim=1) + batch, seq_len = x.shape[0], x.shape[1] + if t.ndim == 0: + t = t.repeat(batch) - # t: conditioning time, c: context (text + masked cond audio), x: noised input audio - t = self.time_embed(t) + with cosyvoice3_batch_flow_profile("cosyvoice3_dit_time_embedding"): + t = self.time_embed(t) x = self.input_embed(x, cond, mu, spks.squeeze(1)) - rope = self.rotary_embed.forward_from_seq_len(seq_len) + with cosyvoice3_batch_flow_profile("cosyvoice3_dit_rope_and_mask"): + rope = self.rotary_embed.forward_from_seq_len(seq_len) + attn_mask = mask[:, 0].bool() if mask.dim() == 3 and mask.shape[1] == 1 else mask.bool() if self.long_skip_connection is not None: residual = x - attn_mask = mask.bool().repeat(1, x.size(1), 1).unsqueeze(dim=1) - - for block in self.transformer_blocks: - x = block(x, t, mask=attn_mask.bool(), rope=rope) + with cosyvoice3_batch_flow_profile("cosyvoice3_dit_transformer_blocks"): + for block_index, block in enumerate(self.transformer_blocks): + with cosyvoice3_batch_flow_profile(f"cosyvoice3_dit_transformer_block_{block_index}"): + x = block(x, t, mask=attn_mask.bool(), rope=rope) if self.long_skip_connection is not None: - x = self.long_skip_connection(torch.cat((x, residual), dim=-1)) + with cosyvoice3_batch_flow_profile("cosyvoice3_dit_long_skip_projection"): + x = self.long_skip_connection(torch.cat((x, residual), dim=-1)) - x = self.norm_out(x, t) - output = self.proj_out(x).transpose(1, 2) + with cosyvoice3_batch_flow_profile("cosyvoice3_dit_final_adalayernorm"): + x = self.norm_out(x, t) + with cosyvoice3_batch_flow_profile("cosyvoice3_dit_final_projection"): + output = self.proj_out(x).transpose(1, 2) return output diff --git a/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/cfm.py b/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/cfm.py index 2ec12a251..c86b4a305 100644 --- a/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/cfm.py +++ b/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/cfm.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project # Adopted from https://github.com/FunAudioLLM/CosyVoice/tree/main/cosyvoice/flow """Conditional Flow Matching (CFM) classes for audio generation.""" @@ -11,6 +11,7 @@ from omegaconf import DictConfig from torch.nn import functional as F from vllm.logger import init_logger +from vllm_omni.model_executor.models.cosyvoice3.runtime import cosyvoice3_batch_flow_profile from vllm_omni.model_executor.models.cosyvoice3.utils import make_pad_mask logger = init_logger(__name__) @@ -74,20 +75,23 @@ class ConditionalCFM(BASECFM): shape: (batch_size, n_feats, mel_timesteps) """ - z = torch.randn_like(mu).to(mu.device).to(mu.dtype) * temperature - cache_size = cache.shape[2] - # fix prompt and overlap part mu and z - if cache_size != 0: - z[:, :, :cache_size] = cache[:, :, :, 0] - mu[:, :, :cache_size] = cache[:, :, :, 1] - z_cache = torch.concat([z[:, :, :prompt_len], z[:, :, -34:]], dim=2) - mu_cache = torch.concat([mu[:, :, :prompt_len], mu[:, :, -34:]], dim=2) - cache = torch.stack([z_cache, mu_cache], dim=-1) + with cosyvoice3_batch_flow_profile("cosyvoice3_cfm_noise_cache"): + z = torch.randn_like(mu).to(mu.device).to(mu.dtype) * temperature + cache_size = cache.shape[2] + # fix prompt and overlap part mu and z + if cache_size != 0: + z[:, :, :cache_size] = cache[:, :, :, 0] + mu[:, :, :cache_size] = cache[:, :, :, 1] + z_cache = torch.concat([z[:, :, :prompt_len], z[:, :, -34:]], dim=2) + mu_cache = torch.concat([mu[:, :, :prompt_len], mu[:, :, -34:]], dim=2) + cache = torch.stack([z_cache, mu_cache], dim=-1) - t_span = torch.linspace(0, 1, n_timesteps + 1, device=mu.device, dtype=mu.dtype) - if self.t_scheduler == "cosine": - t_span = 1 - torch.cos(t_span * 0.5 * torch.pi) - return self.solve_euler(z, t_span=t_span, mu=mu, mask=mask, spks=spks, cond=cond), cache + with cosyvoice3_batch_flow_profile("cosyvoice3_cfm_t_span"): + t_span = torch.linspace(0, 1, n_timesteps + 1, device=mu.device, dtype=mu.dtype) + if self.t_scheduler == "cosine": + t_span = 1 - torch.cos(t_span * 0.5 * torch.pi) + with cosyvoice3_batch_flow_profile(f"cosyvoice3_cfm_euler_{max(1, int(n_timesteps))}_steps"): + return self.solve_euler(z, t_span=t_span, mu=mu, mask=mask, spks=spks, cond=cond), cache def solve_euler(self, x, t_span, mu, mask, spks, cond): """ @@ -111,26 +115,39 @@ class ConditionalCFM(BASECFM): # Do not use concat, it may cause memory format changed and trt infer with wrong results! # NOTE when flow run in amp mode, x.dtype is float32, which cause nan in trt fp16 - # inference, so set dtype=spks.dtype - x_in = torch.zeros([2, 80, x.size(2)], device=x.device, dtype=spks.dtype) - mask_in = torch.zeros([2, 1, x.size(2)], device=x.device, dtype=spks.dtype) - mu_in = torch.zeros([2, 80, x.size(2)], device=x.device, dtype=spks.dtype) - t_in = torch.zeros([2], device=x.device, dtype=spks.dtype) - spks_in = torch.zeros([2, 80], device=x.device, dtype=spks.dtype) - cond_in = torch.zeros([2, 80, x.size(2)], device=x.device, dtype=spks.dtype) + # inference, so set dtype=spks.dtype. The batch is doubled for CFG: + # first B rows are conditioned, second B rows are unconditional. + batch_size = int(x.size(0)) + estimator_batch = 2 * batch_size + estimator_dtype = spks.dtype if spks is not None else x.dtype + with cosyvoice3_batch_flow_profile("cosyvoice3_cfm_cfg_allocate_2b"): + x_in = torch.zeros([estimator_batch, 80, x.size(2)], device=x.device, dtype=estimator_dtype) + mask_in = torch.zeros([estimator_batch, 1, x.size(2)], device=x.device, dtype=estimator_dtype) + mu_in = torch.zeros([estimator_batch, 80, x.size(2)], device=x.device, dtype=estimator_dtype) + t_in = torch.zeros([estimator_batch], device=x.device, dtype=estimator_dtype) + spks_in = torch.zeros([estimator_batch, 80], device=x.device, dtype=estimator_dtype) + cond_in = torch.zeros([estimator_batch, 80, x.size(2)], device=x.device, dtype=estimator_dtype) for step in range(1, len(t_span)): # Classifier-Free Guidance inference introduced in VoiceBox - x_in[:] = x - mask_in[:] = mask - mu_in[0] = mu - t_in[:] = t.unsqueeze(0) - spks_in[0] = spks - cond_in[0] = cond - dphi_dt = self.forward_estimator(x_in, mask_in, mu_in, t_in, spks_in, cond_in) - dphi_dt, cfg_dphi_dt = torch.split(dphi_dt, [x.size(0), x.size(0)], dim=0) - dphi_dt = (1.0 + self.inference_cfg_rate) * dphi_dt - self.inference_cfg_rate * cfg_dphi_dt - x = x + dt * dphi_dt - t = t + dt + with cosyvoice3_batch_flow_profile("cosyvoice3_cfm_cfg_prepare_2b"): + x_in[:batch_size] = x + x_in[batch_size:] = x + mask_in[:batch_size] = mask + mask_in[batch_size:] = mask + mu_in[:batch_size] = mu + t_in[:] = t + if spks is not None: + spks_in[:batch_size] = spks + if cond is not None: + cond_in[:batch_size] = cond + with cosyvoice3_batch_flow_profile("cosyvoice3_cfm_forward_estimator"): + dphi_dt = self.forward_estimator(x_in, mask_in, mu_in, t_in, spks_in, cond_in) + with cosyvoice3_batch_flow_profile("cosyvoice3_cfm_cfg_combine"): + dphi_dt, cfg_dphi_dt = torch.split(dphi_dt, [batch_size, batch_size], dim=0) + dphi_dt = (1.0 + self.inference_cfg_rate) * dphi_dt - self.inference_cfg_rate * cfg_dphi_dt + with cosyvoice3_batch_flow_profile("cosyvoice3_cfm_euler_update"): + x = x + dt * dphi_dt + t = t + dt sol.append(x) if step < len(t_span) - 1: dt = t_span[step + 1] - t @@ -158,12 +175,12 @@ class ConditionalCFM(BASECFM): spks_e = spks.to(io_dtype).contiguous() cond_e = cond.to(io_dtype).contiguous() out_e = torch.empty_like(x_e) - estimator.set_input_shape("x", (2, 80, x_e.size(2))) - estimator.set_input_shape("mask", (2, 1, x_e.size(2))) - estimator.set_input_shape("mu", (2, 80, x_e.size(2))) - estimator.set_input_shape("t", (2,)) - estimator.set_input_shape("spks", (2, 80)) - estimator.set_input_shape("cond", (2, 80, x_e.size(2))) + estimator.set_input_shape("x", tuple(x_e.shape)) + estimator.set_input_shape("mask", tuple(mask_e.shape)) + estimator.set_input_shape("mu", tuple(mu_e.shape)) + estimator.set_input_shape("t", tuple(t_e.shape)) + estimator.set_input_shape("spks", tuple(spks_e.shape)) + estimator.set_input_shape("cond", tuple(cond_e.shape)) data_ptrs = [ x_e.data_ptr(), mask_e.data_ptr(), @@ -211,22 +228,25 @@ class CausalConditionalCFM(ConditionalCFM): shape: (batch_size, n_feats, mel_timesteps) """ - z = ( - torch.randn( - (mu.size(0), mu.size(1), mu.size(2)), - device=mu.device, - dtype=mu.dtype, + with cosyvoice3_batch_flow_profile("cosyvoice3_cfm_noise_cache"): + z = ( + torch.randn( + (mu.size(0), mu.size(1), mu.size(2)), + device=mu.device, + dtype=mu.dtype, + ) + * temperature ) - * temperature - ) # fix prompt and overlap part mu and z - t_span = torch.linspace(0, 1, n_timesteps + 1, device=mu.device, dtype=mu.dtype) + with cosyvoice3_batch_flow_profile("cosyvoice3_cfm_t_span"): + t_span = torch.linspace(0, 1, n_timesteps + 1, device=mu.device, dtype=mu.dtype) - if self.t_scheduler == "cosine": - t_span = 1 - torch.cos(t_span * 0.5 * torch.pi) + if self.t_scheduler == "cosine": + t_span = 1 - torch.cos(t_span * 0.5 * torch.pi) - return self.solve_euler(z, t_span=t_span, mu=mu, mask=mask, spks=spks, cond=cond), None + with cosyvoice3_batch_flow_profile(f"cosyvoice3_cfm_euler_{max(1, int(n_timesteps))}_steps"): + return self.solve_euler(z, t_span=t_span, mu=mu, mask=mask, spks=spks, cond=cond), None class CausalMaskedDiffWithDiT(torch.nn.Module): @@ -299,36 +319,41 @@ class CausalMaskedDiffWithDiT(torch.nn.Module): finalize: bool = False, n_timesteps: int = 10, ): - assert token.shape[0] == 1 - # xvec projection + with cosyvoice3_batch_flow_profile("cosyvoice3_cfm_speaker_embedding"): + embedding = F.normalize(embedding, dim=1) + embedding = self.spk_embed_affine_layer(embedding) - embedding = F.normalize(embedding, dim=1) + with cosyvoice3_batch_flow_profile("cosyvoice3_cfm_token_embedding_lookahead"): + # concat text and prompt_text + codec_token_len = token_len + token, total_token_len = torch.concat([prompt_token, token], dim=1), prompt_token_len + codec_token_len + mask = (~make_pad_mask(total_token_len, max_len=token.shape[1])).unsqueeze(-1).to(embedding) + token = self.input_embedding(torch.clamp(token, min=0)) * mask + # text encode + if finalize is True: + h = self.pre_lookahead_layer(token) + else: + h = self.pre_lookahead_layer( + token[:, : -self.pre_lookahead_len], context=token[:, -self.pre_lookahead_len :] + ) - embedding = self.spk_embed_affine_layer(embedding) + with cosyvoice3_batch_flow_profile("cosyvoice3_cfm_repeat_to_mel_axis"): + h = h.repeat_interleave(self.token_mel_ratio, dim=1) - # concat text and prompt_text - token, token_len = torch.concat([prompt_token, token], dim=1), prompt_token_len + token_len - mask = (~make_pad_mask(token_len)).unsqueeze(-1).to(embedding) - token = self.input_embedding(torch.clamp(token, min=0)) * mask - # text encode - if finalize is True: - h = self.pre_lookahead_layer(token) - else: - h = self.pre_lookahead_layer( - token[:, : -self.pre_lookahead_len], context=token[:, -self.pre_lookahead_len :] - ) - h = h.repeat_interleave(self.token_mel_ratio, dim=1) + batch_size = int(token.shape[0]) + with cosyvoice3_batch_flow_profile("cosyvoice3_cfm_cond_prompt_mel"): + mel_len1, mel_len2 = prompt_feat.shape[1], h.shape[1] - prompt_feat.shape[1] - mel_len1, mel_len2 = prompt_feat.shape[1], h.shape[1] - prompt_feat.shape[1] + # get conditions + conds = torch.zeros([batch_size, mel_len1 + mel_len2, self.output_size], device=token.device).to(h.dtype) + conds[:, :mel_len1] = prompt_feat + conds = conds.transpose(1, 2) - # get conditions - conds = torch.zeros([1, mel_len1 + mel_len2, self.output_size], device=token.device).to(h.dtype) + lookahead = 0 if finalize else int(self.pre_lookahead_len) + valid_h_lens = torch.clamp(total_token_len.to(torch.long) - lookahead, min=0) + mel_lens = torch.clamp(valid_h_lens * int(self.token_mel_ratio), max=mel_len1 + mel_len2) + mask = (~make_pad_mask(mel_lens, max_len=mel_len1 + mel_len2)).to(h) - conds[:, :mel_len1] = prompt_feat - - conds = conds.transpose(1, 2) - - mask = (~make_pad_mask(torch.tensor([mel_len1 + mel_len2]))).to(h) feat, _ = self.decoder( mu=h.transpose(1, 2).contiguous(), mask=mask.unsqueeze(1), @@ -338,6 +363,7 @@ class CausalMaskedDiffWithDiT(torch.nn.Module): streaming=streaming, ) - feat = feat[:, :, mel_len1:] - assert feat.shape[2] == mel_len2 + with cosyvoice3_batch_flow_profile("cosyvoice3_cfm_crop_prompt_mel"): + feat = feat[:, :, mel_len1:] + assert feat.shape[2] == mel_len2 return feat.float(), None diff --git a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py index 558e96b65..0ddc6dafa 100644 --- a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py +++ b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py @@ -39,6 +39,10 @@ from vllm.v1.sample.sampler import Sampler from vllm_omni.data_entry_keys import EmbeddingsStruct, OmniPayloadStruct, to_dict, to_struct from vllm_omni.inputs.mm_processor import OmniMultiModalProcessor +from vllm_omni.model_executor.models.cosyvoice3.runtime import ( + cosyvoice3_batch_flow_debug, + cosyvoice3_batch_flow_enabled, +) from vllm_omni.model_executor.models.cosyvoice3.tokenizer import get_qwen_tokenizer from vllm_omni.model_executor.models.cosyvoice3.utils import ( concat_text_with_prompt_ids, @@ -1070,12 +1074,21 @@ class CosyVoice3Model( request_ids_list = self._split_request_ids(flat_ids, seq_token_counts) num_reqs = max(1, len(request_ids_list)) + debug_batch_flow = cosyvoice3_batch_flow_debug() + if debug_batch_flow: + logger.info( + "CosyVoice3 code2wav debug: forward num_reqs=%d seq_token_counts=%s flat_ids=%d", + num_reqs, + seq_token_counts, + int(flat_ids.numel()), + ) sample_rate = torch.tensor(int(self.config.sample_rate), dtype=torch.int32) empty_audio = torch.zeros((0,), dtype=torch.float32, device=input_ids.device) audios: list[torch.Tensor] = [empty_audio] * num_reqs srs: list[torch.Tensor] = [sample_rate] * num_reqs if not isinstance(runtime_info, list): runtime_info = [] + streaming_flow_items: list[dict[str, object]] = [] for idx, req_ids in enumerate(request_ids_list): raw = runtime_info[idx] if idx < len(runtime_info) and isinstance(runtime_info[idx], dict) else {} @@ -1142,23 +1155,21 @@ class CosyVoice3Model( with self._stream_audio_cache_lock: cache_state = self._stream_vocoder_cache_by_req.get(req_id) - tts_speech, new_cache_state = self.code2wav.forward_streaming( - token=token.unsqueeze(0), - prompt_token=speech_token[:1], - prompt_feat=speech_feat[:1], - embedding=embedding[:1], - cache_state=cache_state, - n_timesteps=10, - token_offset_tokens=token_offset, - finalize=stream_finished, + streaming_flow_items.append( + { + "index": idx, + "req_id": req_id, + "stream_finished": stream_finished, + "token": token.unsqueeze(0), + "prompt_token": speech_token[:1], + "prompt_feat": speech_feat[:1], + "embedding": embedding[:1], + "cache_state": cache_state, + "token_offset_tokens": token_offset, + "finalize": stream_finished, + } ) - - if req_id is not None and hasattr(self, "_stream_vocoder_cache_by_req"): - with self._stream_audio_cache_lock: - if new_cache_state is None or stream_finished: - self._stream_vocoder_cache_by_req.pop(req_id, None) - else: - self._stream_vocoder_cache_by_req[req_id] = new_cache_state + continue else: token_offset = max(0, meta.talker_prefill_offset or 0) if meta else 0 tts_speech = self.code2wav.forward( @@ -1174,6 +1185,60 @@ class CosyVoice3Model( audios[idx] = self._stitch_stream_audio(req_id, audio, stream_finished) + if streaming_flow_items: + if debug_batch_flow: + item_shapes = [ + ( + tuple(item["token"].shape), # type: ignore[union-attr] + tuple(item["prompt_token"].shape), # type: ignore[union-attr] + tuple(item["prompt_feat"].shape), # type: ignore[union-attr] + bool(item.get("finalize", False)), + ) + for item in streaming_flow_items + ] + logger.info( + "CosyVoice3 code2wav debug: streaming_items=%d item_shapes=%s", + len(streaming_flow_items), + item_shapes, + ) + if ( + len(streaming_flow_items) > 1 + and cosyvoice3_batch_flow_enabled() + and hasattr(self.code2wav, "forward_streaming_batch") + ): + streaming_results = self.code2wav.forward_streaming_batch( + streaming_flow_items, + n_timesteps=10, + ) + else: + streaming_results = [ + self.code2wav.forward_streaming( + token=item["token"], # type: ignore[arg-type] + prompt_token=item["prompt_token"], # type: ignore[arg-type] + prompt_feat=item["prompt_feat"], # type: ignore[arg-type] + embedding=item["embedding"], # type: ignore[arg-type] + cache_state=item.get("cache_state"), # type: ignore[arg-type] + n_timesteps=10, + token_offset_tokens=int(item.get("token_offset_tokens", 0)), + finalize=bool(item.get("finalize", False)), + ) + for item in streaming_flow_items + ] + + for item, (tts_speech, new_cache_state) in zip(streaming_flow_items, streaming_results): + idx = int(item["index"]) + req_id = item.get("req_id") + stream_finished = bool(item.get("stream_finished", False)) + if req_id is not None and hasattr(self, "_stream_vocoder_cache_by_req"): + with self._stream_audio_cache_lock: + if new_cache_state is None or stream_finished: + self._stream_vocoder_cache_by_req.pop(req_id, None) + else: + self._stream_vocoder_cache_by_req[req_id] = new_cache_state + + audio = tts_speech.reshape(-1).to(dtype=torch.float32) + audios[idx] = self._stitch_stream_audio(req_id, audio, stream_finished) + return OmniOutput(text_hidden_states=None, multimodal_outputs={"audio": audios, "sr": srs}) else: raise ValueError(f"Unsupported model_stage: {self.model_stage}") diff --git a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py index 49e8e515b..ef706b6c8 100644 --- a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py +++ b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ CosyVoice3 Code2Wav Stage - Converts speech tokens to audio waveforms. @@ -11,6 +11,9 @@ This module contains the code2wav (token-to-waveform) stage which uses: from __future__ import annotations +from collections import Counter +from typing import cast + import numpy as np import torch import torch.nn as nn @@ -27,6 +30,10 @@ from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan import ( CausalHiFTGenerator, ) from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.layers import PreLookaheadLayer +from vllm_omni.model_executor.models.cosyvoice3.runtime import ( + cosyvoice3_batch_flow_debug, + cosyvoice3_batch_flow_profile, +) from vllm_omni.transformers_utils.configs.cosyvoice3 import CosyVoice3Config logger = init_logger(__name__) @@ -159,6 +166,9 @@ class CosyVoice3Code2Wav(nn.Module): token_offset_tokens: int = 0, streaming: bool = True, finalize: bool = False, + token_lens: torch.Tensor | None = None, + prompt_token_lens: torch.Tensor | None = None, + prompt_feat_lens: torch.Tensor | None = None, ) -> torch.Tensor: """Generate mel features via the upstream flow-model inference path.""" flow_weight = next(self.flow_model.parameters()) @@ -169,9 +179,22 @@ class CosyVoice3Code2Wav(nn.Module): prompt_token = prompt_token.to(device=device, dtype=torch.int32) prompt_feat = prompt_feat.to(device=device, dtype=dtype) embedding = embedding.to(device=device, dtype=dtype) - token_len = torch.tensor([token.shape[1]], device=device, dtype=torch.int32) - prompt_token_len = torch.tensor([prompt_token.shape[1]], device=device, dtype=torch.int32) - prompt_feat_len = torch.tensor([prompt_feat.shape[1]], device=device, dtype=torch.int32) + batch_size = int(token.shape[0]) + token_len = ( + token_lens.to(device=device, dtype=torch.int32) + if token_lens is not None + else torch.full((batch_size,), token.shape[1], device=device, dtype=torch.int32) + ) + prompt_token_len = ( + prompt_token_lens.to(device=device, dtype=torch.int32) + if prompt_token_lens is not None + else torch.full((batch_size,), prompt_token.shape[1], device=device, dtype=torch.int32) + ) + prompt_feat_len = ( + prompt_feat_lens.to(device=device, dtype=torch.int32) + if prompt_feat_lens is not None + else torch.full((batch_size,), prompt_feat.shape[1], device=device, dtype=torch.int32) + ) feat, _ = self.flow_model.inference( token=token, @@ -192,37 +215,13 @@ class CosyVoice3Code2Wav(nn.Module): return feat - @torch.inference_mode() - def forward_streaming( + def _stream_hift_from_feat( self, - token: torch.Tensor, - prompt_token: torch.Tensor, - prompt_feat: torch.Tensor, - embedding: torch.Tensor, + feat: torch.Tensor, *, cache_state: dict[str, torch.Tensor] | None = None, - n_timesteps: int = 10, - token_offset_tokens: int = 0, finalize: bool = False, ) -> tuple[torch.Tensor, dict[str, torch.Tensor] | None]: - """Decode streaming audio using cumulative mel + emitted-speech offset. - - This mirrors upstream CosyVoice3 streaming semantics more closely than - waveform-domain overlap-add: keep a cumulative mel history per request, - re-run causal HiFT on the history, and emit only the newly grown speech - suffix. That preserves causal look-right handling without double - trimming or duplicated overlap at chunk boundaries. - """ - feat = self._forward_mel( - token=token, - prompt_token=prompt_token, - prompt_feat=prompt_feat, - embedding=embedding, - n_timesteps=n_timesteps, - token_offset_tokens=token_offset_tokens, - streaming=True, - finalize=finalize, - ) hift_weight = self.hift.m_source.l_linear.weight chunk_mel = feat.to(device=hift_weight.device, dtype=hift_weight.dtype) @@ -257,6 +256,160 @@ class CosyVoice3Code2Wav(nn.Module): } return emitted_speech.reshape(emitted_speech.shape[0], 1, -1), new_state + @torch.inference_mode() + def forward_streaming_batch( + self, + items: list[dict[str, object]], + *, + n_timesteps: int = 10, + ) -> list[tuple[torch.Tensor, dict[str, torch.Tensor] | None]]: + """Batch the flow-matching mel path, then run HiFT per request. + + Items are grouped by prompt condition shape and finalization state. + Codec tokens may have different lengths; those are padded within the + group and passed to the flow as per-row token lengths. + """ + results: list[tuple[torch.Tensor, dict[str, torch.Tensor] | None] | None] = [None] * len(items) + groups: dict[tuple[int, int, int, bool], list[tuple[int, dict[str, object]]]] = {} + for index, item in enumerate(items): + token = item["token"] + prompt_token = item["prompt_token"] + prompt_feat = item["prompt_feat"] + embedding = item["embedding"] + assert isinstance(token, torch.Tensor) + assert isinstance(prompt_token, torch.Tensor) + assert isinstance(prompt_feat, torch.Tensor) + assert isinstance(embedding, torch.Tensor) + key = ( + int(prompt_token.shape[1]), + int(prompt_feat.shape[1]), + int(embedding.shape[1]), + bool(item.get("finalize", False)), + ) + groups.setdefault(key, []).append((index, item)) + + if cosyvoice3_batch_flow_debug(): + group_summary = {key: len(group) for key, group in groups.items()} + group_size_distribution = Counter(group_summary.values()) + batchable_items = sum(size for size in group_summary.values() if size > 1) + logger.info( + "CosyVoice3 code2wav debug: forward_streaming_batch items=%d " + "groups=%s group_size_distribution=%s batchable_items=%d", + len(items), + group_summary, + dict(sorted(group_size_distribution.items())), + batchable_items, + ) + + for _key, group in groups.items(): + if len(group) == 1: + index, item = group[0] + result = self.forward_streaming( + token=item["token"], # type: ignore[arg-type] + prompt_token=item["prompt_token"], # type: ignore[arg-type] + prompt_feat=item["prompt_feat"], # type: ignore[arg-type] + embedding=item["embedding"], # type: ignore[arg-type] + cache_state=item.get("cache_state"), # type: ignore[arg-type] + n_timesteps=n_timesteps, + token_offset_tokens=int(item.get("token_offset_tokens", 0)), + finalize=bool(item.get("finalize", False)), + ) + results[index] = result + continue + + token_tensors = [item["token"] for _, item in group] + assert all(isinstance(token, torch.Tensor) for token in token_tensors) + token_lens = torch.tensor( + [int(token.shape[1]) for token in token_tensors], # type: ignore[union-attr] + dtype=torch.int32, + ) + max_token_len = int(token_lens.max().item()) + padded_tokens = [] + for token in token_tensors: + assert isinstance(token, torch.Tensor) + if int(token.shape[1]) == max_token_len: + padded_tokens.append(token) + else: + pad = torch.zeros( + (token.shape[0], max_token_len - int(token.shape[1])), + device=token.device, + dtype=token.dtype, + ) + padded_tokens.append(torch.cat([token, pad], dim=1)) + tokens = torch.cat(padded_tokens, dim=0) + prompt_tokens = torch.cat([item["prompt_token"] for _, item in group], dim=0) # type: ignore[list-item] + prompt_feats = torch.cat([item["prompt_feat"] for _, item in group], dim=0) # type: ignore[list-item] + embeddings = torch.cat([item["embedding"] for _, item in group], dim=0) # type: ignore[list-item] + prompt_token_lens = torch.full((len(group),), prompt_tokens.shape[1], dtype=torch.int32) + prompt_feat_lens = torch.full((len(group),), prompt_feats.shape[1], dtype=torch.int32) + finalize = bool(group[0][1].get("finalize", False)) + + with cosyvoice3_batch_flow_profile(f"cosyvoice3_flow_batch_b{len(group)}_t{tokens.shape[1]}"): + feat = self._forward_mel( + token=tokens, + prompt_token=prompt_tokens, + prompt_feat=prompt_feats, + embedding=embeddings, + n_timesteps=n_timesteps, + token_offset_tokens=0, + streaming=True, + finalize=finalize, + token_lens=token_lens, + prompt_token_lens=prompt_token_lens, + prompt_feat_lens=prompt_feat_lens, + ) + + for row, (index, item) in enumerate(group): + trim_mel = max(0, int(item.get("token_offset_tokens", 0))) * int(self.token_mel_ratio) + valid_tokens = int(token_lens[row].item()) + if not finalize: + valid_tokens = max(0, valid_tokens - int(self.flow_model.pre_lookahead_len)) + valid_mel = valid_tokens * int(self.token_mel_ratio) + row_feat = feat[row : row + 1, :, :valid_mel] + if trim_mel > 0: + row_feat = row_feat[:, :, trim_mel:] + results[index] = self._stream_hift_from_feat( + row_feat, + cache_state=item.get("cache_state"), # type: ignore[arg-type] + finalize=finalize, + ) + + assert all(result is not None for result in results), "every streaming item must produce exactly one result" + return cast(list[tuple[torch.Tensor, dict[str, torch.Tensor] | None]], results) + + @torch.inference_mode() + def forward_streaming( + self, + token: torch.Tensor, + prompt_token: torch.Tensor, + prompt_feat: torch.Tensor, + embedding: torch.Tensor, + *, + cache_state: dict[str, torch.Tensor] | None = None, + n_timesteps: int = 10, + token_offset_tokens: int = 0, + finalize: bool = False, + ) -> tuple[torch.Tensor, dict[str, torch.Tensor] | None]: + """Decode streaming audio using cumulative mel + emitted-speech offset. + + This mirrors upstream CosyVoice3 streaming semantics more closely than + waveform-domain overlap-add: keep a cumulative mel history per request, + re-run causal HiFT on the history, and emit only the newly grown speech + suffix. That preserves causal look-right handling without double + trimming or duplicated overlap at chunk boundaries. + """ + feat = self._forward_mel( + token=token, + prompt_token=prompt_token, + prompt_feat=prompt_feat, + embedding=embedding, + n_timesteps=n_timesteps, + token_offset_tokens=token_offset_tokens, + streaming=True, + finalize=finalize, + ) + return self._stream_hift_from_feat(feat, cache_state=cache_state, finalize=finalize) + @torch.inference_mode() def forward( self, diff --git a/vllm_omni/model_executor/models/cosyvoice3/runtime.py b/vllm_omni/model_executor/models/cosyvoice3/runtime.py new file mode 100644 index 000000000..d94725148 --- /dev/null +++ b/vllm_omni/model_executor/models/cosyvoice3/runtime.py @@ -0,0 +1,31 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Runtime toggles shared by the CosyVoice3 stages.""" + +import os +from contextlib import nullcontext + +import torch + +_FALSE_ENV_VALUES = ("0", "false", "False", "") + + +def _env_flag(name: str, default: str) -> bool: + return os.environ.get(name, default) not in _FALSE_ENV_VALUES + + +def cosyvoice3_batch_flow_enabled() -> bool: + """Return whether cross-request Stage-1 flow batching is enabled.""" + return _env_flag("COSYVOICE3_BATCH_FLOW", "0") + + +def cosyvoice3_batch_flow_debug() -> bool: + """Return whether Stage-1 batching diagnostics are enabled.""" + return _env_flag("COSYVOICE3_BATCH_FLOW_DEBUG", "0") + + +def cosyvoice3_batch_flow_profile(name: str): + """Create a profiler scope only when batching diagnostics are enabled.""" + if cosyvoice3_batch_flow_debug(): + return torch.profiler.record_function(name) + return nullcontext()