Optimize CosyVoice3 Stage1 flow batching (#4876)

Signed-off-by: gerayking <399geray@gmail.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

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:

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,
     }
 

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

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])

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",

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

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

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}")

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,

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()