[Bugfix] Fix Qwen3-TTS Stage 0 prefix-caching correctness (#3665)

Signed-off-by: Yueqian Lin <linyueqian@outlook.com>
Signed-off-by: Rein Yang <ruiruyang2@gmail.com>
Co-authored-by: Rein Yang <ruiruyang2@gmail.com>

the ci failure case is irrelavent to this PR and occurs rarely

tests/core/test_prefix_cache.py

@@ -13,7 +13,7 @@ DEFAULT_SHAPE = torch.Size([NUM_BLOCKS, BLOCK_SIZE, HIDDEN_SIZE])
 
 
 class MockInputBatch:
-    def __init__(self, num_computed_tokens_cpu):
+    def __init__(self, num_computed_tokens_cpu, block_table=None):
         self.req_ids = ["req1", "req2"]
         self.req_id_to_index = {req_id: i for i, req_id in enumerate(self.req_ids)}
         self.num_computed_tokens_cpu = num_computed_tokens_cpu
@@ -21,11 +21,33 @@ class MockInputBatch:
         # Block table is only mocked for validation of length;
         # we don't actually need to add valid values here since
         # we patch the table when testing.
-        class _DummyBlockTable:
-            pass
+        if block_table is None:
 
-        self.block_table = _DummyBlockTable()
-        self.block_table.block_tables = [None]
+            class _DummyBlockTable:
+                pass
+
+            self.block_table = _DummyBlockTable()
+            self.block_table.block_tables = [None]
+        else:
+
+            class _TensorWrapper:
+                def __init__(self, tensor):
+                    self.cpu = tensor
+
+            class _BlockGroup:
+                def __init__(self, tensor):
+                    self.block_table = _TensorWrapper(tensor)
+
+            class _BlockTable:
+                def __init__(self, tensor):
+                    self._group = _BlockGroup(tensor)
+                    self.block_tables = [self._group.block_table]
+
+                def __getitem__(self, idx):
+                    assert idx == 0
+                    return self._group
+
+            self.block_table = _BlockTable(block_table)
 
 
 def get_omni_pcache_with_mm_tensors(feat_dims, seq_len) -> OmniTensorPrefixCache:
@@ -171,6 +193,94 @@ def test_update_with_multimodal_outputs(feat_dims):
             assert torch.all(slot_states == new_output)
 
 
+def test_deferred_multimodal_cache_write_commits_on_completion():
+    cache = get_omni_pcache()
+    mm_key = "codes.audio"
+    block_table = torch.tensor(
+        [
+            [2, 3, 4, 5],
+            [6, 7, 8, 9],
+        ],
+        dtype=torch.long,
+    )
+    input_batch = MockInputBatch(
+        num_computed_tokens_cpu=torch.tensor([4, 0], dtype=torch.long),
+        block_table=block_table,
+    )
+
+    first_chunk = torch.arange(8, dtype=torch.long).reshape(4, 2)
+    cache.update_omni_tensor_prefix_cache(
+        hidden_states=None,
+        multimodal_outputs={mm_key: first_chunk},
+        num_tokens_unpadded=4,
+        slot_mapping=torch.arange(8, 12),
+        skip_mm_cache_keys={mm_key},
+    )
+
+    rows = cache.mm_outputs_cache[mm_key].view(-1, 2)
+    assert torch.all(rows[8:12] == 0)
+
+    cache.stage_deferred_mm_outputs(
+        query_start_loc=torch.tensor([0, 4], dtype=torch.long),
+        input_batch=input_batch,
+        multimodal_outputs={mm_key: first_chunk},
+        num_scheduled_tokens={"req1": 4, "req2": 0},
+        deferred_mm_cache_keys={mm_key},
+    )
+
+    second_chunk = torch.arange(8, 16, dtype=torch.long).reshape(4, 2)
+    input_batch.num_computed_tokens_cpu = torch.tensor([8, 0], dtype=torch.long)
+    cache.stage_deferred_mm_outputs(
+        query_start_loc=torch.tensor([0, 4], dtype=torch.long),
+        input_batch=input_batch,
+        multimodal_outputs={mm_key: second_chunk},
+        num_scheduled_tokens={"req1": 4, "req2": 0},
+        deferred_mm_cache_keys={mm_key},
+    )
+
+    cache.commit_deferred_mm_outputs({"req1"}, input_batch)
+
+    expected = torch.cat([first_chunk, second_chunk], dim=0)
+    rows = cache.mm_outputs_cache[mm_key].view(-1, 2)
+    assert torch.equal(rows[8:16], expected)
+
+
+def test_deferred_multimodal_cache_can_be_merged_on_full_block_hit():
+    cache = get_omni_pcache()
+    mm_key = "codes.audio"
+    block_table = torch.tensor(
+        [
+            [2, 3, 4, 5],
+            [6, 7, 8, 9],
+        ],
+        dtype=torch.long,
+    )
+    input_batch = MockInputBatch(
+        num_computed_tokens_cpu=torch.tensor([8, 0], dtype=torch.long),
+        block_table=block_table,
+    )
+    cached_codes = torch.arange(16, dtype=torch.long).reshape(8, 2)
+    cache.stage_deferred_mm_outputs(
+        query_start_loc=torch.tensor([0, 8], dtype=torch.long),
+        input_batch=input_batch,
+        multimodal_outputs={mm_key: cached_codes},
+        num_scheduled_tokens={"req1": 8, "req2": 0},
+        deferred_mm_cache_keys={mm_key},
+    )
+    cache.commit_deferred_mm_outputs({"req1"}, input_batch)
+
+    new_codes = torch.tensor([[99, 100]], dtype=torch.long)
+    cache.add_prefix_cached_new_req_id("req1")
+    merged = cache.get_merged_multimodal_states(
+        query_start_loc=torch.tensor([0, 1], dtype=torch.long),
+        input_batch=input_batch,
+        multimodal_outputs={mm_key: new_codes},
+        num_scheduled_tokens={"req1": 1, "req2": 0},
+    )
+
+    assert torch.equal(merged[mm_key]["req1"], torch.cat([cached_codes, new_codes], dim=0))
+
+
 ### Tests for Merging
 def fake_get_cached_block_ids(self, req_idx, *args, **kwargs):
     """Fake block table lookup.

tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_preprocess.py

@@ -1,3 +1,4 @@
+from collections import OrderedDict
 from types import SimpleNamespace
 
 import numpy as np
@@ -15,6 +16,8 @@ from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import (
 def _make_minimal_talker():
     model = Qwen3TTSTalkerForConditionalGeneration.__new__(Qwen3TTSTalkerForConditionalGeneration)
     model.talker_config = SimpleNamespace(codec_pad_id=7, num_code_groups=16)
+    model._ref_audio_artifact_cache_max_entries = 256
+    model._ref_audio_artifact_cache = OrderedDict()
     return model
 
 

vllm_omni/core/prefix_cache.py

@@ -47,6 +47,7 @@ class OmniTensorPrefixCache:
         self.mm_outputs_cache = {}
         self.mm_cache_keys = set()
         self._new_req_cache_hit_ids: set[str] = set()
+        self._deferred_mm_outputs: dict[str, dict[str, list[tuple[int, torch.Tensor]]]] = {}
 
     def maybe_init_missing_mm_cache_keys(self, multimodal_outputs: dict, seq_len: int):
         """Given multimodal outputs from executing the model, dynamically
@@ -62,7 +63,17 @@ class OmniTensorPrefixCache:
         determined by the warmup.
         """
         for key, val in multimodal_outputs.items():
-            if isinstance(val, torch.Tensor) and val.shape[0] == seq_len and key not in self.mm_cache_keys:
+            # Only cache per-token feature tensors: 2D+ with first dim == seq_len.
+            # A 1D tensor of shape (seq_len,) is a broadcast scalar (per-request
+            # metadata such as ref_code_len / codec_streaming), not per-token data;
+            # caching it by slot causes a shape mismatch when a later request has a
+            # different scheduled seq length.
+            if (
+                isinstance(val, torch.Tensor)
+                and val.ndim >= 2
+                and val.shape[0] == seq_len
+                and key not in self.mm_cache_keys
+            ):
                 feat_dim = val.shape[-1]
                 self.mm_outputs_cache[key] = self._get_cache_tensor(
                     dtype=val.dtype,
@@ -89,6 +100,10 @@ class OmniTensorPrefixCache:
         """Clears the cache hit IDs to prepare for a new engine step."""
         self._new_req_cache_hit_ids.clear()
 
+    def has_prefix_cached_new_req_ids(self) -> bool:
+        """Return True when this step contains a newly scheduled prefix hit."""
+        return bool(self._new_req_cache_hit_ids)
+
     @staticmethod
     def _coerce_to_cpu_tensor(maybe_gpu_tensor: torch.Tensor) -> torch.Tensor:
         """Convert GPU tensors -> contiguous CPU tensors if needed."""
@@ -101,6 +116,7 @@ class OmniTensorPrefixCache:
         num_tokens_unpadded: int,
         slot_mapping: torch.Tensor,
         num_tokens_padded: int | None = None,
+        skip_mm_cache_keys: set[str] | None = None,
     ):
         """Updates the hidden cache state for the provided hidden states and multimodal outputs.
 
@@ -110,10 +126,12 @@ class OmniTensorPrefixCache:
             num_tokens_unpadded: Number of tokens without padding
             slot_mapping: Slot mapping for the input sequence
             num_tokens_padded: Total number of tokens including padding
+            skip_mm_cache_keys: Multimodal keys whose CPU cache writes are deferred
         """
         unpadded_slot_mapping = slot_mapping[:num_tokens_unpadded]
         if num_tokens_padded is None:
             num_tokens_padded = num_tokens_unpadded
+        skip_mm_cache_keys = skip_mm_cache_keys or set()
 
         if hidden_states is not None:
             # Slice to unpadded portion before caching
@@ -136,6 +154,8 @@ class OmniTensorPrefixCache:
 
             for mm_out_key, mm_cache in self.mm_outputs_cache.items():
                 if mm_out_key in multimodal_outputs:
+                    if mm_out_key in skip_mm_cache_keys:
+                        continue
                     # Slice to unpadded portion before caching
                     mm_state = multimodal_outputs[mm_out_key][:num_tokens_unpadded]
                     mm_state = OmniTensorPrefixCache._coerce_to_cpu_tensor(mm_state)
@@ -143,6 +163,165 @@ class OmniTensorPrefixCache:
                     flat_cache[unpadded_slot_mapping] = mm_state
             logger.debug("Writing to mm output cache for %s tokens", num_tokens_unpadded)
 
+    def stage_deferred_mm_outputs(
+        self,
+        query_start_loc: torch.Tensor,
+        input_batch: InputBatch,
+        multimodal_outputs: dict[str, torch.Tensor] | None,
+        num_scheduled_tokens: dict[str, int],
+        deferred_mm_cache_keys: set[str],
+    ) -> None:
+        """Keep GPU multimodal chunks until a request finishes.
+
+        The normal prefix-cache write path copies every cached multimodal tensor
+        to CPU on every step.  For model outputs that are only needed by future
+        full-block prefix hits, we can keep detached GPU chunks and materialize
+        the CPU cache once when the request completes.
+        """
+        if not multimodal_outputs or not deferred_mm_cache_keys:
+            return
+
+        for mm_key in deferred_mm_cache_keys:
+            mm_state = multimodal_outputs.get(mm_key)
+            if not isinstance(mm_state, torch.Tensor) or mm_state.ndim < 2:
+                continue
+            for req_id in input_batch.req_ids:
+                req_idx = input_batch.req_id_to_index[req_id]
+                sched = int(num_scheduled_tokens.get(req_id, 0))
+                if sched <= 0:
+                    continue
+                start = int(query_start_loc[req_idx])
+                end = start + sched
+                if start >= int(mm_state.shape[0]):
+                    continue
+                end = min(end, int(mm_state.shape[0]))
+                if end <= start:
+                    continue
+
+                computed_end = int(input_batch.num_computed_tokens_cpu[req_idx])
+                token_start = max(0, computed_end - sched)
+                chunk = mm_state[start:end].detach()
+                self._deferred_mm_outputs.setdefault(req_id, {}).setdefault(mm_key, []).append((token_start, chunk))
+
+    def commit_deferred_mm_outputs(
+        self,
+        finished_req_ids: set[str] | list[str],
+        input_batch: InputBatch,
+    ) -> None:
+        """Write deferred multimodal chunks into the CPU prefix cache.
+
+        This must run before finished requests are removed from ``input_batch``,
+        because the block table is needed to map logical token positions to KV
+        cache slots.
+        """
+        if not finished_req_ids or not self._deferred_mm_outputs:
+            return
+
+        for req_id in tuple(finished_req_ids):
+            per_key = self._deferred_mm_outputs.pop(req_id, None)
+            if not per_key:
+                continue
+            req_idx = input_batch.req_id_to_index.get(req_id)
+            if req_idx is None:
+                continue
+            for mm_key, chunks in per_key.items():
+                self._commit_deferred_mm_key(req_idx, input_batch, mm_key, chunks)
+
+    def discard_deferred_mm_outputs(self, req_id: str) -> None:
+        """Drop deferred chunks for requests that leave without a cache commit."""
+        self._deferred_mm_outputs.pop(req_id, None)
+
+    def _commit_deferred_mm_key(
+        self,
+        req_idx: int,
+        input_batch: InputBatch,
+        mm_key: str,
+        chunks: list[tuple[int, torch.Tensor]],
+    ) -> None:
+        if not chunks:
+            return
+
+        chunks = sorted(chunks, key=lambda item: item[0])
+        first_start = chunks[0][0]
+        expected_start = first_start
+        tensors: list[torch.Tensor] = []
+        for token_start, tensor in chunks:
+            if tensor.numel() == 0:
+                continue
+            if token_start != expected_start:
+                if tensors:
+                    self._write_deferred_mm_tensor(req_idx, input_batch, mm_key, tensors, first_start)
+                first_start = token_start
+                expected_start = token_start
+                tensors = []
+            tensors.append(tensor)
+            expected_start += int(tensor.shape[0])
+
+        if tensors:
+            self._write_deferred_mm_tensor(req_idx, input_batch, mm_key, tensors, first_start)
+
+    def _write_deferred_mm_tensor(
+        self,
+        req_idx: int,
+        input_batch: InputBatch,
+        mm_key: str,
+        tensors: list[torch.Tensor],
+        token_start: int,
+    ) -> None:
+        mm_state = tensors[0] if len(tensors) == 1 else torch.cat(tensors, dim=0)
+        if mm_state.ndim < 2 or mm_state.shape[0] == 0:
+            return
+
+        if mm_key not in self.mm_cache_keys:
+            self.mm_outputs_cache[mm_key] = self._get_cache_tensor(
+                dtype=mm_state.dtype,
+                hidden_size=int(mm_state.shape[-1]),
+            )
+            self.mm_cache_keys.add(mm_key)
+            logger.info(
+                "Initializing deferred multimodal output cache of size %s for key: %s",
+                list(self.mm_outputs_cache[mm_key].shape),
+                mm_key,
+            )
+
+        slots = self._get_slot_ids_for_token_range(
+            req_idx=req_idx,
+            input_batch=input_batch,
+            token_start=token_start,
+            num_tokens=int(mm_state.shape[0]),
+        )
+        if slots.numel() == 0:
+            return
+
+        mm_state_cpu = OmniTensorPrefixCache._coerce_to_cpu_tensor(mm_state[: slots.numel()])
+        mm_cache = self.mm_outputs_cache[mm_key]
+        flat_cache = mm_cache.view(-1, mm_cache.shape[-1])
+        flat_cache[slots] = mm_state_cpu
+
+    def _get_slot_ids_for_token_range(
+        self,
+        req_idx: int,
+        input_batch: InputBatch,
+        token_start: int,
+        num_tokens: int,
+    ) -> torch.Tensor:
+        if num_tokens <= 0:
+            return torch.empty((0,), dtype=torch.long)
+
+        block_table = input_batch.block_table[0].block_table.cpu
+        token_positions = torch.arange(token_start, token_start + num_tokens, dtype=torch.long)
+        block_offsets = token_positions // self.block_size
+        max_blocks = int(block_table.shape[1])
+        valid = block_offsets < max_blocks
+        if not bool(valid.all()):
+            token_positions = token_positions[valid]
+            block_offsets = block_offsets[valid]
+        if token_positions.numel() == 0:
+            return torch.empty((0,), dtype=torch.long)
+
+        block_ids = block_table[req_idx, block_offsets].to(torch.long)
+        return block_ids * self.block_size + (token_positions % self.block_size)
+
     def _coerce_to_payload_dict(
         self,
         element: object,
@@ -151,17 +330,25 @@ class OmniTensorPrefixCache:
         num_scheduled_tokens: dict[str, int],
     ) -> dict[str, object]:
         """Build the multimodal passthrough data per request for
-        the object under consideration. This is identical to the case
-        for no prefix cache when we tensor does have a first dimension
-        matching the seq len.
+        the object under consideration. This mirrors the no-prefix-cache
+        path: a tensor whose first dimension matches the total scheduled
+        token count is sliced per request, so 1D (seq_len,) metadata that
+        is intentionally not cached (e.g. ref_code_len, codec_streaming)
+        is still split per request instead of leaking the whole batch.
         """
+        total_scheduled_tokens = sum(int(num_scheduled_tokens[r]) for r in input_batch.req_ids)
         elem_dict = {}
         for req_id in input_batch.req_ids:
             req_idx = input_batch.req_id_to_index[req_id]
             start = query_start_loc[req_idx]
             end = start + num_scheduled_tokens[req_id]
             elem_dict[req_id] = to_payload_element(
-                element, req_idx, start=start, end=end, pass_lists_through=True, seq_len=None
+                element,
+                req_idx,
+                start=start,
+                end=end,
+                pass_lists_through=True,
+                seq_len=total_scheduled_tokens,
             )
         return elem_dict
 

vllm_omni/deploy/qwen3_tts.yaml

@@ -34,7 +34,12 @@ stages:
     max_num_seqs: 10
     gpu_memory_utilization: 0.3
     trust_remote_code: true
-    enable_prefix_caching: false
+    # Prefix caching on by default. The HS opt-out (talker only consumes the
+    # last token's hidden state) and deferred mm-output writes added in this
+    # PR make per-step bookkeeping negligible, and any text repetition across
+    # requests (same voice or shared ref_audio) is a net throughput win.
+    # Disable explicitly if memory headroom is tight on a single-GPU deploy.
+    enable_prefix_caching: true
     async_scheduling: true
     max_num_batched_tokens: 512
     max_model_len: 4096

vllm_omni/engine/orchestrator.py

@@ -107,7 +107,7 @@ def build_engine_core_request_from_tokens(
         pooling_params=pooling_params,
         arrival_time=arrival_time,
         lora_request=getattr(params, "lora_request", None),
-        cache_salt=None,
+        cache_salt=prompt.get("cache_salt"),
         data_parallel_rank=None,
         prompt_embeds=prompt_embeds,
         resumable=resumable,

vllm_omni/entrypoints/omni_base.py

@@ -277,6 +277,10 @@ class OmniBase(PDDisaggregationMixin):
         try:
             if req_state is None or req_state.metrics is None:
                 return
+            if self.log_stats:
+                # Emit per-request orchestrator timing (including e2e_total_ms)
+                # before dropping request state.
+                req_state.metrics.build_and_log_summary()
         except Exception:
             logger.exception(
                 "[%s] Failed to build/log summary for req=%s",

vllm_omni/entrypoints/openai/serving_speech.py

@@ -1,5 +1,6 @@
 import asyncio
 import base64
+import hashlib
 import io
 import json
 import math
@@ -186,6 +187,55 @@ def _validate_path_within_directory(file_path: Path, directory: Path) -> bool:
         return False
 
 
+def _conditioning_cache_salt(request, tts_params: dict | None = None) -> str:
+    """Stable hash of the real Stage 0 conditioning for the prefix cache.
+
+    The talker's vLLM prompt is placeholder token ids; the real inputs are
+    rebuilt from text / ref_audio / ref_text into inputs_embeds. vLLM hashes
+    token ids (folded with cache_salt) for prefix caching, so without a salt
+    every request collides and a hit could reuse KV from a semantically
+    different input. Tying the salt to the conditioning keeps a hit safe:
+    identical conditioning may share the prefix, any difference never does.
+
+    Raw request fields alone are not enough for uploaded voices: the request
+    only carries the voice *name* (ref_audio/ref_text/task_type are resolved
+    from stored voice data into ``tts_params``). Delete + re-upload under the
+    same name leaves every raw field identical, so the resolved conditioning
+    must also be folded in. ``voice_created_at`` bumps on every (re-)upload,
+    which uniquely identifies the resolved reference artifact together with
+    the voice name; the decoded ref_audio array itself need not be hashed.
+    """
+    h = hashlib.sha256()
+    for part in (
+        request.input,
+        request.task_type,
+        request.language,
+        request.voice,
+        request.ref_text,
+        request.ref_audio,
+        request.instructions,
+        request.x_vector_only_mode,
+        request.speaker_embedding,
+    ):
+        h.update(b"\x00")
+        if part is not None:
+            h.update(repr(part).encode("utf-8"))
+    # Fold resolved conditioning that is auto-derived for uploaded voices and
+    # absent from the raw request.
+    for key in (
+        "voice_created_at",
+        "task_type",
+        "speaker",
+        "ref_text",
+        "x_vector_only_mode",
+    ):
+        h.update(b"\x00")
+        value = tts_params.get(key) if tts_params is not None else None
+        if value is not None:
+            h.update(repr(value).encode("utf-8"))
+    return h.hexdigest()[:32]
+
+
 class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
     _diffusion_mode: bool = False
     _tts_executor: ThreadPoolExecutor | None = None
@@ -1500,6 +1550,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         request_id: str,
         response_format: str = "pcm",
         raw_request: Request | None = None,
+        request_start_s: float | None = None,
     ):
         """Generate audio chunks for streaming response.
 
@@ -1520,6 +1571,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         prev_count = 0
         sample_rate_val = 24000
         first_chunk = True
+        first_audio_chunk_s: float | None = None
+        stream_start_s = request_start_s if request_start_s is not None else time.perf_counter()
 
         try:
             async for res in generator:
@@ -1571,11 +1624,40 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                         stream_format="audio",
                         base64_encode=False,
                     )
