[Model] Batch Qwen3-TTS Talker prompt projection and settled-row preprocessing (#8477)

Signed-off-by: Sy03 <1370724210@qq.com>

vllm_omni/model_executor/models/qwen3_tts/prompt_embeds_builder.py

@@ -370,6 +370,11 @@ class Qwen3TTSPromptEmbedsBuilder:
         # CustomVoice prompt pieces that do not depend on the text, per
         # (device, speaker, language): bounded by the model's voices.
         self._custom_voice_prompt_cache: dict[tuple, tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = {}
+        # Speaker codec id and embedding per (device, speaker).
+        self._speaker_embed_cache: dict[tuple[str, str], tuple[int, torch.Tensor]] = {}
+        # Text-conditioning embeddings projected by ``preprocess_infos_batch``
+        # for the current step's new requests, keyed by request id.
+        self._batched_text_embeds: dict[str, torch.Tensor] = {}
 
         self._ref_audio_artifact_cache_max_entries = int(ref_audio_artifact_cache_max_entries)
         self._ref_audio_artifact_cache: OrderedDict[str, dict[str, torch.Tensor | bool]] = OrderedDict()
@@ -752,6 +757,18 @@ class Qwen3TTSPromptEmbedsBuilder:
             return bool(value)
         return False
 
+    @staticmethod
+    def _non_streaming_mode(info_dict: dict[str, Any], task_type: str) -> bool:
+        value = info_dict.get("non_streaming_mode")
+        if isinstance(value, list):
+            value = value[0] if value else None
+        if isinstance(value, bool):
+            return value
+        # Match official inference defaults:
+        # - CustomVoice/VoiceDesign: non_streaming_mode=True
+        # - Base: non_streaming_mode=False
+        return task_type in ("CustomVoice", "VoiceDesign")
+
     @staticmethod
     def _needs_initial_prompt_preprocess(info_dict: dict[str, Any]) -> bool:
         embed = info_dict.get("embed")
@@ -806,15 +823,16 @@ class Qwen3TTSPromptEmbedsBuilder:
         req_infos: list[dict[str, Any]],
         device: torch.device,
     ) -> None:
-        """Batch Base voice-clone ref-audio codec extraction for current prefill requests.
+        """Batch per-request prompt work for current prefill requests.
 
-        Tokenizes assistant / ref text and runs the talker-supplied
-        ``encode_ref_audio_batch`` callable once per sample-rate group,
-        stashing the precomputed artifacts under
+        For Base voice cloning, tokenizes assistant / ref text and runs the
+        talker-supplied ``encode_ref_audio_batch`` callable once per
+        sample-rate group, stashing the precomputed artifacts under
         :data:`PRECOMPUTED_TEXT_IDS_KEY`, :data:`PRECOMPUTED_REF_IDS_KEY`,
-        :data:`PRECOMPUTED_REF_CODE_KEY`, and :data:`NORMALIZED_REF_AUDIO_KEY`
-        so the per-request :meth:`build_prompt_embeds` call can skip the
-        equivalent serial work.
+        :data:`PRECOMPUTED_REF_CODE_KEY`, and :data:`NORMALIZED_REF_AUDIO_KEY`.
+        For non-streaming CustomVoice / VoiceDesign, projects every request's
+        text at once (:meth:`_project_texts_batch`). The per-request
+        :meth:`build_prompt_embeds` call then skips the equivalent serial work.
         """
         pending_text: list[tuple[dict[str, Any], str]] = []
         pending_ref_text: list[tuple[dict[str, Any], str]] = []
@@ -919,6 +937,54 @@ class Qwen3TTSPromptEmbedsBuilder:
                 if isinstance(cache_key, str) and cache_key:
                     self.put_ref_audio_artifacts(cache_key, ref_code=ref_code_t)
 
+        self._batched_text_embeds = self._project_texts_batch(req_infos)
+
+    def _project_texts_batch(self, req_infos: list[dict[str, Any]]) -> dict[str, torch.Tensor]:
+        """Embed and project the text of new non-streaming CustomVoice / VoiceDesign requests at once.
+
+        Each request otherwise uploads its text ids and runs the text
+        embedding + projection MLP on its own. Here every request's ids go up
+        in one pinned host-to-device copy and the MLP runs once over the
+        concatenated text tokens; the rows are independent, so each request's
+        slice equals its own per-request result. The device ids replace
+        :data:`PRECOMPUTED_TEXT_IDS_KEY` so ``build_prompt_embeds`` skips the
+        per-request copy too.
+        """
+        items: list[tuple[dict[str, Any], str, np.ndarray]] = []
+        for info_dict in req_infos:
+            req_id = info_dict.get("req_id")
+            if not isinstance(req_id, str) or not self._needs_initial_prompt_preprocess(info_dict):
+                continue
+            task_type = first_value(info_dict.get("task_type"), "CustomVoice")
+            if task_type not in ("CustomVoice", "VoiceDesign") or not self._non_streaming_mode(info_dict, task_type):
+                continue
+            ids = first_value(info_dict.get(PRECOMPUTED_TEXT_IDS_KEY))
+            if isinstance(ids, torch.Tensor) and ids.device.type == "cpu":
+                ids = ids.reshape(-1).numpy()
+            elif isinstance(ids, list) and ids and all(isinstance(v, (int, np.integer)) for v in ids):
+                ids = np.asarray(ids)
+            elif not (isinstance(ids, np.ndarray) and ids.ndim == 1):
+                continue
+            # The assistant template wraps the text in 3 leading and 5 trailing tokens.
+            if ids.size <= 8:
+                continue
+            items.append((info_dict, req_id, ids.astype(np.int64, copy=False)))
+        if not items:
+            return {}
+        full_ids = [ids for _info, _req_id, ids in items]
+        text_ids = [ids[3:-5] for ids in full_ids]
+        packed = to_device_nonblocking(torch.from_numpy(np.concatenate(full_ids + text_ids)), self._device())
+        num_full = sum(ids.size for ids in full_ids)
+        projected = self._text_projection(self._text_embedding(packed[num_full:].view(1, -1)))
+        embeds: dict[str, torch.Tensor] = {}
+        full_offset = text_offset = 0
+        for (info_dict, req_id, ids), text in zip(items, text_ids, strict=True):
+            info_dict[PRECOMPUTED_TEXT_IDS_KEY] = packed[full_offset : full_offset + ids.size].view(1, -1)
+            embeds[req_id] = projected[:, text_offset : text_offset + text.size]
+            full_offset += ids.size
+            text_offset += text.size
+        return embeds
+
     # -------------------- ICL prompt assembly --------------------
 
     def _generate_icl_prompt(
@@ -1010,18 +1076,7 @@ class Qwen3TTSPromptEmbedsBuilder:
 
         text = (info_dict.get("text") or [""])[0]
         language = (info_dict.get("language") or ["Auto"])[0]
-        non_streaming_mode_val = info_dict.get("non_streaming_mode")
-        if isinstance(non_streaming_mode_val, list):
-            non_streaming_mode_raw = non_streaming_mode_val[0] if non_streaming_mode_val else None
-        else:
-            non_streaming_mode_raw = non_streaming_mode_val
-        if isinstance(non_streaming_mode_raw, bool):
-            non_streaming_mode = non_streaming_mode_raw
-        else:
-            # Match official inference defaults:
-            # - CustomVoice/VoiceDesign: non_streaming_mode=True
-            # - Base: non_streaming_mode=False
-            non_streaming_mode = task_type in ("CustomVoice", "VoiceDesign")
+        non_streaming_mode = self._non_streaming_mode(info_dict, task_type)
 
         # Text ids for assistant template (always).
         tok = self.get_text_tokenizer()
@@ -1412,15 +1467,19 @@ class Qwen3TTSPromptEmbedsBuilder:
             )
             if not speaker:
                 raise ValueError("CustomVoice requires additional_information.speaker.")
-            spk_id_map = {k.lower(): v for k, v in (getattr(talker_config, "spk_id", None) or {}).items()}
-            if speaker not in spk_id_map:
-                raise ValueError(f"Unsupported speaker: {speaker}")
-            spk_id = spk_id_map[speaker]
-            # Keep it at least 1D; embedding on a 0-d tensor can return 1D.
-            spk_tensor = index_to_device([spk_id], input_ids.device)
-            spk_embed = codec_embed(spk_tensor)
-            if spk_embed.ndim in (1, 2):
-                spk_embed = spk_embed.view(1, 1, -1)
+            speaker_key = (str(input_ids.device), speaker)
+            cached_speaker = self._speaker_embed_cache.get(speaker_key)
+            if cached_speaker is None:
+                spk_id_map = {k.lower(): v for k, v in (getattr(talker_config, "spk_id", None) or {}).items()}
+                if speaker not in spk_id_map:
+                    raise ValueError(f"Unsupported speaker: {speaker}")
+                spk_id = spk_id_map[speaker]
+                # Keep it at least 1D; embedding on a 0-d tensor can return 1D.
+                spk_embed = codec_embed(index_to_device([spk_id], input_ids.device))
+                if spk_embed.ndim in (1, 2):
+                    spk_embed = spk_embed.view(1, 1, -1)
+                cached_speaker = self._speaker_embed_cache[speaker_key] = (spk_id, spk_embed)
+            spk_id, spk_embed = cached_speaker
             speaker_embed = spk_embed
             codec_input = torch.cat([codec_input_0, speaker_embed, codec_input_1], dim=1)
 
@@ -1452,7 +1511,9 @@ class Qwen3TTSPromptEmbedsBuilder:
             talker_prompt = prefix
 
             if non_streaming_mode:
-                text_all = text_projection(text_embedding(input_ids[:, 3:-5]))
+                text_all = self._batched_text_embeds.pop(str(info_dict.get("req_id")), None)
+                if text_all is None:
+                    text_all = text_projection(text_embedding(input_ids[:, 3:-5]))
                 text_all = torch.cat([text_all, tts_eos_embed], dim=1)
                 talker_prompt = torch.cat([talker_prompt, text_all + pad_row, tail], dim=1)
                 trailing_text_hidden = tts_pad_embed
@@ -1477,7 +1538,9 @@ class Qwen3TTSPromptEmbedsBuilder:
             talker_prompt = torch.cat((role_embed, codec_prefix), dim=1)
 
             if non_streaming_mode:
-                text_all = text_projection(text_embedding(input_ids[:, 3:-5]))
+                text_all = self._batched_text_embeds.pop(str(info_dict.get("req_id")), None)
+                if text_all is None:
+                    text_all = text_projection(text_embedding(input_ids[:, 3:-5]))
                 text_all = torch.cat([text_all, tts_eos_embed], dim=1)
                 pad_ids = torch.full(
                     (1, int(text_all.shape[1])),

vllm_omni/worker_v2/model_states/eager_mtp.py

@@ -111,6 +111,10 @@ class EagerMTPState:
         if self._audio_buffer is not None:
             self._audio_buffer.finish(req_ids)
 
+    def has_pending_replay(self) -> bool:
+        """Whether any preempted request still has saved inputs to replay."""
+        return bool(self._restore_audio)
+
     def replay_inputs(self, req_id: str, req_idx: int, offset: int, ids: torch.Tensor, embeds: torch.Tensor) -> bool:
         """Rebuild freed Talker KV from the exact conditioned inputs, without rerunning MTP."""
         if req_id not in self._restore_audio:
@@ -139,12 +143,15 @@ class EagerMTPState:
             return
         history_slices: list[torch.Tensor] = []
         input_slices: list[torch.Tensor] = []
+        histories = self._talker_inputs
+        on_cpu = embeds.is_cpu
         for i, req_id in enumerate(input_batch.req_ids):
             start, end = input_batch.query_start_loc_np[i : i + 2]
             offset = int(input_batch.num_computed_tokens_np[i])
             length = offset + int(end - start)
-            history = self._talker_inputs.get(req_id)
-            if history is None or history.embeds.device != embeds.device:
+            history = histories.get(req_id)
+            # A suspended request's inputs wait on the host; bring them back.
+            if history is None or history.embeds.is_cpu != on_cpu:
                 storage = torch.empty(
                     self.owner.vllm_config.model_config.max_model_len,
                     embeds.shape[-1],
@@ -405,13 +412,15 @@ class EagerMTPState:
                         p = int(ref.shape[0])
                 pos_list.append(p)
                 positions[req_id] = p + 1
-        meta_list = (
-            [i for i, _r, _q, _p in entries]
-            + [r for _i, r, _q, _p in entries]
-            + [int(p) for _i, _r, _q, p in entries]
-            + [int(q in fa_requests) for _i, _r, q, _p in entries]
-            + pos_list
-        )
+        # One transpose of the entry tuples instead of a Python pass per field.
+        rows_i, rows_req_idx, rows_req_id, rows_prefill = zip(*entries) if entries else ((), (), (), ())
+        meta_list = [
+            *rows_i,
+            *rows_req_idx,
+            *map(int, rows_prefill),
+            *[int(req_id in fa_requests) for req_id in rows_req_id],
+            *pos_list,
+        ]
         meta = index_to_device(meta_list, device)
         rows = meta[bsz : 2 * bsz]
         sampled = sampled_token_ids.reshape(input_batch.num_reqs, -1)
@@ -430,7 +439,7 @@ class EagerMTPState:
             req_indices, batch_ids, batch_emb, batch_hidden, batch_step, mtp_batch_descriptor_dispatcher
         )
         assert codes is not None
-        has_prefill = any(p for _i, _r, _q, p in entries)
+        has_prefill = any(rows_prefill)
         fa_in_kernel = isinstance(first_audio, torch.Tensor) and not has_prefill
         valid = eager_post(
             meta, bsz, last_tokens, layer0, frame_embeds.reshape(frame_embeds.shape[0], -1).contiguous(),
@@ -438,8 +447,7 @@ class EagerMTPState:
             first_audio if fa_in_kernel else None, self._first_audio_valid if fa_in_kernel else None,
             owner.model._codebook_vocab_size,
         )  # fmt: skip
-        for _i, req_idx, req_id, _prefill in entries:
-            owner._eager_ready[req_idx] = req_id
+        owner._eager_ready.update(zip(rows_req_idx, rows_req_id))
         if pos_list:
             assert stream is not None
             assert isinstance(stream_out, torch.Tensor)
@@ -452,10 +460,11 @@ class EagerMTPState:
             for t in (frame_codes, meta, last_tokens, valid, stream_out):
                 t.record_stream(side)
             with torch.cuda.stream(side):
-                for _i, req_idx, req_id, _p in entries:
-                    saved = self._restore_audio.pop(req_id, None)
-                    if saved is not None:
-                        stream.restore_slot(req_idx, saved)
+                if self._restore_audio:
+                    for _i, req_idx, req_id, _p in entries:
+                        saved = self._restore_audio.pop(req_id, None)
+                        if saved is not None:
+                            stream.restore_slot(req_idx, saved)
                 model.prime_stream_decoder(primes)
                 pcm = decode(frame_codes, rows, meta[4 * bsz : 5 * bsz])
                 stream_out.index_copy_(0, last_tokens, pcm.reshape(bsz, -1).to(stream_out.dtype))

vllm_omni/worker_v2/model_states/omni_model_state.py

@@ -16,7 +16,7 @@ from __future__ import annotations
 import inspect
 import threading
 import types
-from collections.abc import Callable
+from collections.abc import Callable, Iterable
 from dataclasses import replace
 from typing import Any, cast
 
@@ -657,6 +657,44 @@ class OmniModelState(DefaultModelState):
                 else:
                     del updates[key]
 
+    @staticmethod
+    def _split_settled_rows(
+        input_batch: InputBatch,
+        req_indices: list[int],
+        settled: dict[int, str],
+    ) -> tuple[list[tuple[int, int, int, str]], list[tuple[int, int]]] | None:
+        """Vectorized settled-row test of ``run_preprocess``.
+
+        In steady decode almost every row is settled: it schedules one token
+        and ``settled`` maps its slot to its request id. ``settled`` entries
+        are dropped together with the slot buffer, so comparing against the
+        row's request id equals the per-row ``buf["req_id"]`` test. Returns the
+        settled rows and the remaining ``(row, slot)`` pairs, both in row
+        order, or ``None`` when no row is settled.
+        """
+        n = input_batch.num_reqs
+        slots = np.asarray(input_batch.idx_mapping_np[:n])
+        owners = np.full(max(int(slots.max()), max(settled)) + 1, None, dtype=object)
+        owners[np.fromiter(settled.keys(), dtype=np.int64, count=len(settled))] = np.array(
+            list(settled.values()), dtype=object
+        )
+        req_ids = np.empty(n, dtype=object)
+        req_ids[:] = input_batch.req_ids[:n]
+        is_settled = (np.asarray(input_batch.num_scheduled_tokens[:n]) == 1) & (owners[slots] == req_ids)
+        if not is_settled.any():
+            return None
+        rows = np.flatnonzero(is_settled)
+        settled_rows = list(
+            zip(
+                rows.tolist(),
+                slots[rows].tolist(),
+                np.asarray(input_batch.query_start_loc_np[:n])[rows].tolist(),
+                req_ids[rows].tolist(),
+            )
+        )
+        remaining = [(i, req_indices[i]) for i in np.flatnonzero(~is_settled).tolist()]
+        return settled_rows, remaining
+
     def run_preprocess(
         self,
         input_batch: InputBatch,
@@ -694,7 +732,7 @@ class OmniModelState(DefaultModelState):
         mtp_batches: list[tuple[int, int, tuple[torch.Tensor, torch.Tensor]]] = []
         prepacked_mtp_inputs: tuple[torch.Tensor, torch.Tensor] | None = None
 
-        req_indices = [int(input_batch.idx_mapping_np[i]) for i in range(input_batch.num_reqs)]
+        req_indices = np.asarray(input_batch.idx_mapping_np[: input_batch.num_reqs]).tolist()
         self._stage_batched_preprocess_inputs(req_indices, embeds.device)
 
         preprocess_entries: list[tuple[int, int, int, int, dict[str, Any], bool]] = []
@@ -704,7 +742,12 @@ class OmniModelState(DefaultModelState):
         # Declared identity decode has no per-row updates or hook work.
         skip_decode_rows = self._decode_preprocess_is_identity and not self._eager_mtp
         is_prefilling_np = getattr(input_batch, "is_prefilling_np", None) if skip_decode_rows else None
-        for i, req_idx in enumerate(req_indices):
+        rows: Iterable[tuple[int, int]] = enumerate(req_indices)
+        if settled and is_prefilling_np is None and not self._eager_state.has_pending_replay():
+            split = self._split_settled_rows(input_batch, req_indices, settled)
+            if split is not None:
+                settled_rows, rows = split
+        for i, req_idx in rows:
             if is_prefilling_np is not None and not is_prefilling_np[i]:
                 continue
             buf = self.intermediate_buffer.buffers[req_idx]