Optimize CosyVoice3 Stage1 flow batching (#4876)

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

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