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

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

tests/model_executor/models/qwen3_tts/test_qwen3_tts_prompt_text_batch.py

@@ -0,0 +1,140 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from typing import Any
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import (
+    Qwen3TTSConfig,
+    Qwen3TTSTalkerConfig,
+)
+from vllm_omni.model_executor.models.qwen3_tts.prompt_embeds_builder import (
+    PRECOMPUTED_TEXT_IDS_KEY,
+    Qwen3TTSPromptEmbedsBuilder,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_batch_preprocess_projects_new_non_streaming_texts_once():
+    torch.manual_seed(0)
+    text_embedding = torch.nn.Embedding(64, 4)
+    projection = torch.nn.Linear(4, 4)
+    projected_shapes: list[tuple[int, ...]] = []
+
+    def text_projection(x: torch.Tensor) -> torch.Tensor:
+        projected_shapes.append(tuple(x.shape))
+        return projection(x)
+
+    builder = Qwen3TTSPromptEmbedsBuilder.__new__(Qwen3TTSPromptEmbedsBuilder)
+    builder._text_embedding = text_embedding
+    builder._text_projection = text_projection
+    builder._batched_text_embeds = {}
+
+    ids_a = list(range(1, 13))
+    ids_b = list(range(20, 31))
+    # Serving stores the assistant-template ids wrapped in a one-element list.
+    serving_ids = [ids_a]
+    buf: dict[str, dict[str, Any]] = {
+        "a": {"req_id": "a", "task_type": ["CustomVoice"], PRECOMPUTED_TEXT_IDS_KEY: [ids_a]},
+        "b": {"req_id": "b", "task_type": ["VoiceDesign"], PRECOMPUTED_TEXT_IDS_KEY: [torch.tensor(ids_b)]},
+        "streaming": {
+            "req_id": "streaming",
+            "task_type": ["CustomVoice"],
+            "non_streaming_mode": [False],
+            PRECOMPUTED_TEXT_IDS_KEY: serving_ids,
+        },
+        "base": {"req_id": "base", "task_type": ["Base"], PRECOMPUTED_TEXT_IDS_KEY: serving_ids},
+        "built": {
+            "req_id": "built",
+            "task_type": ["CustomVoice"],
+            "embed": {"prefill": torch.zeros(1)},
+            PRECOMPUTED_TEXT_IDS_KEY: serving_ids,
+        },
+    }
+
+    builder.preprocess_infos_batch(req_infos=list(buf.values()), device=torch.device("cpu"))
+
+    # One embedding + projection over both requests' text tokens (template stripped).
+    assert projected_shapes == [(1, (len(ids_a) - 8) + (len(ids_b) - 8), 4)]
+    assert set(builder._batched_text_embeds) == {"a", "b"}
+    with torch.no_grad():
+        for req_id, ids in (("a", ids_a), ("b", ids_b)):
+            expected = projection(text_embedding(torch.tensor([ids[3:-5]])))
+            torch.testing.assert_close(builder._batched_text_embeds[req_id], expected)
+            assert buf[req_id][PRECOMPUTED_TEXT_IDS_KEY].tolist() == [ids]
+    for skipped in ("streaming", "base", "built"):
+        assert buf[skipped][PRECOMPUTED_TEXT_IDS_KEY] is serving_ids
+
+
+class _Tokenizer:
+    def __call__(self, text: str, **kwargs: Any) -> dict[str, torch.Tensor]:
+        return {"input_ids": torch.tensor([[2, 3, 4, 5]])}
+
+
+def _prompt_builder() -> Qwen3TTSPromptEmbedsBuilder:
+    """Use the real constructor and prompt assembly with small embedding tables."""
+    torch.manual_seed(42)
+    config = Qwen3TTSConfig(tts_bos_token_id=50, tts_eos_token_id=51, tts_pad_token_id=52)
+    talker_config = Qwen3TTSTalkerConfig(
+        codec_nothink_id=10,
+        codec_think_id=11,
+        codec_think_bos_id=12,
+        codec_think_eos_id=13,
+        codec_pad_id=14,
+        codec_bos_id=15,
+        codec_language_id={"english": 16},
+        spk_id={"vivian": 17, "serena": 18},
+    )
+    builder = Qwen3TTSPromptEmbedsBuilder(
+        config=config,
+        talker_config=talker_config,
+        model_path="",
+        text_embedding=torch.nn.Embedding(64, 4),
+        text_projection=torch.nn.Linear(4, 4),
+        codec_embed=torch.nn.Embedding(64, 4),
+        residual_code_embeddings=lambda: [],
+        speaker_encoder=torch.nn.Identity(),
+        tts_pad_embed=torch.zeros(1, 4),
+        encode_ref_audio_batch=lambda *args, **kwargs: [],
+    )
+    builder._text_tokenizer = _Tokenizer()
+    return builder
+
+
+@pytest.mark.parametrize("task_type", ["CustomVoice", "VoiceDesign"])
+@pytest.mark.parametrize("non_streaming", [False, True])
+def test_batched_prompt_consumption_matches_serial_and_clears_cache(task_type: str, non_streaming: bool):
+    batched, serial = _prompt_builder(), _prompt_builder()
+
+    def request(req_id: str, ids: list[int], speaker: str) -> dict[str, Any]:
+        return {
+            "req_id": req_id,
+            "text": ["hello"],
+            "task_type": [task_type],
+            "language": ["English"],
+            "speaker": [speaker],
+            "instruct": ["Speak calmly"],
+            "non_streaming_mode": [non_streaming],
+            PRECOMPUTED_TEXT_IDS_KEY: [ids],
+        }
+
+    with torch.inference_mode():
+        # The same request id can be used again after its earlier prompt was consumed.
+        for ids_a, ids_b in [(list(range(12)), list(range(20, 31))), (list(range(10, 24)), list(range(30, 46)))]:
+            pending = [request("a", ids_a, "Vivian"), request("b", ids_b, "Serena")]
+            expected = [request("a", ids_a, "Vivian"), request("b", ids_b, "Serena")]
+            batched.preprocess_infos_batch(req_infos=pending, device=torch.device("cpu"))
+            assert set(batched._batched_text_embeds) == ({"a", "b"} if non_streaming else set())
+            for actual_info, expected_info in zip(pending, expected, strict=True):
+                actual = batched.build_prompt_embeds(task_type=task_type, info_dict=actual_info)
+                reference = serial.build_prompt_embeds(task_type=task_type, info_dict=expected_info)
+                torch.testing.assert_close(actual[0], reference[0])
+                torch.testing.assert_close(actual[1], reference[1])
+                assert actual[2:] == reference[2:]
+                assert PRECOMPUTED_TEXT_IDS_KEY not in actual_info
+            assert not batched._batched_text_embeds
+        batched.preprocess_infos_batch(req_infos=[], device=torch.device("cpu"))
+        assert not batched._batched_text_embeds

tests/worker_v2/test_omni_model_state.py

@@ -24,6 +24,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 class _DummyInputBatch:
     is_prefilling_np: np.ndarray
+    req_ids: list[str]
 
     input_ids: SimpleNamespace
     query_start_loc: torch.Tensor
@@ -701,6 +702,58 @@ def test_publish_sampled_embeddings_is_opt_in() -> None:
     assert state.publish_sampled_embeddings(SimpleNamespace(num_reqs=1), torch.tensor([[1]])) is None
 
 
+def test_split_settled_rows_matches_per_row_check():
+    batch = _DummyInputBatch(np.array([3, 0, 2, 1, 4]))
+    batch.req_ids = ["c", "a", "x", "b", "e"]
+    batch.num_scheduled_tokens = [1, 1, 1, 4, 1]
+    batch.query_start_loc_np = [0, 1, 2, 3, 7]
+    req_indices = batch.idx_mapping_np.tolist()
+    # Slot 2 still names a finished request, slot 1 schedules a prefill
+    # chunk, slot 4 is not settled and slot 7 is not scheduled.
+    settled = {3: "c", 0: "a", 2: "old", 1: "b", 7: "z"}
+    split = OmniModelState._split_settled_rows(batch, req_indices, settled)
+    assert split is not None
+    settled_rows, remaining = split
+    assert settled_rows == [(0, 3, 0, "c"), (1, 0, 1, "a")]
+    assert remaining == [(2, 2), (3, 1), (4, 4)]
+    assert OmniModelState._split_settled_rows(batch, req_indices, {2: "old"}) is None
+
+
+def test_settled_row_split_preserves_mixed_batch_and_pending_replay(monkeypatch):
+    """Exercise run_preprocess, including the scalar fallback during replay."""
+
+    def prepare():
+        state = _make_eager_state()
+        _fill_buffers(state, "prefill", "decode")
+        state._eager_fastpath = True
+        state._eager_settled = {1: "decode"}
+        state._eager_ready = {1: "decode"}
+        state._eager_embeds[1].fill_(7)
+        state.model.eager_settled_text_step = lambda: torch.ones(_EAGER_DIM)
+        state.model.eager_decode_settled = lambda info: False
+        state.model.preprocess = lambda input_ids, input_embeds, **info: (input_ids, input_embeds, {})
+        batch = _EagerBatch([3, 1])
+        batch.req_ids = ["prefill", "decode"]
+        inputs = {"input_ids": torch.zeros(4, dtype=torch.long), "inputs_embeds": torch.zeros(4, _EAGER_DIM)}
+        return state, batch, inputs
+
+    vector, batch, inputs = prepare()
+    vector.run_preprocess(batch, inputs)
+    reference, ref_batch, ref_inputs = prepare()
+    monkeypatch.setattr(reference, "_split_settled_rows", lambda *args: None)
+    reference.run_preprocess(ref_batch, ref_inputs)
+    torch.testing.assert_close(inputs["inputs_embeds"], ref_inputs["inputs_embeds"], rtol=0, atol=0)
+    assert vector._eager_rows[1] == reference._eager_rows[1]
+    assert inputs["inputs_embeds"][3].tolist() == [8.0] * _EAGER_DIM
+
+    replay, replay_batch, replay_inputs = prepare()
+    replay._eager_state._restore_audio["prefill"] = [torch.ones(1)]
+    assert replay._eager_state.has_pending_replay()
+    monkeypatch.setattr(replay, "_split_settled_rows", MagicMock(side_effect=AssertionError("replay must skip split")))
+    replay.run_preprocess(replay_batch, replay_inputs)
+    torch.testing.assert_close(replay_inputs["inputs_embeds"], ref_inputs["inputs_embeds"], rtol=0, atol=0)
+
+
 @pytest.mark.parametrize("prefilling", [False, True])
 def test_identity_preprocess_skips_only_decode_rows(prefilling):
     state = _make_state(has_preprocess=True)

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]