+                    if first_audio_chunk_s is None:
+                        first_audio_chunk_s = time.perf_counter()
                     yield self.create_audio(audio_obj).audio_data
+            total_ms = (time.perf_counter() - stream_start_s) * 1000.0
+            if first_audio_chunk_s is not None:
+                first_chunk_ms = (first_audio_chunk_s - stream_start_s) * 1000.0
+                logger.info(
+                    "[SpeechE2E] request_id=%s stream=true status=ok total_ms=%.2f first_chunk_ms=%.2f",
+                    request_id,
+                    total_ms,
+                    first_chunk_ms,
+                )
+            else:
+                logger.info(
+                    "[SpeechE2E] request_id=%s stream=true status=ok total_ms=%.2f first_chunk_ms=NA",
+                    request_id,
+                    total_ms,
+                )
         except asyncio.CancelledError:
+            total_ms = (time.perf_counter() - stream_start_s) * 1000.0
+            logger.info(
+                "[SpeechE2E] request_id=%s stream=true status=cancelled total_ms=%.2f",
+                request_id,
+                total_ms,
+            )
             logger.info("Streaming request %s cancelled by client", request_id)
             raise
         except EngineDeadError as e:
+            total_ms = (time.perf_counter() - stream_start_s) * 1000.0
+            logger.error(
+                "[SpeechE2E] request_id=%s stream=true status=engine_dead total_ms=%.2f",
+                request_id,
+                total_ms,
+            )
             logger.error(
                 "EngineDeadError during streaming speech for %s: %s",
                 request_id,
@@ -1589,6 +1671,13 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                 )
             raise
         except Exception as e:
