[Perf] Deduplicate AR prefix cache hidden-state CPU staging (#3734)

Signed-off-by: TaffyOfficial <2324465096@qq.com>
Co-authored-by: TaffyOfficial <2324465096@qq.com>

tests/core/test_prefix_cache.py

@@ -157,6 +157,75 @@ def test_update_no_multimodal():
         assert torch.all(slot_states == new_states)
 
 
+def test_update_uses_precomputed_hidden_states_cpu():
+    """Precomputed CPU staging should be the only hidden-state cache source."""
+    cache = get_omni_pcache()
+
+    num_tokens_unpadded = 8
+    slot_offset = 8
+    slot_mapping = torch.arange(slot_offset, slot_offset + num_tokens_unpadded)
+    hidden_states = torch.zeros((num_tokens_unpadded, HIDDEN_SIZE), dtype=DTYPE)
+    hidden_states_cpu = torch.rand((num_tokens_unpadded, HIDDEN_SIZE), dtype=DTYPE)
+
+    cache.update_omni_tensor_prefix_cache(
+        hidden_states=hidden_states,
+        hidden_states_cpu=hidden_states_cpu,
+        multimodal_outputs=None,
+        num_tokens_unpadded=num_tokens_unpadded,
+        slot_mapping=slot_mapping,
+    )
+
+    hs_rows = cache.hidden_states_cache.view(NUM_BLOCKS * BLOCK_SIZE, HIDDEN_SIZE)
+    for slot_idx, new_states in zip(slot_mapping, hidden_states_cpu):
+        slot_states = hs_rows[slot_idx]
+        assert torch.all(slot_states == new_states)
+
+
+@pytest.mark.parametrize(
+    "hidden_states_cpu",
+    [
+        torch.rand((4, HIDDEN_SIZE), dtype=DTYPE),
+        torch.rand((8, HIDDEN_SIZE), dtype=OTHER_DTYPE),
+        torch.rand((8, HIDDEN_SIZE + 1), dtype=DTYPE),
+        torch.rand((HIDDEN_SIZE, 8), dtype=DTYPE).t(),
+    ],
+)
+def test_update_rejects_invalid_hidden_states_cpu(hidden_states_cpu):
+    cache = get_omni_pcache()
+
+    num_tokens_unpadded = 8
+    slot_mapping = torch.arange(num_tokens_unpadded)
+    hidden_states = torch.rand((num_tokens_unpadded, HIDDEN_SIZE), dtype=DTYPE)
+
+    with pytest.raises(RuntimeError):
+        cache.update_omni_tensor_prefix_cache(
+            hidden_states=hidden_states,
+            hidden_states_cpu=hidden_states_cpu,
+            multimodal_outputs=None,
+            num_tokens_unpadded=num_tokens_unpadded,
+            slot_mapping=slot_mapping,
+        )
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
+def test_update_rejects_gpu_hidden_states_cpu():
+    cache = get_omni_pcache()
+
+    num_tokens_unpadded = 8
+    slot_mapping = torch.arange(num_tokens_unpadded)
+    hidden_states = torch.rand((num_tokens_unpadded, HIDDEN_SIZE), dtype=DTYPE)
+    hidden_states_cpu = torch.rand((num_tokens_unpadded, HIDDEN_SIZE), dtype=DTYPE, device="cuda")
+
+    with pytest.raises(RuntimeError):
+        cache.update_omni_tensor_prefix_cache(
+            hidden_states=hidden_states,
+            hidden_states_cpu=hidden_states_cpu,
+            multimodal_outputs=None,
+            num_tokens_unpadded=num_tokens_unpadded,
+            slot_mapping=slot_mapping,
+        )
+
+
 @pytest.mark.parametrize(
     "feat_dims",
     [
@@ -359,6 +428,76 @@ def test_get_merged_hidden_states(num_tokens_padded, mocker):
     assert torch.all(req2_merged_states == req2_new_states)
 
 
+def test_get_merged_hidden_states_uses_precomputed_hidden_states_cpu(mocker):
+    cache = get_omni_pcache()
+
+    orig_num_tokens_unpadded = 8
+    slot_offset = 8
+    orig_slot_mapping = torch.arange(slot_offset, slot_offset + orig_num_tokens_unpadded)
+    orig_hidden_states = torch.rand((orig_num_tokens_unpadded, HIDDEN_SIZE), dtype=DTYPE)
+
+    cache.update_omni_tensor_prefix_cache(
+        hidden_states=orig_hidden_states,
+        multimodal_outputs=None,
+        num_tokens_unpadded=orig_num_tokens_unpadded,
+        slot_mapping=orig_slot_mapping,
+    )
+
+    num_new_toks_req1 = 3
+    num_new_toks_req2 = 2
+    cache.add_prefix_cached_new_req_id("req1")
+    num_scheduled_tokens = {
+        "req1": num_new_toks_req1,
+        "req2": num_new_toks_req2,
+    }
+    new_hidden_states = torch.zeros(
+        (num_new_toks_req1 + num_new_toks_req2, HIDDEN_SIZE),
+        dtype=DTYPE,
+    )
+    hidden_states_cpu = torch.rand_like(new_hidden_states)
+    req1_new_states = hidden_states_cpu[:num_new_toks_req1]
+    req2_new_states = hidden_states_cpu[-num_new_toks_req2:]
+    input_batch = MockInputBatch(num_computed_tokens_cpu=torch.Tensor([orig_num_tokens_unpadded, 0]))
+
+    mocker.patch(
+        "vllm_omni.core.prefix_cache.OmniTensorPrefixCache._get_cached_block_ids",
+        new=fake_get_cached_block_ids,
+    )
+    merged_states = cache.get_merged_hidden_states(
+        query_start_loc=[0, num_new_toks_req1],
+        input_batch=input_batch,
+        hidden_states=new_hidden_states,
+        hidden_states_cpu=hidden_states_cpu,
+        num_scheduled_tokens=num_scheduled_tokens,
+    )
+
+    assert torch.all(merged_states["req1"][:orig_num_tokens_unpadded] == orig_hidden_states)
+    assert torch.all(merged_states["req1"][-num_new_toks_req1:] == req1_new_states)
+    assert torch.all(merged_states["req2"] == req2_new_states)
+
+
+def test_get_merged_hidden_states_rejects_short_hidden_states_cpu(mocker):
+    cache = get_omni_pcache()
+    num_scheduled_tokens = {
+        "req1": 3,
+        "req2": 2,
+    }
+    input_batch = MockInputBatch(num_computed_tokens_cpu=torch.Tensor([0, 0]))
+
+    mocker.patch(
+        "vllm_omni.core.prefix_cache.OmniTensorPrefixCache._get_cached_block_ids",
+        new=fake_get_cached_block_ids,
+    )
+    with pytest.raises(RuntimeError):
+        cache.get_merged_hidden_states(
+            query_start_loc=[0, 3],
+            input_batch=input_batch,
+            hidden_states=torch.rand((5, HIDDEN_SIZE), dtype=DTYPE),
+            hidden_states_cpu=torch.rand((4, HIDDEN_SIZE), dtype=DTYPE),
+            num_scheduled_tokens=num_scheduled_tokens,
+        )
+
+
 @pytest.mark.parametrize("num_tokens_padded", [None, 16])
 @pytest.mark.parametrize(
     "feat_dims",

tests/worker/test_gpu_ar_model_runner.py

@@ -1,8 +1,9 @@
 from types import SimpleNamespace
 
 import pytest
+import torch
 
-from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner
+from vllm_omni.worker.gpu_ar_model_runner import ExecuteModelState, GPUARModelRunner
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -41,3 +42,83 @@ def test_resolve_pooler_payload_req_ids_downstream_stage_uses_filtered_requests(
 
     assert engine_output_type == "latent"
     assert payload_req_ids == ["r2"]
+
+
+@pytest.mark.parametrize("query_start_loc_attr", ["method", "tensor_attr"])
+def test_sample_tokens_tail_only_prefix_cache_uses_staged_cpu_hidden_states(monkeypatch, query_start_loc_attr):
+    runner = object.__new__(GPUARModelRunner)
+    runner.execute_model_state = ExecuteModelState(
+        SimpleNamespace(
+            total_num_scheduled_tokens=3,
+            num_scheduled_tokens={"r1": 1, "r2": 2},
+        ),
+        None,
+        None,
+        None,
+        torch.zeros((3, 2), dtype=torch.float32),
+        torch.tensor([[1.0, 10.0], [2.0, 20.0], [3.0, 30.0]]),
+        None,
+        None,
+        None,
+        None,
+        {},
+        None,
+    )
+    runner.kv_connector_output = None
+    runner.input_batch = SimpleNamespace(
+        req_ids=["r1", "r2"],
+        req_id_to_index={"r1": 0, "r2": 1},
+        sampling_metadata=SimpleNamespace(no_penalties=True),
+        vocab_size=10,
+        num_tokens_no_spec=None,
+    )
+    query_start_loc = torch.tensor([0, 1], dtype=torch.long)
+    if query_start_loc_attr == "method":
+        runner.query_start_loc = query_start_loc
+    else:
+        runner.query_start_loc = SimpleNamespace(cpu=query_start_loc)
+    runner.omni_prefix_cache = object()
+    runner.speculative_config = None
+    runner.routed_experts_initialized = False
+    runner.requests = {}
+    runner.supports_mm_inputs = False
+    runner.use_async_scheduling = False
+    runner._omni_num_scheduled_tokens_np = None
+
+    monkeypatch.setattr(
+        GPUARModelRunner, "_sample", lambda self, logits, spec_decode_metadata: SimpleNamespace(sampled_token_ids=[])
+    )
+    monkeypatch.setattr(GPUARModelRunner, "_update_states_after_model_execute", lambda *args, **kwargs: None)
+    monkeypatch.setattr(
+        GPUARModelRunner,
+        "_bookkeeping_sync",
+        lambda *args, **kwargs: (
+            0,
+            None,
+            [],
+            None,
+            ["r1", "r2"],
+            {"r1": 0, "r2": 1},
+            [],
+        ),
+    )
+    monkeypatch.setattr(GPUARModelRunner, "eplb_step", lambda self: None)
+    monkeypatch.setattr(GPUARModelRunner, "_resolve_pooler_payload_req_ids", lambda self, req_ids: ("audio", req_ids))
+    monkeypatch.setattr(GPUARModelRunner, "_deferred_prefix_cache_mm_keys", lambda self: set())
+    monkeypatch.setattr(GPUARModelRunner, "_model_needs_full_prefix_hidden_states", lambda self: False)
+    monkeypatch.setattr(
+        GPUARModelRunner,
+        "_maybe_get_combined_prefix_cache_tensors",
+        lambda *args, **kwargs: (None, None),
+    )
+    monkeypatch.setattr(GPUARModelRunner, "_process_additional_information_updates", lambda *args, **kwargs: None)
+    monkeypatch.setattr(GPUARModelRunner, "_should_accumulate_full_payload_output", lambda self: False)
+    monkeypatch.setattr(GPUARModelRunner, "get_omni_connector_output", lambda self: None)
+
+    output = GPUARModelRunner.sample_tokens(runner, grammar_output=None)
+
+    assert torch.equal(output.pooler_output[0]["hidden"], torch.tensor([[1.0, 10.0]]))
+    assert torch.equal(
+        output.pooler_output[1]["hidden"],
+        torch.tensor([[2.0, 20.0], [3.0, 30.0]]),
+    )

tests/worker/test_gpu_generation_model_runner.py

@@ -1,7 +1,10 @@
 import pytest
 import torch
 
-from vllm_omni.worker.gpu_generation_model_runner import GPUGenerationModelRunner
+from vllm_omni.worker.gpu_generation_model_runner import (
+    ExecuteModelState,
+    GPUGenerationModelRunner,
+)
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -16,7 +19,8 @@ class _DummyInputBatch:
 
 def _make_runner(multimodal_outputs):
     runner = object.__new__(GPUGenerationModelRunner)
-    runner.execute_model_state = (
+    runner.execute_model_state = ExecuteModelState(
+        None,
         None,
         None,
         None,

vllm_omni/core/prefix_cache.py

@@ -109,6 +109,35 @@ class OmniTensorPrefixCache:
         """Convert GPU tensors -> contiguous CPU tensors if needed."""
         return maybe_gpu_tensor.detach().cpu().contiguous()
 
+    @staticmethod
+    def _resolve_hidden_states_cpu(
+        hidden_states: torch.Tensor,
+        num_tokens: int,
+        hidden_states_cpu: torch.Tensor | None,
+    ) -> torch.Tensor:
+        if hidden_states_cpu is None:
+            return OmniTensorPrefixCache._coerce_to_cpu_tensor(hidden_states[:num_tokens])
+        if hidden_states_cpu.device.type != "cpu":
+            raise RuntimeError("hidden_states_cpu must be a CPU tensor.")
+        if not hidden_states_cpu.is_contiguous():
+            raise RuntimeError("hidden_states_cpu must be contiguous.")
+        if hidden_states_cpu.dtype != hidden_states.dtype:
+            raise RuntimeError(
+                "hidden_states_cpu has an incompatible dtype: "
+                f"got {hidden_states_cpu.dtype}, expected {hidden_states.dtype}."
+            )
+        if hidden_states_cpu.shape[1:] != hidden_states.shape[1:]:
+            raise RuntimeError(
+                "hidden_states_cpu has an incompatible feature shape: "
+                f"got {tuple(hidden_states_cpu.shape[1:])}, expected {tuple(hidden_states.shape[1:])}."
+            )
+        if hidden_states_cpu.shape[0] < num_tokens:
+            raise RuntimeError(
+                "hidden_states_cpu does not cover the requested hidden states "
+                f"slice: got {hidden_states_cpu.shape[0]} tokens, need {num_tokens}."
+            )
+        return hidden_states_cpu[:num_tokens]
+
     def update_omni_tensor_prefix_cache(
         self,
         hidden_states: torch.Tensor | None,
@@ -117,6 +146,7 @@ class OmniTensorPrefixCache:
         slot_mapping: torch.Tensor,
         num_tokens_padded: int | None = None,
         skip_mm_cache_keys: set[str] | None = None,
+        hidden_states_cpu: torch.Tensor | None = None,
     ):
         """Updates the hidden cache state for the provided hidden states and multimodal outputs.
 
@@ -127,6 +157,9 @@ class OmniTensorPrefixCache:
             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
+            hidden_states_cpu: Optional pre-staged CPU view of hidden_states.
+                When provided, it must be contiguous, live on CPU, match the
+                feature shape of hidden_states, and cover num_tokens_unpadded.
         """
         unpadded_slot_mapping = slot_mapping[:num_tokens_unpadded]
         if num_tokens_padded is None:
@@ -135,9 +168,11 @@ class OmniTensorPrefixCache:
 
         if hidden_states is not None:
             # Slice to unpadded portion before caching
-            hidden_states = hidden_states[:num_tokens_unpadded]
-            # Ensure that hidden states are on the CPU
-            hidden_states = OmniTensorPrefixCache._coerce_to_cpu_tensor(hidden_states)
+            hidden_states = OmniTensorPrefixCache._resolve_hidden_states_cpu(
+                hidden_states,
+                num_tokens_unpadded,
+                hidden_states_cpu,
+            )
             # View the cache as 2D so that we can treat our slots as row indices
             flat_cache = self.hidden_states_cache.view(-1, self.hidden_states_cache.shape[-1])
             flat_cache[unpadded_slot_mapping] = hidden_states
@@ -388,11 +423,27 @@ class OmniTensorPrefixCache:
             )
         return combined_multimodal_outputs
 
-    def get_merged_hidden_states(self, *args, **kwargs) -> dict[str, torch.Tensor]:
-        """Get the merged hidden states."""
+    def get_merged_hidden_states(
+        self,
+        query_start_loc: torch.Tensor,
+        input_batch: InputBatch,
+        hidden_states: torch.Tensor,
+        num_scheduled_tokens: dict[str, int],
+        hidden_states_cpu: torch.Tensor | None = None,
+    ) -> dict[str, torch.Tensor]:
+        """Get merged hidden states, optionally reusing pre-staged CPU states.
+
+        When provided, hidden_states_cpu follows the same contract as
+        update_omni_tensor_prefix_cache: CPU, contiguous, same dtype and feature
+        shape as hidden_states, and covering every scheduled-token span derived
+        from query_start_loc and num_scheduled_tokens.
+        """
         return self._get_merged_tensors(
-            *args,
-            **kwargs,
+            query_start_loc=query_start_loc,
+            input_batch=input_batch,
+            hidden_states=hidden_states,
+            num_scheduled_tokens=num_scheduled_tokens,
+            staged_cpu_tensor=hidden_states_cpu,
             cache=self.hidden_states_cache,
         )
 
@@ -403,12 +454,14 @@ class OmniTensorPrefixCache:
         cache: torch.Tensor,
         hidden_states: torch.Tensor,
         num_scheduled_tokens: dict[str, int],
+        staged_cpu_tensor: torch.Tensor | None = None,
     ) -> dict[str, torch.Tensor]:
         """When hidden state caching is enabled, takes the input hidden_states,
         which only correspond to the scheduled tokens, and returns a mapping
         from request IDs to their full hidden states. This is accomplished by
         looking up the block IDs & scheduled token counts to split the
-        hidden_states.
+        hidden_states. staged_cpu_tensor can supply the already materialized CPU
+        scheduled-token view; None preserves the legacy per-call CPU conversion.
         """
         # We do not support hybrid caches at the moment.
         if len(input_batch.block_table.block_tables) > 1:
@@ -418,7 +471,17 @@ class OmniTensorPrefixCache:
             )
 
         combined_hidden_states = {}
-        hidden_states = OmniTensorPrefixCache._coerce_to_cpu_tensor(hidden_states)
+        required_tokens = 0
+        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]
+            required_tokens = max(required_tokens, int(end))
+        hidden_states = OmniTensorPrefixCache._resolve_hidden_states_cpu(
+            hidden_states,
+            required_tokens,
+            staged_cpu_tensor,
+        )
         for req_id in input_batch.req_ids:
             req_idx = input_batch.req_id_to_index[req_id]
 

vllm_omni/entrypoints/openai/api_server.py

@@ -17,6 +17,7 @@ from collections.abc import AsyncIterator
 from contextlib import asynccontextmanager
 from http import HTTPStatus
 from numbers import Integral
+from pathlib import Path
 from typing import Annotated, Any, Literal, cast
 
 import httpx
@@ -143,6 +144,51 @@ MAX_UINT32_SEED = 2**32 - 1
 profiler_router = APIRouter()
 
 
+def _load_model_chat_template_json(model: str) -> str | None:
+    """Load a model-level chat_template.json from a local path or HF cache.
+
+    Some multimodal HF repos, including Qwen3-Omni, ship the chat template as a
+    separate file instead of embedding it in tokenizer_config.json. Transformers
+    4.44+ no longer supplies a default template, so serving must pass that model
+    template explicitly when the user did not provide --chat-template.
+    """
+    candidate = Path(model) / "chat_template.json"
+    template_path: str | None = str(candidate) if candidate.is_file() else None
+
+    if template_path is None:
+        try:
+            from huggingface_hub import hf_hub_download
+
+            template_path = hf_hub_download(
+                repo_id=model,
+                filename="chat_template.json",
+                local_files_only=True,
+            )
+        except Exception:
+            return None
+
+    try:
+        with open(template_path, encoding="utf-8") as f:
+            payload = json.load(f)
+    except Exception as exc:
+        logger.warning("Failed to load chat template from %s: %s", template_path, exc)
+        return None
+
+    if isinstance(payload, dict):
+        template = payload.get("chat_template")
+    elif isinstance(payload, str):
+        template = payload
+    else:
+        template = None
+
+    if not isinstance(template, str) or not template.strip():
+        logger.warning("Ignoring malformed chat template payload in %s", template_path)
+        return None
+
+    logger.info("Loaded chat template from %s", template_path)
+    return template
+
+
 def _should_enable_profiler_endpoints(stage_configs: list | None) -> bool:
     """Check if any stage has profiler_config set in its engine_args."""
     if not stage_configs:
@@ -673,6 +719,14 @@ async def omni_init_app_state(
     logger.info("Supported tasks: %s", supported_tasks)
 
     resolved_chat_template = load_chat_template(args.chat_template)
+    if resolved_chat_template is None:
+        try:
+            tokenizer = await engine_client.get_tokenizer()
+        except Exception as exc:
+            logger.debug("Could not inspect tokenizer chat_template before serving init: %s", exc)
+            tokenizer = None
+        if tokenizer is None or getattr(tokenizer, "chat_template", None) is None:
+            resolved_chat_template = _load_model_chat_template_json(args.model)
 
     if args.tool_server == "demo":
         tool_server: ToolServer | None = DemoToolServer()

vllm_omni/worker/gpu_ar_model_runner.py

@@ -52,6 +52,7 @@ class ExecuteModelState(NamedTuple):
     spec_decode_metadata: Any
     spec_decode_common_attn_metadata: Any
     hidden_states: torch.Tensor
+    hidden_states_cpu: torch.Tensor | None
     sample_hidden_states: torch.Tensor
     aux_hidden_states: list[torch.Tensor] | None
     ec_connector_output: Any
@@ -282,6 +283,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
     def _maybe_update_prefix_cache(
         self,
         hidden_states: torch.Tensor,
+        hidden_states_cpu: torch.Tensor | None,
         multimodal_outputs: dict,
         num_tokens_unpadded: int,
         num_tokens_padded: int,
@@ -292,7 +294,10 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
         """
         # Cache hidden states if we've enabled hidden state prefix caching
         # unless this isn't the last pipeline parallelism rank.
-        if self.omni_prefix_cache is not None and get_pp_group().is_last_rank:
+        is_last_pp_rank = get_pp_group().is_last_rank
+        if hidden_states_cpu is not None and not is_last_pp_rank:
+            raise RuntimeError("hidden_states_cpu staging is only valid on the last pipeline parallel rank.")
+        if self.omni_prefix_cache is not None and is_last_pp_rank:
             # If this happens, it generally means the model is not following the correct
             # interface yet and is therefore currently not compatible with prefix cache.
             if multimodal_outputs is not None and not isinstance(multimodal_outputs, dict):
@@ -314,11 +319,13 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
                 slot_mapping=slot_mapping_cpu,
                 num_tokens_padded=num_tokens_padded,
                 skip_mm_cache_keys=self._deferred_prefix_cache_mm_keys(),
+                hidden_states_cpu=hidden_states_cpu,
             )
 
     def _maybe_get_combined_prefix_cache_tensors(
         self,
         hidden_states: torch.Tensor,
+        hidden_states_cpu: torch.Tensor | None,
         multimodal_outputs: dict,
         num_scheduled_tokens: dict[str, int],
     ) -> tuple[dict[str, torch.Tensor] | None, dict | None]:
@@ -328,7 +335,12 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
         # Prior to applying the post-processing func, extract
         # the prefix cached hidden states and multimodal states.
         combined_hidden_states, combined_multimodal_outputs = None, None
+        is_last_pp_rank = get_pp_group().is_last_rank
+        if hidden_states_cpu is not None and not is_last_pp_rank:
+            raise RuntimeError("hidden_states_cpu staging is only valid on the last pipeline parallel rank.")
         if self.omni_prefix_cache is not None:
+            if not is_last_pp_rank:
+                raise RuntimeError("Omni prefix-cache tensor merge is only valid on the last pipeline parallel rank.")
             if (
                 not self._model_needs_full_prefix_hidden_states()
                 and not self.omni_prefix_cache.has_prefix_cached_new_req_ids()
@@ -339,6 +351,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
                     query_start_loc=self.query_start_loc.cpu,
                     input_batch=self.input_batch,
                     hidden_states=hidden_states,
+                    hidden_states_cpu=hidden_states_cpu,
                     num_scheduled_tokens=num_scheduled_tokens,
                 )
             combined_multimodal_outputs = self.omni_prefix_cache.get_merged_multimodal_states(
@@ -662,11 +675,15 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
                 aux_hidden_states = None
 
             hidden_states, multimodal_outputs = self.extract_multimodal_outputs(model_output)
+            hidden_states_cpu = None
+            if self.omni_prefix_cache is not None and get_pp_group().is_last_rank:
+                hidden_states_cpu = hidden_states[:num_tokens_unpadded].detach().to("cpu").contiguous()
 
             # Cache hidden states & multimodal outputs if we've enabled hidden state
             # prefix caching unless this isn't the last pipeline parallelism rank.
             self._maybe_update_prefix_cache(
                 hidden_states=hidden_states,
+                hidden_states_cpu=hidden_states_cpu,
                 multimodal_outputs=multimodal_outputs,
                 num_tokens_unpadded=num_tokens_unpadded,
                 num_tokens_padded=num_tokens_padded,
@@ -737,6 +754,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
             spec_decode_metadata,
             spec_decode_common_attn_metadata,
             hidden_states,
+            hidden_states_cpu,
             sample_hidden_states,
             aux_hidden_states,
             ec_connector_output,
@@ -838,6 +856,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
             spec_decode_metadata,
             spec_decode_common_attn_metadata,
             hidden_states,
+            staged_hidden_states_cpu,
             sample_hidden_states,
             aux_hidden_states,
             ec_connector_output,
@@ -963,7 +982,14 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
         downstream_req_id_set = set(downstream_req_ids)
         hidden_states_cpu = None
         req_hidden_states_cpu: dict[str, torch.Tensor] | None = None
-        if needs_pooler_payload:
+        needs_scheduled_hidden_payload = needs_pooler_payload and (
+            self.omni_prefix_cache is None or not self._model_needs_full_prefix_hidden_states()
+        )
+        if needs_scheduled_hidden_payload and self.omni_prefix_cache is not None:
+            if staged_hidden_states_cpu is None:
+                raise RuntimeError("Prefix-cache hidden-state payload requires staged CPU hidden states.")
+            hidden_states_cpu = staged_hidden_states_cpu
+        elif needs_scheduled_hidden_payload:
             num_valid_tokens = min(
                 int(scheduler_output.total_num_scheduled_tokens),
                 int(hidden_states.shape[0]),
@@ -980,6 +1006,8 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
                 dtype=np.int32,
             )
         query_start_loc_cpu = self.query_start_loc.cpu
+        if callable(query_start_loc_cpu):
+            query_start_loc_cpu = 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:
@@ -1004,6 +1032,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
                     combined_multimodal_outputs,
                 ) = self._maybe_get_combined_prefix_cache_tensors(
                     hidden_states,
+                    staged_hidden_states_cpu,
                     multimodal_outputs,
                     scheduler_output.num_scheduled_tokens,
                 )

vllm_omni/worker/gpu_generation_model_runner.py

@@ -355,6 +355,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin
             None,
             None,
             None,
+            None,
             ec_connector_output,
             cudagraph_stats,
             multimodal_outputs,
@@ -398,6 +399,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin
             spec_decode_metadata,
             spec_decode_common_attn_metadata,
             hidden_states,
+            _hidden_states_cpu,
             sample_hidden_states,
             aux_hidden_states,
             ec_connector_output,