diff --git a/tests/model_executor/models/qwen3_tts/test_qwen3_tts_prompt_text_batch.py b/tests/model_executor/models/qwen3_tts/test_qwen3_tts_prompt_text_batch.py new file mode 100644 index 000000000..6b58a13a4 --- /dev/null +++ b/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 diff --git a/tests/worker_v2/test_omni_model_state.py b/tests/worker_v2/test_omni_model_state.py index 86764b691..05adbf413 100644 --- a/tests/worker_v2/test_omni_model_state.py +++ b/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) diff --git a/vllm_omni/model_executor/models/qwen3_tts/prompt_embeds_builder.py b/vllm_omni/model_executor/models/qwen3_tts/prompt_embeds_builder.py index 7a63a3072..0a8eb4003 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/prompt_embeds_builder.py +++ b/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])), diff --git a/vllm_omni/worker_v2/model_states/eager_mtp.py b/vllm_omni/worker_v2/model_states/eager_mtp.py index 03fb8c0f1..220f40540 100644 --- a/vllm_omni/worker_v2/model_states/eager_mtp.py +++ b/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)) diff --git a/vllm_omni/worker_v2/model_states/omni_model_state.py b/vllm_omni/worker_v2/model_states/omni_model_state.py index 34c3c1cdf..80eb56b83 100644 --- a/vllm_omni/worker_v2/model_states/omni_model_state.py +++ b/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]