+            total_ms = (time.perf_counter() - stream_start_s) * 1000.0
+            logger.exception(
+                "[SpeechE2E] request_id=%s stream=true status=error total_ms=%.2f error=%s",
+                request_id,
+                total_ms,
+                e,
+            )
             logger.exception("Streaming speech generation failed for %s: %s", request_id, e)
             raise
 
@@ -2195,6 +2284,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                     tts_params["seed"] = [sampling_params_list[0].seed]
                 prompt = tokens_input(prompt_token_ids=[1])
                 prompt["additional_information"] = tts_params
+                prompt["cache_salt"] = _conditioning_cache_salt(request, tts_params)
             else:
                 tts_params = self._build_tts_params(request)
                 # Resolve ref_audio (explicit or auto-set for uploaded voices)
@@ -2210,6 +2300,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                 ph_len = await self._estimate_prompt_len_async(tts_params)
                 prompt = tokens_input(prompt_token_ids=[1] * ph_len)
                 prompt["additional_information"] = tts_params
+                prompt["cache_salt"] = _conditioning_cache_salt(request, tts_params)
         else:
             # Qwen omni models (Qwen3-Omni, Qwen2.5-Omni) use a "talker"
             # stage whose preprocess requires chat-templated tokens.  The
@@ -2651,6 +2742,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             return error_check_ret
 
         request_id = f"speech-{random_uuid()}"
