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 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01MpXe2LeRrUQGjZfKFaocMZ
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>

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")

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,