From a1cc583665a9f75743f6a9ae761bd96cc5ddb897 Mon Sep 17 00:00:00 2001
From: Yuekai Zhang <zhangyuekai@foxmail.com>
Date: Thu, 4 Jun 2026 16:58:53 +0800
Subject: [PATCH] fix: cosyvoice3 batch>1 inference (#3910)

Signed-off-by: Yuekai Zhang <zhangyuekai@foxmail.com>
---
 .../models/cosyvoice3/cosyvoice3.py           | 66 +++++++++++++++++--
 1 file changed, 61 insertions(+), 5 deletions(-)

diff --git a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py
index 02326f456..d20312a97 100644
--- a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py
+++ b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py
@@ -643,14 +643,70 @@ class CosyVoice3Model(
     ) -> torch.Tensor:
         if self.model_stage == "cosyvoice3_talker":
             if is_multimodal is not None and any(is_multimodal):
+                # Per-request rearrange for the talker prompt layout
+                # [SOS_ph, TASK_ID_ph, AUDIO_ph * pstoken_len, ...text_tokens...]
+                # → [SOS_emb, ...text_embs..., TASK_ID_emb, AUDIO_embs].
+                #
+                # In batched prefill ``input_ids`` is the flat concatenation of
+                # N sequences and ``multimodal_embeddings`` is a list with one
+                # tensor per multimodal request. Each request's audio
+                # placeholders form one contiguous ``True`` group in
+                # ``is_multimodal``; the 2 positions immediately before each
+                # group are the SOS / TASK_ID placeholders. Walking the groups
+                # lets us locate every request's segment without needing a
+                # ``seq_token_counts`` kwarg (which only the generation runner
+                # injects, not the AR runner used here).
                 embed_tokens = self.model.llm.model.embed_tokens(input_ids)
                 sos = self.model.speech_embedding.weight[self.model.sos].reshape(1, -1)
                 task_id = self.model.speech_embedding.weight[self.model.task_id].reshape(1, -1)
-                prompt_speech_token_emb = multimodal_embeddings[0]
-                pstoken_len = prompt_speech_token_emb.shape[0]  # Get length from tensor shape
-                embed_tokens = torch.cat(
-                    [sos, embed_tokens[2 + pstoken_len :], task_id, prompt_speech_token_emb], dim=0
-                )
+
+                is_mm = is_multimodal.to(torch.bool).reshape(-1)
+                # vLLM v1's multimodal placeholder range covers the FULL
+                # ``[SOS_ph, TASK_ID_ph, AUDIO_ph * pstoken_len]`` block
+                # (length ``2 + pstoken_len``, see ``_get_prompt_updates``'s
+                # ``insertion_end`` which inserts ``[1] * (1 + 1 + token_len)``
+                # and ``mm_position.length`` matches that). So each contiguous
+                # ``True`` group in ``is_mm`` starts at a request's SOS
+                # placeholder, NOT at its first audio placeholder.
+                prev_false = torch.cat([torch.ones(1, dtype=torch.bool, device=is_mm.device), ~is_mm[:-1]])
+                group_starts = (is_mm & prev_false).nonzero(as_tuple=True)[0].tolist()
+                if len(group_starts) != len(multimodal_embeddings):
+                    raise RuntimeError(
+                        f"cosyvoice3 talker: found {len(group_starts)} placeholder "
+                        f"blocks in is_multimodal but {len(multimodal_embeddings)} "
+                        f"multimodal_embeddings tensors — these must match 1:1."
+                    )
+
+                total_tokens = int(embed_tokens.shape[0])
+                # vLLM v1 packs cached (decode) requests first, then new
+                # (prefill) requests. The decode tokens at positions
+                # ``[0:group_starts[0]]`` get a speech_embedding lookup (the
+                # talker is generating codec tokens here, so input ids are
+                # codec ids, not text ids). Positions starting at
+                # ``group_starts[0]`` are the first new prefill request's
+                # SOS placeholder.
+                segments: list[torch.Tensor] = []
+                first_prefill = group_starts[0]
+                if first_prefill > 0:
+                    decode_ids = input_ids[:first_prefill].to(dtype=torch.long)
+                    segments.append(self.model.speech_embedding.weight[decode_ids])
+
+                for i, req_start in enumerate(group_starts):
+                    pstoken_len_i = int(multimodal_embeddings[i].shape[0])
+                    # Real text starts after the 2 + pstoken_len_i leading
+                    # placeholders of this request's prompt.
+                    text_start = req_start + 2 + pstoken_len_i
+                    # Next request's segment starts at its own SOS placeholder,
+                    # which is exactly the next True group's first index.
+                    if i + 1 < len(group_starts):
+                        req_end = group_starts[i + 1]
+                    else:
+                        req_end = total_tokens
+                    segments.append(sos)
+                    segments.append(embed_tokens[text_start:req_end])
+                    segments.append(task_id)
+                    segments.append(multimodal_embeddings[i])
+                embed_tokens = torch.cat(segments, dim=0)
             else:
                 embed_tokens = self.model.speech_embedding.weight[input_ids]
             return embed_tokens