+        request_start_s = time.perf_counter()
         if raw_request:
             raw_request.state.request_metadata = RequestResponseMetadata(
                 request_id=request_id,
@@ -2683,20 +2775,58 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                         request_id,
                         response_format,
                         raw_request=raw_request,
+                        request_start_s=request_start_s,
                     ),
                     media_type=media_type,
                 )
 
             audio_bytes, media_type = await self._generate_audio_bytes(request, request_id=request_id)
+            total_ms = (time.perf_counter() - request_start_s) * 1000.0
+            logger.info(
+                "[SpeechE2E] request_id=%s stream=false status=ok total_ms=%.2f response_bytes=%d",
+                request_id,
+                total_ms,
+                len(audio_bytes) if isinstance(audio_bytes, (bytes, bytearray)) else len(str(audio_bytes)),
+            )
             return Response(content=audio_bytes, media_type=media_type)
 
         except asyncio.CancelledError:
+            total_ms = (time.perf_counter() - request_start_s) * 1000.0
+            logger.info(
+                "[SpeechE2E] request_id=%s stream=%s status=cancelled total_ms=%.2f",
+                request_id,
+                bool(request.stream),
+                total_ms,
+            )
             return self.create_error_response("Client disconnected")
         except (EngineGenerateError, EngineDeadError):
