From a58825ed6cc4ef814e5ec8afeffade476a7de431 Mon Sep 17 00:00:00 2001 From: tzhouam Date: Wed, 9 Sep 2026 23:55:36 +0800 Subject: [PATCH] perf(omnivoice): tokenize the prompt once per request apply() already tokenized the prompt to build the engine prompt ids, but the shared OmniMultiModalProcessor._apply_hf_processor_main then called _call_hf_processor again, which rebuilt the control-token prompt and ran the text tokenizer a second time before encoding the reference audio. Upstream's _cached_apply_hf_processor invokes _apply_hf_processor_main for every request, so this also happened for text-only requests and for requests whose reference audio was already cached. Split the old _call_hf_processor into _encode_prompt_text (used once, in apply) and _encode_reference_audio, and override _apply_hf_processor_main to encode the reference audio only, returning just the passthrough data when there is no audio. _get_mm_fields_config registers only ref_audio_tokens and ref_audio_len, so dropping input_ids from the media result changes no consumer. The now-unreferenced _call_hf_processor is deleted; the prompt-text kwarg transport is left as is. Validation: ruff and pre-commit on both files; the module imports cleanly; new tests/model_executor/models/omnivoice/test_processor_prompt_tokenization.py covers text-only apply (one tokenizer call), the audio path (audio encoded, no re-tokenization) and the empty-items path (no tokenizer, no encoder); tests/model_executor/models/omnivoice (106), tests/inputs/test_mm_processor.py and tests/entrypoints/openai_api/test_tts_{adapter,detection}.py (162) pass. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01MpXe2LeRrUQGjZfKFaocMZ Signed-off-by: tzhouam --- .../test_processor_prompt_tokenization.py | 102 ++++++++++++++++++ .../models/omnivoice/omnivoice.py | 73 ++++++------- 2 files changed, 139 insertions(+), 36 deletions(-) create mode 100644 tests/model_executor/models/omnivoice/test_processor_prompt_tokenization.py diff --git a/tests/model_executor/models/omnivoice/test_processor_prompt_tokenization.py b/tests/model_executor/models/omnivoice/test_processor_prompt_tokenization.py new file mode 100644 index 000000000..9c6533c76 --- /dev/null +++ b/tests/model_executor/models/omnivoice/test_processor_prompt_tokenization.py @@ -0,0 +1,102 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""OmniVoice processor: the prompt is tokenized exactly once per request. + +Upstream ``_cached_apply_hf_processor`` calls ``_apply_hf_processor_main`` for +every request, including text-only requests and requests whose reference audio +is already cached (empty ``mm_items``). These tests pin that only ``apply`` +touches the text tokenizer and that the media path only encodes the audio. +""" + +from types import SimpleNamespace +from unittest.mock import Mock + +import numpy as np +import pytest +import torch +from vllm.multimodal.parse import MultiModalDataItems +from vllm.multimodal.processing import ProcessorInputs + +from vllm_omni.inputs.mm_processor import OmniMultiModalProcessor +from vllm_omni.model_executor.models.omnivoice.omnivoice import OmniVoiceMultiModalProcessor + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +_PROMPT_IDS = [5, 6, 7] + + +def _make_processor(monkeypatch: pytest.MonkeyPatch) -> OmniVoiceMultiModalProcessor: + processor = object.__new__(OmniVoiceMultiModalProcessor) + tokenizer = SimpleNamespace(decode=Mock(return_value="hello world")) + processor.info = SimpleNamespace( + ctx=SimpleNamespace( + get_hf_config=lambda *_args, **_kwargs: SimpleNamespace(sample_rate=24000), + model_config=SimpleNamespace(model="fake-model"), + ), + get_tokenizer=lambda: tokenizer, + ) + processor.text_tokenizer = Mock(return_value=SimpleNamespace(input_ids=torch.tensor([_PROMPT_IDS]))) + processor.audio_tokenizer = SimpleNamespace(encode=Mock(return_value=torch.zeros(8, 4, dtype=torch.long))) + processor.feature_extractor = None + monkeypatch.setattr(processor, "_ensure_cached_runtime_components", lambda *_args, **_kwargs: None) + return processor + + +def _items(counts: dict[str, int], mm_data: dict[str, object]) -> SimpleNamespace: + return SimpleNamespace( + get_all_counts=lambda: counts, + select=lambda _keys: mm_data, + ) + + +def test_apply_tokenizes_the_prompt_once(monkeypatch: pytest.MonkeyPatch) -> None: + processor = _make_processor(monkeypatch) + captured: dict[str, ProcessorInputs] = {} + + def fake_parent_apply(self, inputs, timing_ctx): + captured["inputs"] = inputs + return "sentinel" + + monkeypatch.setattr(OmniMultiModalProcessor, "apply", fake_parent_apply) + inputs = ProcessorInputs( + prompt=[1, 2], + mm_data_items=MultiModalDataItems({}), + hf_processor_mm_kwargs={"ref_text": "reference"}, + ) + + assert processor.apply(inputs, timing_ctx=None) == "sentinel" + assert captured["inputs"].prompt == _PROMPT_IDS + assert processor.text_tokenizer.call_count == 1 + (text_prompt,), _ = processor.text_tokenizer.call_args + assert text_prompt.endswith("<|text_start|>reference hello world<|text_end|>") + assert captured["inputs"].hf_processor_mm_kwargs["ref_text"] == "reference" + + +def test_media_path_encodes_audio_without_retokenizing(monkeypatch: pytest.MonkeyPatch) -> None: + processor = _make_processor(monkeypatch) + audio = np.zeros(2400, dtype=np.float32) + monkeypatch.setattr(processor, "_get_hf_mm_data", lambda _items: ({"audios": [audio]}, {})) + + out = processor._apply_hf_processor_main(_items({"audio": 1}, {}), {"ref_text": "reference"}) + + assert tuple(out["ref_audio_tokens"].shape) == (8, 4) + assert list(out["ref_audio_len"]) == [4] + assert "input_ids" not in out + assert processor.audio_tokenizer.encode.call_count == 1 + assert processor.text_tokenizer.call_count == 0 + + +def test_media_path_with_no_audio_touches_no_tokenizer(monkeypatch: pytest.MonkeyPatch) -> None: + """Text-only requests and fully cached audio both reach this path with empty items.""" + processor = _make_processor(monkeypatch) + monkeypatch.setattr(processor, "_get_hf_mm_data", lambda _items: ({}, {})) + + out = processor._apply_hf_processor_main(_items({}, {}), {"ref_text": "reference"}) + + assert len(out.keys()) == 0 + assert processor.text_tokenizer.call_count == 0 + assert processor.audio_tokenizer.encode.call_count == 0 + + +def test_legacy_call_hf_processor_is_gone() -> None: + assert not hasattr(OmniVoiceMultiModalProcessor, "_call_hf_processor") diff --git a/vllm_omni/model_executor/models/omnivoice/omnivoice.py b/vllm_omni/model_executor/models/omnivoice/omnivoice.py index 26e15c3dd..a0b0d725e 100644 --- a/vllm_omni/model_executor/models/omnivoice/omnivoice.py +++ b/vllm_omni/model_executor/models/omnivoice/omnivoice.py @@ -74,17 +74,9 @@ class OmniVoiceMultiModalProcessor(OmniMultiModalProcessor[OmniVoiceMultiModalPr def apply(self, inputs: ProcessorInputs, timing_ctx): tokenizer = self.info.get_tokenizer() prompt_text = tokenizer.decode(inputs.prompt, skip_special_tokens=False) - prompt_inputs = self._call_hf_processor( - prompt_text, - {}, - inputs.hf_processor_mm_kwargs, - {}, - ) - prompt_ids = prompt_inputs["input_ids"] - if isinstance(prompt_ids, torch.Tensor): - prompt_ids = prompt_ids.reshape(-1).tolist() - else: - prompt_ids = list(prompt_ids) + # Tokenize once here; _apply_hf_processor_main below only encodes the + # reference audio, so the media path never re-tokenizes the text. + prompt_ids = self._encode_prompt_text(prompt_text, inputs.hf_processor_mm_kwargs).reshape(-1).tolist() inputs = replace( inputs, prompt=prompt_ids, @@ -122,23 +114,39 @@ class OmniVoiceMultiModalProcessor(OmniMultiModalProcessor[OmniVoiceMultiModalPr self._cached_model_dir = model_dir - def _call_hf_processor( + def _apply_hf_processor_main( self, - prompt: str, - mm_data: Mapping[str, object], - mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], + mm_items: MultiModalDataItems, + hf_processor_mm_kwargs: Mapping[str, object], ) -> BatchFeature: + """Encode the reference audio, if any; the prompt was tokenized in apply(). + + Upstream calls this for every request, including text-only ones and + ones whose audio is already cached (empty ``mm_items``), so it must not + touch the text tokenizer. + """ + valid_items = mm_items.select({key for key, count in mm_items.get_all_counts().items() if count > 0}) + mm_data, passthrough = self._get_hf_mm_data(valid_items) + audio = self._get_reference_audio(mm_data) + if audio is None: + return BatchFeature(dict(passthrough)) + processed = self._encode_reference_audio(*audio) + processed.update(passthrough) + return processed + + def _get_reference_audio(self, mm_data: Mapping[str, object]) -> tuple[object, int] | None: + audio = mm_data.get("audio", None) + if audio is None: + audios = mm_data.get("audios") + if audios is not None: + audio = audios[0], self.info.ctx.get_hf_config().sample_rate + return audio + + def _encode_prompt_text(self, prompt: str, mm_kwargs: Mapping[str, object]) -> torch.Tensor: config = self.info.ctx.get_hf_config() model_dir = self.info.ctx.model_config.model self._ensure_cached_runtime_components(model_dir, config) - audio = mm_data.get("audio", None) - if audio is None: - audio = mm_data.get("audios") - if audio is not None: - audio = audio[0], config.sample_rate - # Build text prompt with control tokens lang = mm_kwargs.get("lang", None) instruct = mm_kwargs.get("instruct", None) @@ -161,19 +169,14 @@ class OmniVoiceMultiModalProcessor(OmniMultiModalProcessor[OmniVoiceMultiModalPr full_text = prompt text_prompt = f"{style_text}<|text_start|>{full_text}<|text_end|>" - text_tokens = self.text_tokenizer(text_prompt, return_tensors="pt").input_ids.squeeze(0) # [N_text] + return self.text_tokenizer(text_prompt, return_tensors="pt").input_ids.squeeze(0) # [N_text] - if audio is None: - # Text-only path (auto voice mode) - return BatchFeature( - { - "input_ids": text_tokens, - "input_len": [len(text_tokens)], - } - ) + def _encode_reference_audio(self, audio_signal: object, sr: int) -> BatchFeature: + """Voice cloning: encode the reference audio to 8-codebook tokens.""" + config = self.info.ctx.get_hf_config() + model_dir = self.info.ctx.model_config.model + self._ensure_cached_runtime_components(model_dir, config) - # Voice cloning: encode reference audio to tokens - audio_signal, sr = audio if isinstance(audio_signal, np.ndarray): audio_signal = torch.from_numpy(audio_signal).float() if audio_signal.dim() == 1: @@ -194,14 +197,12 @@ class OmniVoiceMultiModalProcessor(OmniMultiModalProcessor[OmniVoiceMultiModalPr if ref_audio_tokens.dim() == 3: ref_audio_tokens = ref_audio_tokens.squeeze(0) # [8, T_ref] - ft = BatchFeature( + return BatchFeature( { - "input_ids": text_tokens, "ref_audio_tokens": ref_audio_tokens, # [8, T_ref] "ref_audio_len": [ref_audio_tokens.shape[1]], } ) - return ft def _get_mm_fields_config( self,