diff --git a/tests/core/test_prefix_cache.py b/tests/core/test_prefix_cache.py index 326eb6b34..2111556f2 100644 --- a/tests/core/test_prefix_cache.py +++ b/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", diff --git a/tests/worker/test_gpu_ar_model_runner.py b/tests/worker/test_gpu_ar_model_runner.py index 87af8982e..c36600be8 100644 --- a/tests/worker/test_gpu_ar_model_runner.py +++ b/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]]), + ) diff --git a/tests/worker/test_gpu_generation_model_runner.py b/tests/worker/test_gpu_generation_model_runner.py index 15880d7a9..7bc0f0ba5 100644 --- a/tests/worker/test_gpu_generation_model_runner.py +++ b/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, diff --git a/vllm_omni/core/prefix_cache.py b/vllm_omni/core/prefix_cache.py index f80074173..79f0990fd 100644 --- a/vllm_omni/core/prefix_cache.py +++ b/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] diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index d37a2374a..daeb22cbd 100644 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/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() diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index 6f959cf46..b267b8948 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/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, ) diff --git a/vllm_omni/worker/gpu_generation_model_runner.py b/vllm_omni/worker/gpu_generation_model_runner.py index 1dcdd31b7..e4b7944d1 100644 --- a/vllm_omni/worker/gpu_generation_model_runner.py +++ b/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,