+            total_ms = (time.perf_counter() - request_start_s) * 1000.0
+            logger.error(
+                "[SpeechE2E] request_id=%s stream=%s status=engine_error total_ms=%.2f",
+                request_id,
+                bool(request.stream),
+                total_ms,
+            )
             raise  # Propagate to the global Omni exception handler
         except ValueError as e:
+            total_ms = (time.perf_counter() - request_start_s) * 1000.0
+            logger.warning(
+                "[SpeechE2E] request_id=%s stream=%s status=bad_request total_ms=%.2f error=%s",
+                request_id,
+                bool(request.stream),
+                total_ms,
+                e,
+            )
             return self.create_error_response(e)
         except Exception as e:
+            total_ms = (time.perf_counter() - request_start_s) * 1000.0
+            logger.exception(
+                "[SpeechE2E] request_id=%s stream=%s status=error total_ms=%.2f error=%s",
+                request_id,
+                bool(request.stream),
+                total_ms,
+                e,
+            )
             logger.exception("Speech generation failed: %s", e)
             return self.create_error_response(f"Speech generation failed: {e}")
 

vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py

@@ -2,8 +2,10 @@ from __future__ import annotations
 
 import base64
 import copy
+import hashlib
 import io
 import os
+from collections import OrderedDict
 from collections.abc import Callable, Iterable, Mapping
 from typing import Any
 from urllib.parse import urlparse
@@ -40,6 +42,7 @@ logger = init_logger(__name__)
 _TRAILING_TEXT_COMPACT_MIN_FRAMES = 64
 _PRECOMPUTED_REF_CODE_KEY = "precomputed_ref"
 _NORMALIZED_REF_AUDIO_KEY = "_qwen3_tts_normalized_ref_audio"
+_REF_AUDIO_CACHE_KEY = "_qwen3_tts_ref_audio_cache_key"
 _PRECOMPUTED_TEXT_IDS_KEY = "_qwen3_tts_text_ids"
 _PRECOMPUTED_REF_IDS_KEY = "_qwen3_tts_ref_ids"
 
@@ -354,6 +357,17 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
         self.have_multimodal_outputs = True
         self.has_preprocess = True
         self.has_postprocess = True
+        # Qwen3-TTS postprocess() only reads hidden_states[-1, :]. On a prefix-
+        # cache hit, the last hidden state is in the newly computed tail, so
+        # reconstructing the full cached_prefix + new_tail span is wasted work.
+        # Opt out of the per-step GPU->CPU hidden-state cache write and merged-
+        # tensor read; postprocess receives the tail-only slice instead, which
+        # avoids ~18 ms merge + ~6 ms write per step (Sy0307 profile, #3665).
+        self.requires_full_prefix_cached_hidden_states = False
+        # ``codes.audio`` is only needed for future prefix-hit reconstruction
+        # after a request has produced codec rows. Keep per-step rows on GPU and
+        # materialize the CPU OmniTensorPrefixCache entry once at completion.
+        self.deferred_prefix_cache_mm_keys = {"codes.audio"}
         # Used by OmniGPUModelRunner for the GPU-side MTP fast-path.
         self.mtp_hidden_size = int(self.talker_config.hidden_size)
         # OmniGPUModelRunner will store talker_mtp output under this key in
@@ -443,6 +457,8 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
         self._speech_tokenizer: Qwen3TTSTokenizer | None = None
 
         self._speaker_cache = get_speaker_cache()
+        self._ref_audio_artifact_cache_max_entries = 256
+        self._ref_audio_artifact_cache: OrderedDict[str, dict[str, torch.Tensor | bool]] = OrderedDict()
         raw_subtalker_sampling = getattr(vllm_config.model_config, "subtalker_sampling_params", None)
         self._subtalker_sampling_params: dict[str, Any] = (
             dict(raw_subtalker_sampling) if isinstance(raw_subtalker_sampling, Mapping) else {}
@@ -634,9 +650,11 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
                     info_update["meta"]["ref_code_len"] = int(ref_code_len)
                 # Always return a span_len slice; if the scheduled placeholder is longer, pad with tts_pad_embed.
                 # This preserves placeholder/embedding alignment.
-                offset = 0
-                s = 0
-                e = span_len
+                offset = int(info_dict.get("_omni_num_computed_tokens", 0) or 0)
+                if offset < 0:
+                    offset = 0
+                s = max(0, min(offset, int(prompt_embeds_cpu.shape[0])))
+                e = max(0, min(offset + span_len, int(prompt_embeds_cpu.shape[0])))
                 take = prompt_embeds_cpu[s:e]
                 if int(take.shape[0]) < span_len:
                     pad_n = int(span_len - int(take.shape[0]))
@@ -1393,9 +1411,50 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
             # 12Hz: likely [T, Q] or [B, T, Q]
             if ref_code.ndim == 3:
                 ref_code = ref_code[0]
-            return ref_code.to(device=next(self.parameters()).device, dtype=torch.long)
+            out = ref_code.to(device=next(self.parameters()).device, dtype=torch.long)
+            return out
         raise ValueError("SpeechTokenizer.encode did not return audio_codes tensor")
 
+    @staticmethod
+    def _make_ref_audio_cache_key(wav: np.ndarray, sr: int) -> str:
+        wav_f32 = wav.astype(np.float32, copy=False).reshape(-1)
+        h = hashlib.sha1()
+        h.update(int(sr).to_bytes(4, byteorder="little", signed=False))
+        h.update(int(wav_f32.size).to_bytes(8, byteorder="little", signed=False))
+        h.update(wav_f32.tobytes(order="C"))
+        return h.hexdigest()
+
+    def _get_ref_audio_artifacts(self, cache_key: str) -> dict[str, torch.Tensor | bool] | None:
+        entry = self._ref_audio_artifact_cache.get(cache_key)
+        if entry is None:
+            return None
+        # Touch for LRU.
+        self._ref_audio_artifact_cache.move_to_end(cache_key, last=True)
+        return entry
+
+    def _put_ref_audio_artifacts(
+        self,
+        cache_key: str,
+        *,
+        ref_code: torch.Tensor | None = None,
+        ref_spk_embedding: torch.Tensor | None = None,
+    ) -> None:
+        if not cache_key:
+            return
+        entry = self._ref_audio_artifact_cache.get(cache_key)
+        if entry is None:
+            entry = {}
+        if isinstance(ref_code, torch.Tensor):
+            entry["ref_code"] = ref_code.detach().to("cpu", dtype=torch.long).contiguous()
+        if isinstance(ref_spk_embedding, torch.Tensor):
+            entry["ref_spk_embedding"] = ref_spk_embedding.detach().to("cpu", dtype=torch.bfloat16).reshape(-1)
+        if not entry:
+            return
+        self._ref_audio_artifact_cache[cache_key] = entry
+        self._ref_audio_artifact_cache.move_to_end(cache_key, last=True)
+        while len(self._ref_audio_artifact_cache) > self._ref_audio_artifact_cache_max_entries:
+            self._ref_audio_artifact_cache.popitem(last=False)
+
     @staticmethod
     def _first_value(value: object, default: object = None) -> object:
         if isinstance(value, list):
@@ -1544,6 +1603,14 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
                 # invalid ref_audio errors with its existing messages.
                 continue
             info_dict[_NORMALIZED_REF_AUDIO_KEY] = (wav, sr)
+            cache_key = self._make_ref_audio_cache_key(wav, sr)
+            info_dict[_REF_AUDIO_CACHE_KEY] = cache_key
+            cached = self._get_ref_audio_artifacts(cache_key)
+            if cached is not None:
+                cached_ref_code = self._coerce_ref_code_tensor(cached.get("ref_code"), device=device)
+                if isinstance(cached_ref_code, torch.Tensor):
+                    info_dict.setdefault("codes", {})[_PRECOMPUTED_REF_CODE_KEY] = cached_ref_code
+                    continue
             groups.setdefault(int(sr), []).append((info_dict, wav, int(sr)))
 
         if pending_text or pending_ref_text:
@@ -1587,6 +1654,9 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
                     continue
                 info_dict.setdefault("codes", {})[_PRECOMPUTED_REF_CODE_KEY] = ref_code_t
                 info_dict[_NORMALIZED_REF_AUDIO_KEY] = (wav, item_sr)
+                cache_key = info_dict.get(_REF_AUDIO_CACHE_KEY)
+                if isinstance(cache_key, str) and cache_key:
+                    self._put_ref_audio_artifacts(cache_key, ref_code=ref_code_t)
 
     def _generate_icl_prompt(
         self,
@@ -1773,6 +1843,7 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
             voice_clone_prompt = _normalize_voice_clone_prompt(info_dict.get("voice_clone_prompt"))
             ref_audio_wav: np.ndarray | None = None
             ref_audio_sr: int | None = None
+            ref_audio_cache_key: str | None = None
 
             def _get_ref_audio() -> tuple[np.ndarray, int]:
                 nonlocal ref_audio_wav, ref_audio_sr
@@ -1792,6 +1863,32 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
                     ref_audio_wav, ref_audio_sr = self._normalize_ref_audio(ref_audio_list[0])
                 return ref_audio_wav, ref_audio_sr
 
+            def _get_ref_audio_cache_key() -> str:
+                nonlocal ref_audio_cache_key
+                if isinstance(ref_audio_cache_key, str) and ref_audio_cache_key:
+                    return ref_audio_cache_key
+                cached_key = info_dict.get(_REF_AUDIO_CACHE_KEY)
+                if isinstance(cached_key, str) and cached_key:
+                    ref_audio_cache_key = cached_key
+                    return ref_audio_cache_key
+                wav_np, sr = _get_ref_audio()
+                ref_audio_cache_key = self._make_ref_audio_cache_key(wav_np, sr)
+                info_dict[_REF_AUDIO_CACHE_KEY] = ref_audio_cache_key
+                return ref_audio_cache_key
+
+            def _try_get_ref_audio_cache_key() -> str | None:
+                nonlocal ref_audio_cache_key
+                if isinstance(ref_audio_cache_key, str) and ref_audio_cache_key:
+                    return ref_audio_cache_key
+                cached_key = info_dict.get(_REF_AUDIO_CACHE_KEY)
+                if isinstance(cached_key, str) and cached_key:
+                    ref_audio_cache_key = cached_key
+                    return ref_audio_cache_key
+                ref_audio_list = info_dict.get("ref_audio")
+                if isinstance(ref_audio_list, list) and ref_audio_list:
+                    return _get_ref_audio_cache_key()
+                return None
+
             # Speaker cache: only for uploaded (named) speakers
             _speaker_cache_key = None
             if voice_clone_prompt is None:
@@ -1849,12 +1946,24 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
                 if isinstance(ref_code_t, torch.Tensor):
                     ref_code_len = int(ref_code_t.shape[0])
             if ref_code_t is None and in_context_mode:
-                # Compute ref_code from ref_audio if not provided.
-                wav_np, sr = _get_ref_audio()
-                ref_code_t = self._encode_ref_audio_to_code(wav_np, sr).to(device=input_ids.device)
-                ref_code_len = int(ref_code_t.shape[0])
+                cache_key = _get_ref_audio_cache_key()
+                cached = self._get_ref_audio_artifacts(cache_key)
+                cached_ref_code = None
+                if isinstance(cached, dict):
+                    cached_ref_code = cached.get("ref_code")
+                ref_code_t = self._coerce_ref_code_tensor(cached_ref_code, device=input_ids.device)
+                if isinstance(ref_code_t, torch.Tensor):
+                    ref_code_len = int(ref_code_t.shape[0])
+                else:
+                    wav_np, sr = _get_ref_audio()
+                    ref_code_t = self._encode_ref_audio_to_code(wav_np, sr).to(device=input_ids.device)
+                    ref_code_len = int(ref_code_t.shape[0])
+                    self._put_ref_audio_artifacts(cache_key, ref_code=ref_code_t)
             if isinstance(ref_code_t, torch.Tensor):
                 ref_code_prompt = ref_code_t
+                cache_key = _try_get_ref_audio_cache_key()
+                if isinstance(cache_key, str) and cache_key:
+                    self._put_ref_audio_artifacts(cache_key, ref_code=ref_code_prompt)
 
             # Speaker embedding: use prompt embed if provided; otherwise extract from audio.
             # NOTE: Do NOT use _as_singleton here — the embedding may be a plain
@@ -1869,8 +1978,17 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
                 # Plain list/array from API (survived msgspec IPC serialization).
                 speaker_embed = torch.tensor(spk, dtype=torch.bfloat16, device=input_ids.device).view(1, 1, -1)
             else:
-                wav_np, sr = _get_ref_audio()
-                speaker_embed = self._extract_speaker_embedding(wav_np, sr).view(1, 1, -1)
+                cache_key = _get_ref_audio_cache_key()
+                cached = self._get_ref_audio_artifacts(cache_key)
+                cached_spk = None
+                if isinstance(cached, dict):
+                    cached_spk = cached.get("ref_spk_embedding")
+                if isinstance(cached_spk, torch.Tensor):
+                    speaker_embed = cached_spk.to(device=input_ids.device, dtype=torch.bfloat16).view(1, 1, -1)
+                else:
+                    wav_np, sr = _get_ref_audio()
+                    speaker_embed = self._extract_speaker_embedding(wav_np, sr).view(1, 1, -1)
+                    self._put_ref_audio_artifacts(cache_key, ref_spk_embedding=speaker_embed)
 
             # Cache miss: store extraction result
             if _speaker_cache_key is not None and speaker_embed is not None:

vllm_omni/platforms/npu/worker/npu_ar_model_runner.py

@@ -155,6 +155,11 @@ class NPUARModelRunner(OmniNPUModelRunner):
         finally:
             set_cudagraph_capturing_enabled(False)
 
+    def _model_needs_full_prefix_hidden_states(self) -> bool:
+        """See gpu_ar_model_runner._model_needs_full_prefix_hidden_states."""
+        model = getattr(self, "model", None)
+        return bool(getattr(model, "requires_full_prefix_cached_hidden_states", True))
+
     def _maybe_update_prefix_cache(
         self,
         hidden_states: torch.Tensor,
@@ -169,8 +174,9 @@ class NPUARModelRunner(OmniNPUModelRunner):
                     type(multimodal_outputs),
                 )
 
+            hs_for_cache = hidden_states if self._model_needs_full_prefix_hidden_states() else None
             self.omni_prefix_cache.update_omni_tensor_prefix_cache(
-                hidden_states=hidden_states,
+                hidden_states=hs_for_cache,
                 multimodal_outputs=flatten_payload(multimodal_outputs) if multimodal_outputs else multimodal_outputs,
                 num_tokens_unpadded=num_tokens_unpadded,
                 slot_mapping=self.input_batch.block_table[0].slot_mapping.cpu,
@@ -185,12 +191,13 @@ class NPUARModelRunner(OmniNPUModelRunner):
     ) -> tuple[dict[str, torch.Tensor] | None, dict | None]:
         combined_hidden_states, combined_multimodal_outputs = None, None
         if self.omni_prefix_cache is not None:
-            combined_hidden_states = self.omni_prefix_cache.get_merged_hidden_states(
-                query_start_loc=self.query_start_loc.cpu,
-                input_batch=self.input_batch,
-                hidden_states=hidden_states,
-                num_scheduled_tokens=num_scheduled_tokens,
-            )
+            if self._model_needs_full_prefix_hidden_states():
+                combined_hidden_states = self.omni_prefix_cache.get_merged_hidden_states(
+                    query_start_loc=self.query_start_loc.cpu,
+                    input_batch=self.input_batch,
+                    hidden_states=hidden_states,
+                    num_scheduled_tokens=num_scheduled_tokens,
+                )
             combined_multimodal_outputs = self.omni_prefix_cache.get_merged_multimodal_states(
                 query_start_loc=self.query_start_loc.cpu,
                 input_batch=self.input_batch,

vllm_omni/utils/mm_outputs.py

@@ -68,12 +68,13 @@ def to_payload_element(
             if we need to avoid splitting nonempty lists prior to calling
             postprocess, which is the case for prefix cache.
         seq_len: Optional sequence length (i.e., dim 0 of hidden states).
-            This should be set to None in the prefix caching case, because
-            the condition that would be executed here is the same as the
-            criteria for being added to the multimodal outputs cache.
+            When set, a tensor whose first dimension equals seq_len is
+            sliced per request. The prefix cache passthrough also passes
+            the total scheduled token count here so 1D (seq_len,) metadata
+            that is intentionally not cached is still split per request.
     """
-    # Prefix cache won't hit this case because this is the condition
-    # for being a mm_cache_key in the multimodal outputs tensor.
+    # Cached per-token tensors are merged elsewhere; here a first dim
+    # equal to seq_len means a per-request slice is required.
     if seq_len is not None and isinstance(element, torch.Tensor) and element.shape[0] == seq_len:
         return element[start:end].contiguous()
     # Every other case is shared between prefix cache (passthrough data)

vllm_omni/worker/gpu_ar_model_runner.py

@@ -203,6 +203,24 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
         finally:
             set_cudagraph_capturing_enabled(False)
 
+    def _model_needs_full_prefix_hidden_states(self) -> bool:
+        """Opt-out hook for models whose postprocess only consumes the tail.
+
+        When False, we skip both the per-step GPU->CPU hidden-state write into
+        OmniTensorPrefixCache and the merged-tensor reconstruction on hits;
+        postprocess receives the normal scheduled-token slice instead. Models
+        that need the full cached_prefix + new_tail span (default) are not
+        affected.
+        """
+        model = getattr(self, "model", None)
+        return bool(getattr(model, "requires_full_prefix_cached_hidden_states", True))
+
+    def _deferred_prefix_cache_mm_keys(self) -> set[str]:
+        """Model-declared multimodal keys whose prefix-cache writes are deferred."""
+        model = getattr(self, "model", None)
+        keys = getattr(model, "deferred_prefix_cache_mm_keys", ())
+        return set(keys or ())
+
     def _maybe_update_prefix_cache(
         self,
         hidden_states: torch.Tensor,
@@ -225,12 +243,14 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
                     type(multimodal_outputs),
                 )
 
+            hs_for_cache = hidden_states if self._model_needs_full_prefix_hidden_states() else None
             self.omni_prefix_cache.update_omni_tensor_prefix_cache(
-                hidden_states=hidden_states,
+                hidden_states=hs_for_cache,
                 multimodal_outputs=flatten_payload(multimodal_outputs) if multimodal_outputs else multimodal_outputs,
                 num_tokens_unpadded=num_tokens_unpadded,
                 slot_mapping=self.input_batch.block_table[0].slot_mapping.cpu,
                 num_tokens_padded=num_tokens_padded,
+                skip_mm_cache_keys=self._deferred_prefix_cache_mm_keys(),
             )
 
     def _maybe_get_combined_prefix_cache_tensors(
@@ -246,12 +266,18 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
         # the prefix cached hidden states and multimodal states.
         combined_hidden_states, combined_multimodal_outputs = None, None
         if self.omni_prefix_cache is not None:
-            combined_hidden_states = self.omni_prefix_cache.get_merged_hidden_states(
-                query_start_loc=self.query_start_loc.cpu,
-                input_batch=self.input_batch,
-                hidden_states=hidden_states,
-                num_scheduled_tokens=num_scheduled_tokens,
-            )
+            if (
+                not self._model_needs_full_prefix_hidden_states()
+                and not self.omni_prefix_cache.has_prefix_cached_new_req_ids()
+            ):
+                return None, None
+            if self._model_needs_full_prefix_hidden_states():
+                combined_hidden_states = self.omni_prefix_cache.get_merged_hidden_states(
+                    query_start_loc=self.query_start_loc.cpu,
+                    input_batch=self.input_batch,
+                    hidden_states=hidden_states,
+                    num_scheduled_tokens=num_scheduled_tokens,
+                )
             combined_multimodal_outputs = self.omni_prefix_cache.get_merged_multimodal_states(
                 query_start_loc=self.query_start_loc.cpu,
                 input_batch=self.input_batch,
@@ -313,6 +339,12 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
                 if flush_ids:
                     self.flush_full_payload_outputs(flush_ids)
 
+        if self.omni_prefix_cache is not None and scheduler_output.finished_req_ids:
+            self.omni_prefix_cache.commit_deferred_mm_outputs(
+                set(scheduler_output.finished_req_ids),
+                self.input_batch,
+            )
+
         if self.routed_experts_initialized:
             capturer = get_global_experts_capturer()
             if capturer is not None and hasattr(capturer, "finalize_pending_copy"):
@@ -884,9 +916,22 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
                 dtype=np.int32,
             )
         query_start_loc_cpu = self.query_start_loc.cpu
+        if self.omni_prefix_cache is not None:
+            deferred_mm_cache_keys = self._deferred_prefix_cache_mm_keys()
+            if deferred_mm_cache_keys:
+                self.omni_prefix_cache.stage_deferred_mm_outputs(
+                    query_start_loc=query_start_loc_cpu,
+                    input_batch=self.input_batch,
+                    multimodal_outputs=(
+                        flatten_payload(multimodal_outputs) if multimodal_outputs else multimodal_outputs
+                    ),
+                    num_scheduled_tokens=scheduler_output.num_scheduled_tokens,
+                    deferred_mm_cache_keys=deferred_mm_cache_keys,
+                )
 
         pooler_output: list[dict[str, object]] | None = None
         if needs_pooler_payload:
+            mm_cpu = None
             # Prior to applying the post-processing func, extract
             # the prefix cached hidden states and multimodal states.
             if self.omni_prefix_cache is not None:
@@ -899,8 +944,8 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
                     scheduler_output.num_scheduled_tokens,
                 )
             # Otherwise we don't have the mm CPU data yet, so we still need to build it
-            if self.omni_prefix_cache is None:
-                mm_cpu = build_mm_cpu(flatten_payload(multimodal_outputs))
+            if self.omni_prefix_cache is None or combined_multimodal_outputs is None:
+                mm_cpu = build_mm_cpu(flatten_payload(multimodal_outputs) if multimodal_outputs else multimodal_outputs)
 
             self._process_additional_information_updates(
                 hidden_states,

vllm_omni/worker/gpu_model_runner.py

@@ -332,6 +332,8 @@ class OmniGPUModelRunner(GPUModelRunner):
             self.requests.pop(req_id, None)
             self.model_intermediate_buffer.pop(req_id, None)
             self.num_prompt_logprobs.pop(req_id, None)
+            if self.omni_prefix_cache is not None:
+                self.omni_prefix_cache.discard_deferred_mm_outputs(req_id)
             if hasattr(self, "_downstream_payload_cache"):
                 self._downstream_payload_cache.pop(req_id, None)
             if hasattr(self, "_talker_mtp_generators"):