From 470ec60848fa98e030298706cceff88ae8f691a8 Mon Sep 17 00:00:00 2001 From: Zheng Wengang Date: Sat, 19 Sep 2026 14:23:59 +0100 Subject: [PATCH] [Core] Refactor omni prefix cache into Manager/Controller (#6654) Signed-off-by: ZhengWG Co-authored-by: amy-why-3459 Co-authored-by: Hongsheng Liu --- .../omni_async_output_materialization.md | 25 +- docs/design/feature/prefix_caching.md | 193 +- tests/core/test_prefix_cache.py | 2672 +++++++++++------ tests/core/test_prefix_cache_async_write.py | 119 - tests/core/test_prefix_cache_runner_mixin.py | 182 ++ tests/e2e/online_serving/test_qwen3_omni.py | 112 +- ...st_runner_assisted_unified_graph_static.py | 4 +- .../npu/test_npu_prefix_cache_wiring.py | 158 + .../test_fa3_scheduler_metadata_size.py | 12 +- tests/worker/test_gpu_ar_model_runner.py | 41 +- vllm_omni/core/prefix_cache.py | 737 ----- vllm_omni/core/prefix_cache/__init__.py | 58 + vllm_omni/core/prefix_cache/block_pool.py | 84 + vllm_omni/core/prefix_cache/controller.py | 976 ++++++ vllm_omni/core/prefix_cache/group_view.py | 208 ++ vllm_omni/core/prefix_cache/interface.py | 171 ++ vllm_omni/core/prefix_cache/manager.py | 1475 +++++++++ vllm_omni/core/prefix_cache/runner_mixin.py | 148 + vllm_omni/deploy/qwen3_tts.yaml | 11 +- .../npu/worker/npu_ar_model_runner.py | 124 +- .../npu/worker/npu_generation_model_runner.py | 1 + .../platforms/npu/worker/npu_model_runner.py | 55 +- vllm_omni/worker/gpu_ar_model_runner.py | 240 +- .../worker/gpu_generation_model_runner.py | 1 + vllm_omni/worker/gpu_model_runner.py | 52 +- 25 files changed, 5710 insertions(+), 2149 deletions(-) delete mode 100644 tests/core/test_prefix_cache_async_write.py create mode 100644 tests/core/test_prefix_cache_runner_mixin.py create mode 100644 tests/platforms/npu/test_npu_prefix_cache_wiring.py delete mode 100644 vllm_omni/core/prefix_cache.py create mode 100644 vllm_omni/core/prefix_cache/__init__.py create mode 100644 vllm_omni/core/prefix_cache/block_pool.py create mode 100644 vllm_omni/core/prefix_cache/controller.py create mode 100644 vllm_omni/core/prefix_cache/group_view.py create mode 100644 vllm_omni/core/prefix_cache/interface.py create mode 100644 vllm_omni/core/prefix_cache/manager.py create mode 100644 vllm_omni/core/prefix_cache/runner_mixin.py diff --git a/docs/design/feature/omni_async_output_materialization.md b/docs/design/feature/omni_async_output_materialization.md index 50a2e7c52..eff3024d8 100644 --- a/docs/design/feature/omni_async_output_materialization.md +++ b/docs/design/feature/omni_async_output_materialization.md @@ -318,18 +318,18 @@ The same behavior applies to the Base and VoiceDesign Qwen3-TTS checkpoints because they use the same Talker implementation. !!! warning - Do not pass `--no-async-chunk` or enable prefix caching when you want this - optimization. Either change causes the runner to fall back to synchronous - output construction. Enabling async chunk alone cannot activate the feature - for a model that has not opted into async Omni output. + Do not pass `--no-async-chunk` when you want this optimization. Disabling + async chunk causes the runner to fall back to synchronous output + construction. Enabling async chunk alone cannot activate the feature for a + model that has not opted into async Omni output. !!! note "Prefix cache compatibility" - Async Omni output materialization and Omni prefix caching cannot currently - run together. `_should_use_async_omni_output()` returns `False` whenever - `self.omni_prefix_cache` is present, so prefix caching selects synchronous - output materialization. Supporting both features requires snapshotting or - otherwise synchronizing prefix-cache merge and update state before the - background output path can consume it safely. + Prefix cache and async Omni output can run together. `save_outputs` returns + a `step_id`; the output builder later calls `materialize(sid, req_ids)` with + the save-time request list, which may be a step late relative to the live + `input_batch`. `_should_use_async_omni_output()` does not disable itself + when the cache is present. See + [Automatic Prefix Caching](prefix_caching.md#implementation). ## Compatibility and Fallbacks @@ -341,10 +341,9 @@ runtime conditions hold: | AR async scheduling is enabled | The optimization relies on the scheduler advancing while the prior output is materialized | | `async_chunk` is enabled | The feature targets incremental downstream Omni payloads | | The model stage opts in with `use_async_omni_output` | Models must declare that their output lifecycle is safe to defer | -| Omni prefix cache is disabled | Prefix-cache merge and update ordering currently requires synchronous materialization | | Speculative decoding is disabled | Speculative output state is not included in this deferred path | | Routed-expert output is disabled | Routed-expert extraction currently requires the synchronous path | -| Postprocess is absent or explicitly runs eagerly | State needed by the next decode step must be updated before the runner returns | +| Postprocess is absent or explicitly runs eagerly | State needed by the next decode step must be updated before the runner returns. Eager postprocess sees the scheduled slice (`None, None` combined tensors), not the prefix-cache-merged full prompt; it must only consume the tail | These checks are evaluated per stage on runners based on `GPUARModelRunner`. An unsupported combination does not prevent serving; it only falls back to @@ -380,6 +379,8 @@ so Omni payload materialization remains synchronous. - `tests/worker/test_gpu_ar_model_runner.py`: Snapshot, guard, connector ordering, and background error-propagation tests. - [Async Chunk](async_chunk.md): Inter-stage chunking and scheduling design. +- [Automatic Prefix Caching](prefix_caching.md#implementation): Prefix-cache + `step_id` / `materialize` contract used by the deferred output builder. - [Qwen3-Omni optimization blog](https://vllm.ai/blog/2026-07-01-qwen3-omni-optimization): Optimization context and controlled performance results. - [PR #4476](https://github.com/vllm-project/vllm-omni/pull/4476): Feature diff --git a/docs/design/feature/prefix_caching.md b/docs/design/feature/prefix_caching.md index a8fe4fc50..5aeef0f77 100644 --- a/docs/design/feature/prefix_caching.md +++ b/docs/design/feature/prefix_caching.md @@ -1,6 +1,5 @@ # Automatic Prefix Caching in Omni Models - --- ## Table of Contents @@ -9,6 +8,8 @@ - [High-Level Approach](#high-level-approach) - [Example](#example) - [What About Multimodal Inputs?](#what-about-multimodal-inputs) +- [Implementation](#implementation) +- [Related Files](#related-files) --- @@ -27,6 +28,7 @@ vLLM implements automatic prefix caching for managing its kv-cache, which is bes This document describes vLLM-Omni's mechanism for caching tensor outputs that are meant to be passed between stages, when requests have common prefixes, similar to the way in which vLLM has prefix caching for the kv-cache. This works in conjunction with vLLM's multimodal encoder caching, but is distinct. See the final section for a concrete example for how they tie together in practice. ### High-Level Approach + !!! note "Note 2" Prior to reading this section, it's recommended to take a look at the design documents in vLLM for [Automatic Prefix Caching](https://docs.vllm.ai/en/latest/features/automatic_prefix_caching/), which will make some of the concepts more clear. @@ -38,8 +40,10 @@ The main focus of vLLM-Omni's approach to prefix caching stage outputs is to bui With this in mind, consider the set of blocks in a 2D layout, where the row represents the index of blocks being considered, and the columns represent the slots corresponding to tokens within each block. Since we know the `num_blocks` and `block_size` from our kv cache config, if we want to cache a tensor with feature size `D`, we can preallocate a CPU tensor of size `(num_blocks, block_size, D)`, and use the same block index and slot mapping to retrieve the corresponding feature vector. +Host footprint: each cached key costs `num_blocks × block_size × D × dtype_bytes` of **pinned** CPU memory (pinned so device→host can overlap compute), allocated on the first `save_outputs` that sees the key — the first real request pays the `cudaHostAlloc`. Measured on a Qwen3-Omni deployment: thinker `__hidden_states__` `[15092, 16, 2048]` bf16 ≈ 0.92 GiB (plus the same again for each `hidden_states.layer_*` key a model exposes), Qwen3-TTS talker `codes.audio` `[16180, 16, 16]` ≈ 33 MiB. Budget host RAM for the stage accordingly. ### Example + !!! note "Note 3" Prefix caching in vLLM-Omni currently is only supported on AutoRegressive stages with one kv-cache group. Configure it with the pipeline-wide `enable_prefix_caching` field in the deploy config. @@ -56,7 +60,7 @@ The prefix cache flow is then outlined below. 2. Say we process the request `The quick brown fox was tired and slept beneath the shady tree`, which is 12 tokens and evenly divides into 3 blocks as shown below. -``` +```text [ The quick brown fox ] [ was tired and slept ] [beneath the shady tree ] Block 1: |<--- block tokens ---->| Block 2: |<------- prefix ------>| |<--- block tokens --->| @@ -65,10 +69,9 @@ Block 3: |<------------------ prefix -------------------->| |<--- block tokens - When the request processes, we inspect the multimodal outputs and identify the `mm_feature` tensor, which will be of shape `(seq_len, feature_dim)`, i.e., `(12, 16)` in this example. We note that the first axis is dependent on the `seq_len` and add a new cache_tensor of shape `(num_blocks, block_size, feature_dim)` to our multimodal cache for tensors. +1. If we lay out the cache as a 2D tensor of shape (`num_blocks`, `block_size`), we'll have something like the following: -3. If we lay out the cache as a 2D tensor of shape (`num_blocks`, `block_size`), we'll have something like the following: - -``` +```text 0: [ The quick brown fox ] 1: [ was tired and slept ] 2: [beneath the shady tree ] @@ -78,7 +81,8 @@ When the request processes, we inspect the multimodal outputs and identify the ` ``` Or, if we flatten it down to 1D, -``` + +```text 0: The 1: quick 2: brown @@ -90,7 +94,8 @@ Or, if we flatten it down to 1D, ``` which we can think of as row indices into the hidden states tensor if we view it as the 2D shape `(num_blocks x block_size, feature_dim)`. That is, the analogous flattened (from 3D -> 2D) mapping of the cache for hidden states becomes the following. -``` + +```text 0: 1: 2: @@ -103,10 +108,9 @@ which we can think of as row indices into the hidden states tensor if we view it Similarly, for the multimodal outputs cache, the flattened coordinates are the same, but the `mm_feature` maps to vectors of length `16` instead of the hidden size of `2`. Note that in practice, we may have multiple multimodal output tensors per forward pass, which may have different names and different feature dimensions. +1. Now, say that we receive a new request `The quick brown fox jumped over the dog`. -4. Now, say that we receive a new request `The quick brown fox jumped over the dog`. - -``` +```text [ The quick brown fox ] [ jumped over the dog ] Block 1: |<--- block tokens ---->| Block 2: |<------- prefix ------>| |<--- block tokens --->| @@ -116,7 +120,7 @@ Here, we will have a cache hit for `Block 1` which will be detected by vLLM base Since we have the block indices / slot mappings from the kv cache manager, we can simply mirror the mappings and leverage the same indices for the cached hidden states and multimodal outputs. This allows us to look up the correct tensors from our externally maintained 3D caches. -``` +```text 0: [ The quick brown fox ] < already in the cache 1: [ was tired and slept ] 2: [beneath the shady tree ] @@ -129,11 +133,11 @@ Since we have the block indices / slot mappings from the kv cache manager, we ca Finally, to pass the full hidden states and multimodal outputs to the next stage, we simply concatenate the cached contents with the corresponding new tensors computed from the current forward call. - ### What About Multimodal Inputs? + It's also useful to consider the case about how Omni prefix caching is handled when we have multimodal inputs that don't cleanly end on block boundaries, as well as how this works with multimodal encoder caching in vLLM. For example: -``` +```text [ Im0 Im1 Im2 Im3 ] [ Im4 Im5 foo ] Block 1: |<--- block tokens ---->| Block 2: |<------- prefix ------>| |<--- block tokens --->| @@ -147,9 +151,9 @@ In reality, this isn't a big problem for correctness, because vLLM also maintain - The hash describing the image data starting at position 0 and with length 6 - In vLLM's encoder cache, a mapping from the image hash above to the encoder output - To understand what happens, say we get the following input as a second request: -``` + +```text [ Im0 Im1 Im2 Im3 ] [ Im4 Im5 bar baz ] Block 1: |<--- block tokens ---->| Block 2: |<------- prefix ------>| |<--- block tokens --->| @@ -162,3 +166,162 @@ Because we have multimodal data in a scheduled span that isn't fully precomputed When we pass our multimodal tensors to the language model component in the same stage, we'll then expect the same outputs, because the prefix caching behaviors in vLLM-Omni / vLLM match, so the LLM will use vLLM's KV cache manager's prefix caching to correctly handle the attention information for `Block 1` while calculating the outputs for `Block 2`, giving us the correct results for processing `Block 2` with the context of `Block 1`. Finally, we look up the output hidden states/multimodal tensors corresponding to the prefix cache hit `Block 1` and concatenate it with the forward pass result to get the final result, which is expected to be identical to the full hidden states when prefix caching is disabled. + +### Implementation + +The block/slot model is `vllm_omni/core/prefix_cache/`. +`OmniPrefixCacheManager` owns slot occupancy, the request-task table, +hit spans, and merge. +`OmniPrefixCacheController` moves data: a reusable `StagingBufferPool` for +this step's device→host copy, and writes into the durable `PrefixBlockPool`. +The state lock covers those tables only. + +Miss is not an error (this step's forward slice only). A hit span whose +hidden rows are absent is fatal. For mm keys the rule is looser: a model may +emit a key only for some requests, so a hit span with no rows behind an mm key +reads zeros from the pool for those positions rather than raising. Abort still +writes: once a hash entered this step's batch it must land in the cache. + +`enable_prefix_caching` is refused on KV-consumer stages (`kv_role` of +`kv_consumer` or `kv_both`). KV received from a producer is reported as +`num_computed_tokens` too, and the manager cannot tell it from a local hit. +Producer-only stages are unaffected. Pooling stages never save, so they get +no cache. This gate is one function, `stage_prefix_cache_config`, called by +both the GPU and the NPU model runner at kv-cache init. + +Which stages may set `enable_prefix_caching: true`: + +| Stage | `enable_prefix_caching` | Why | +| --- | --- | --- | +| AR stage with one full-attention kv group whose hidden states / per-token mm feed the next stage (Qwen3-Omni thinker and talker) | supported | The case the cache is built for: full-prompt hidden states are merged from the pool on a hit. | +| AR stage that sets `requires_full_prefix_cached_hidden_states = False`, optionally with `deferred_prefix_cache_mm_keys` (Qwen3-TTS talker, Higgs v3 talker) | supported | Hidden is not cached; deferred codec rows are written once on finish. | +| Pooling stage | ignored | Never saves; the gate returns no config. | +| `kv_role: kv_consumer` / `kv_both` | refused at kv-cache init (`OmniPrefixCacheUnmatchError`) | Producer KV shows up as `num_computed_tokens` and is indistinguishable from a local hit. | +| Speculative decoding on the stage | refused at kv-cache init | Under async scheduling vLLM keeps `num_computed_tokens_cpu` optimistic (all drafts accepted) during the forward and corrects it afterwards; `step_slots_cpu` would mirror rows at the wrong slots. Not verified; refused as a whole. | +| `prefix_match_unit` smaller than `block_size` | refused at kv-cache init | Sub-block hits make `num_computed_tokens` unaligned; the hit registry only mirrors whole blocks. | +| Hybrid / sliding-window / multi-group kv cache (e.g. a talker with `attention_type: sliding_recompute`) | refused at kv-cache init | The cache mirrors exactly one full-attention block table. | +| Attention backend whose kernel block size differs from `--block-size` (FlashInfer / FlashMLA / CutlassMLA with a block size they do not list natively), or decode context parallel | refused at first step (`FullAttentionGroupView`) | `step_slots_cpu` computes `table[req, pos // block_size] * block_size + pos % block_size` over allocator block ids; hybrid kernel blocks and DCP token striping change that row layout. FlashAttention / Triton accept any multiple of 16 and never split blocks. | +| Codec decoder / Code2Wav stages (Qwen3-Omni stage 2, Qwen3-TTS stage 1) | keep `false` | Nothing downstream consumes their hidden states; the cache would only add device→host copies. Not validated. | +| Diffusion stages | n/a | No vLLM KV cache to mirror. | + +Hit spans come from `scheduled_new_reqs` only, as in the pre-refactor cache: + +- A new request with a (partial) prefix hit is the normal path: the hit + blocks are read from the pool, the rest is this step's rows, and the + divergent tail is written under its own blocks. The gather starts on the + prefetch thread at `new_step_starts`; a same-step hit (vLLM hashes blocks + at schedule time, so `b` can hit blocks `a` computes in the same forward) + cannot plan until `a`'s write is registered and starts at `save_outputs` + instead. Either way it runs before the next step can hand those blocks + to a new tenant. vLLM frees blocks one step before `finished_req_ids` + arrives, so a finished request whose last step materializes late may + still find a hit slot reassigned. Each `(slot, key)` carries a write + version, bumped whenever a new write claims it. A planned read captures + that version and is registered in `_pending_reads`; a later write that + reclaims those slots copy-on-writes the still-`COMMITTED` rows into the + ref before it overwrites the pool, so a delayed fetch serves the original + tenant. A version mismatch with no preserved copy raises for live and + finished alike (the pool rows are a newer tenant's). Production defaults + stay opt-in until preempt/resume hit spans are reconstructed. +- `async_chunk` continuation: when the next upstream chunk arrives, the same + request id re-enters `scheduled_new_reqs` with `num_computed_tokens` equal + to what it already computed itself. Ids already in `live_reqs` are skipped + for hit marking: those rows were delivered in earlier steps and re-emitting + them would duplicate output. A `delivered_upto` span for this case is + Phase 2. +- Preemption + reschedule: vLLM resets `num_computed_tokens` to 0 on + preemption and re-runs prefix matching on resume, so the resumed request + can come back with a fresh hit. With the V1 model runner it arrives + through `scheduled_cached_reqs` (id in `resumed_req_ids`, `new_block_ids` + replaces the table); with the V2 runner it re-enters `scheduled_new_reqs` + while still in `live_reqs`. Neither path marks an omni hit span: the + resumed request gets only the rows it recomputes, and its still-open + deferred write keeps appending (a slot written twice keeps the later + chunk). Cache integrity holds either way — the hit blocks already have + rows, from this request or the one it hit. Stages that need full prompt + hidden states should be sized so preemption does not occur while prefix + caching is on. Same as before this refactor; tracked for Phase 2. + +Two write paths, split by `ModelCachePolicy.deferred_keys`: + +- Immediate (`JOIN_NEXT_STEP`): save launches a whole-step device→host into + a staging slot; the committer waits that event and writes the CPU pool. + The next save waits `done` (the pool write): a reused slot must never + leave a pending pool write behind, or a delayed hit read of the old rows + would find nothing recoverable. +- Deferred (`JOIN_ON_FINISH`): mm whose first dim is this step's token count + stays on a per-request GPU clone; `_WriteChunk`s append across steps; + finish/abort (or GPU-byte-budget pressure) forces the copy. One open + WriteTask per request; a budget flush may close it mid-request, in which + case the next save opens a new one (`write_n` + 1) and a hit reads both, + so one long request cannot pin the whole budget. + +A `WriteTask` moves through `TaskState` only via `transition()`, one step at +a time along a strict chain: +`PENDING` (registered, not queued) → `QUEUED` (on a copy queue) → +`COPYING` (one thread owns the copy stage) → `HOST_READY` (host rows ready) → +`WRITTEN` (in the CPU pool); `FAILED` is reachable from any non-terminal +state. Skipping a step raises. The worker claims `COPYING` in the same `_wake` +critical section as the queue pop, so `escalate` never re-queues a task it can +see is already claimed. `host_ready` / `done` are wait primitives set by the +transitions into `HOST_READY` / `WRITTEN` / `FAILED`. GPU-clone bytes are +charged once per clone on a `_BudgetTicket` pinned by every task that views +it, and uncharged when the last holder releases (idempotent per tid). + +Mm whose first dim equals the unpadded scheduled length *or* the CUDA-graph +padded length is registered on first sighting. Talker `codes.audio` is a +cat of scheduled rows and stays unpadded while hidden is padded; both must +open a pool key, whether immediate or deferred. +Leftover mm (lists, `codes.ref`, any tensor whose first dim is not this +step's token count) is copied to CPU at save without truncating that first +dim. That leftover copy is the async-builder read replica for this step; it +does not write the pool or carry abort/preempt occupancy. + +```python +cache.register_policy(ModelCachePolicy.from_model(model)) # load_model +cache.new_step_starts(scheduler_output) # before _update_states +sid = cache.save_outputs(hidden, mm_outputs, num_tokens_unpadded=n, + num_tokens_padded=n_pad) +outs = cache.materialize(sid, req_ids) # or discard_step(sid) +``` + +Each step id is consumed exactly once. `req_ids` must be a subset of the save +snapshot. At most `staging_depth` unused step ids may exist at once: every +`save_outputs` claims one staging slot, including saves with only leftover +mm that copy no device→host page. A slot is also held by each immediate +write that views it, until the committer's pool write. A later save waits +for `materialize`/`discard_step` or that pool write to free a slot; +`staging_claim_timeout_s` then errors with the unused ids and the task count. +`join`/`join_host_ready` use the same bound and raise with the stuck task's +id and state instead of hanging the caller. +`materialize` may run on the async output builder after the engine has +entered the next step; leftover mm (not written to the pool) is copied +to CPU at `save_outputs` so the builder never reads live CUDA-graph +buffers. See +[Async Omni Output Materialization](omni_async_output_materialization.md). + +The cache is constructed only on the last pipeline-parallel rank +(`_ensure_omni_prefix_cache`). Other ranks skip it: they never call +`save_outputs`, so a hit table there would resolve to absent slots. + +Threads, locks, and what each may block on: + +| Thread | Role | May block on | Must not hold while blocked | +| --- | --- | --- | --- | +| Engine | `new_step_starts`, `save_outputs` | previous-step `join` (`done`) / finished-write `join_host_ready`; `reserve()` GPU-byte flush; staging-slot claim; `dispatch()` / finish-abort `escalate()` (eager mode: the copy + pool write run inline) | `_state_lock` | +| Async output builder | `materialize` (may overlap the next engine step) | this step's `step_d2h_event`; `join` (`done`); deferred `fetch_host` | `_state_lock` | +| Committer | `_worker_loop`: wait device→host / deferred copy / pool write | `_wake.wait`; `step_d2h_event` or copy-stream sync | never takes `_state_lock` | +| Prefetch pool | hit-span gather during forward | `join` (`done`); deferred `fetch_host` | `_state_lock` | + +| Lock | Covers | Does not cover | +| --- | --- | --- | +| manager `_state_lock` | occupancy tables, step contexts, hit spans, task registration, pool-key publish (`install_key`) | join, GPU-byte flush, copy, `step_d2h_event` wait, pool-key allocation, eager `dispatch()` / `escalate()` | +| controller `_lock` / `_wake` | task registry, queues, GPU-clone byte budget | device→host / pool-write body (released before `synchronize`) | +| `WriteTask.lock` | `state`, `reassigned`, `append_chunk` | waiting on `host_ready` / `done` (those are events) | + +### Related Files + +- `vllm_omni/core/prefix_cache/` +- `vllm_omni/worker/gpu_model_runner.py` (`_ensure_omni_prefix_cache`) +- `tests/core/test_prefix_cache.py` +- [Async Omni Output Materialization](omni_async_output_materialization.md) diff --git a/tests/core/test_prefix_cache.py b/tests/core/test_prefix_cache.py index 4f91a95b7..171f36f71 100644 --- a/tests/core/test_prefix_cache.py +++ b/tests/core/test_prefix_cache.py @@ -1,917 +1,1801 @@ -import math -from typing import NamedTuple +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Unit tests for vllm_omni/core/prefix_cache (manager + controller + pool). + +CPU-only: the controller runs in eager mode. Uses a fake group view, so no +vLLM runtime is required (runnable without a vllm install via +``pytest --confcutdir=tests/core``). +""" + +import ast +import logging +import sys +import threading +import time +from pathlib import Path +from types import SimpleNamespace import pytest import torch -from vllm_omni.core.prefix_cache import OmniTensorPrefixCache +try: # pragma: no cover - shim only matters on vllm-less dev machines + import vllm # noqa: F401 +except ModuleNotFoundError: + # Bypass vllm_omni/__init__ (which imports vllm): register namespace + # parents so the pure-torch prefix_cache subpackage imports directly. + _root = Path(__file__).resolve().parents[2] + for _pkg in ("vllm_omni", "vllm_omni.core", "vllm_omni.utils"): + if _pkg not in sys.modules: + _m = __import__("types").ModuleType(_pkg) + _m.__path__ = [str(_root / _pkg.replace(".", "/"))] + sys.modules[_pkg] = _m + _vllm = __import__("types").ModuleType("vllm") + _vllm_logger = __import__("types").ModuleType("vllm.logger") + _vllm_logger.init_logger = logging.getLogger + _vllm.logger = _vllm_logger + sys.modules["vllm"] = _vllm + sys.modules["vllm.logger"] = _vllm_logger + +from vllm_omni.core.prefix_cache.controller import StagingBufferHolder +from vllm_omni.core.prefix_cache.group_view import ( + FullAttentionGroupView, + check_prefix_cache_kv_groups, + stage_prefix_cache_config, +) +from vllm_omni.core.prefix_cache.interface import ( + HIDDEN_KEY, + ModelCachePolicy, + OmniPrefixCacheStagingTimeoutError, + OmniPrefixCacheUnmatchError, + PrefixCacheConfig, + WriteSchedule, +) +from vllm_omni.core.prefix_cache.manager import ( + OmniPrefixCacheManager, + _is_step_token_tensor, + _snapshot_leftover_mm_cpu, +) pytestmark = [pytest.mark.core_model, pytest.mark.cpu] - -DEFAULT_SEQ_LEN = 15 -NUM_BLOCKS = 10 +NUM_BLOCKS = 16 BLOCK_SIZE = 4 -HIDDEN_SIZE = 2 +HIDDEN = 8 DTYPE = torch.float32 -OTHER_DTYPE = torch.float16 -DEFAULT_SHAPE = torch.Size([NUM_BLOCKS, BLOCK_SIZE, HIDDEN_SIZE]) -class MockInputBatch: - 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 +class FakeView: + """Duck-typed group view backed by plain dicts.""" - # 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. - if block_table is None: + def __init__(self): + self.block_size = BLOCK_SIZE + self.req_blocks: dict[str, list[int]] = {} + self.order: list[str] = [] + self.computed: dict[str, int] = {} + self.step_slot_mapping: torch.Tensor | None = None - class _DummyBlockTable: - pass + def slots_for(self, req_id, token_start, token_end): + blocks = self.req_blocks[req_id] + slots = [] + for pos in range(token_start, token_end): + slots.append(blocks[pos // BLOCK_SIZE] * BLOCK_SIZE + pos % BLOCK_SIZE) + return torch.tensor(slots, dtype=torch.long) - self.block_table = _DummyBlockTable() - self.block_table.block_tables = [None] + def batch_req_ids(self) -> list[str]: + return list(self.order) + + def step_slots_cpu(self, req_ids, num_scheduled) -> torch.Tensor: + parts = [] + for r in req_ids: + n = int(num_scheduled.get(r, 0)) + if n <= 0: + continue + start = self.computed.get(r, 0) + parts.append(self.slots_for(r, start, start + n)) + return torch.cat(parts) if parts else torch.empty((0,), dtype=torch.long) + + +class FakeNewReq: + def __init__(self, req_id, num_computed_tokens=0, block_ids=None): + self.req_id = req_id + self.num_computed_tokens = num_computed_tokens + self.block_ids = block_ids + + +class FakeSchedOut: + def __init__(self, new_reqs=(), finished=(), num_scheduled=None): + self.scheduled_new_reqs = list(new_reqs) + self.finished_req_ids = set(finished) + self.num_scheduled_tokens = dict(num_scheduled or {}) + + +def make_manager(view=None, policy=None, **cfg_kwargs) -> tuple[OmniPrefixCacheManager, FakeView]: + view = view or FakeView() + config = PrefixCacheConfig(num_blocks=NUM_BLOCKS, block_size=BLOCK_SIZE, **cfg_kwargs) + mgr = OmniPrefixCacheManager(config, view, eager=True) + if policy is not None: + mgr.register_policy(policy) + return mgr, view + + +def run_step( + mgr, + view, + reqs: dict[str, tuple[list[int], int, int]], + new_hits=None, + finished=(), + mm=None, + num_tokens_padded=None, +) -> int: + """One step: reqs = req_id -> (blocks, sched_start_pos, sched_tokens).""" + view.order = list(reqs.keys()) + new_reqs = [] + num_sched = {} + slot_parts = [] + hidden_parts = [] + for req_id, (blocks, start_pos, sched) in reqs.items(): + view.req_blocks[req_id] = blocks + view.computed[req_id] = start_pos + num_sched[req_id] = sched + slots = view.slots_for(req_id, start_pos, start_pos + sched) + slot_parts.append(slots) + hidden_parts.append(slots.to(DTYPE).unsqueeze(1).expand(sched, HIDDEN).clone()) + hit = (new_hits or {}).get(req_id, 0) + new_reqs.append(FakeNewReq(req_id, num_computed_tokens=hit, block_ids=[list(blocks)])) + view.step_slot_mapping = torch.cat(slot_parts) + hidden = torch.cat(hidden_parts) + sched_out = FakeSchedOut(new_reqs=new_reqs, finished=finished, num_scheduled=num_sched) + mgr.new_step_starts(sched_out) + n = int(view.step_slot_mapping.numel()) + padded = n if num_tokens_padded is None else int(num_tokens_padded) + return mgr.save_outputs(hidden, mm or {}, num_tokens_unpadded=n, num_tokens_padded=padded) + + +def expected_rows(slots: torch.Tensor) -> torch.Tensor: + return slots.to(DTYPE).unsqueeze(1).expand(slots.numel(), HIDDEN) + + +def plan_fetch(mgr, slots, key, *, req_id): + with mgr._state_lock: + src = mgr._slot_ref(slots, key, req_id) + with torch.inference_mode(): + return mgr._fetch_source(src) + + +def _assert_leftover_shapes(inp, got, n: int) -> None: + """Every leftover tensor keeps its shape unless dim0 is this step's token count.""" + if isinstance(inp, torch.Tensor): + if inp.ndim >= 1 and n > 0 and inp.shape[0] == n: + assert got.shape == (n, *inp.shape[1:]) 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) + assert got.shape == inp.shape + return + if isinstance(inp, dict): + assert set(got) == set(inp) + for k in inp: + _assert_leftover_shapes(inp[k], got[k], n) + return + if isinstance(inp, (list, tuple)): + assert len(got) == len(inp) + for a, b in zip(inp, got): + _assert_leftover_shapes(a, b, n) -def get_omni_pcache_with_mm_tensors(feat_dims, seq_len) -> OmniTensorPrefixCache: - """Build an OmniTensorPrefixCache and init mm tensors.""" - cache = get_omni_pcache() - mm_outputs = get_multimodal_outputs(feat_dims, seq_len) - cache.maybe_init_missing_mm_cache_keys(mm_outputs, seq_len) - return cache +def _module_scope_imported_modules(path: Path) -> list[str]: + """Top-level import module names, skipping ``if TYPE_CHECKING`` blocks.""" + tree = ast.parse(path.read_text()) + names: list[str] = [] + + def _from_stmt(node: ast.AST) -> None: + if isinstance(node, ast.Import): + names.extend(alias.name for alias in node.names) + elif isinstance(node, ast.ImportFrom) and node.module: + names.append(node.module) + + for node in tree.body: + if isinstance(node, ast.If) and isinstance(node.test, ast.Name) and node.test.id == "TYPE_CHECKING": + continue + _from_stmt(node) + return names -def get_omni_pcache() -> OmniTensorPrefixCache: - """Build an OmniTensorPrefixCache, but don't init mm tensors.""" - cache = OmniTensorPrefixCache( +def test_prefix_cache_package_has_no_module_scope_vllm_import(): + """This package must import without vllm. See the package docstring.""" + pkg = Path(__file__).resolve().parents[2] / "vllm_omni" / "core" / "prefix_cache" + leaked: list[str] = [] + for path in sorted(pkg.glob("*.py")): + for mod in _module_scope_imported_modules(path): + if mod == "vllm" or mod.startswith("vllm."): + leaked.append(f"{path.name}: {mod}") + assert not leaked, leaked + + +def test_no_hit_passthrough(): + mgr, view = make_manager() + sid = run_step(mgr, view, {"a": ([0, 1], 0, 8)}) + outs = mgr.materialize(sid, ["a"]) + assert torch.equal(outs.hidden_states["a"], expected_rows(view.slots_for("a", 0, 8))) + + +def test_hit_merge_from_mirror(): + mgr, view = make_manager() + s1 = run_step(mgr, view, {"a": ([0, 1], 0, 8)}) + mgr.materialize(s1, ["a"]) + s2 = run_step(mgr, view, {"b": ([0, 1, 2], 8, 4)}, new_hits={"b": 8}, finished=["a"]) + outs = mgr.materialize(s2, ["b"]) + merged = outs.hidden_states["b"] + assert merged.shape == (12, HIDDEN) + assert torch.equal(merged[:8], expected_rows(view.slots_for("b", 0, 8))) + assert torch.equal(merged[8:], expected_rows(view.slots_for("b", 8, 12))) + + +@pytest.mark.parametrize("a_finished_before_b", [False, True]) +def test_partial_match_hits_shared_prefix_and_writes_divergent_block(a_finished_before_b): + """`a` occupies [0, 1, 2]; `b` shares the first two blocks and diverges on + the third ([0, 1, 3]). `b` must merge 8 hit rows from the shared blocks + plus its own 4 rows, land block 3 in the pool, and leave `a`'s block 2 + untouched — whether `a` already finished or is still running.""" + mgr, view = make_manager() + s1 = run_step(mgr, view, {"a": ([0, 1, 2], 0, 12)}) + mgr.materialize(s1, ["a"]) + a_tail = view.slots_for("a", 8, 12) + assert torch.equal(plan_fetch(mgr, a_tail, HIDDEN_KEY, req_id="a"), expected_rows(a_tail)) + + finished = ["a"] if a_finished_before_b else () + s2 = run_step(mgr, view, {"b": ([0, 1, 3], 8, 4)}, new_hits={"b": 8}, finished=finished) + outs = mgr.materialize(s2, ["b"]) + merged = outs.hidden_states["b"] + assert merged.shape == (12, HIDDEN) + assert torch.equal(merged[:8], expected_rows(view.slots_for("b", 0, 8))) + assert torch.equal(merged[8:], expected_rows(view.slots_for("b", 8, 12))) + # Divergent third block landed; the shared prefix and a's own tail did not move. + b_tail = view.slots_for("b", 8, 12) + assert torch.equal(plan_fetch(mgr, b_tail, HIDDEN_KEY, req_id="b"), expected_rows(b_tail)) + assert torch.equal(plan_fetch(mgr, a_tail, HIDDEN_KEY, req_id="a"), expected_rows(a_tail)) + shared = view.slots_for("a", 0, 8) + assert torch.equal(plan_fetch(mgr, shared, HIDDEN_KEY, req_id="a"), expected_rows(shared)) + + if not a_finished_before_b: + # `a` finishes later; releasing its task must not disturb b's rows. + s3 = run_step(mgr, view, {"b": ([0, 1, 3, 4], 12, 1)}, finished=["a"]) + mgr.materialize(s3, ["b"]) + assert torch.equal(plan_fetch(mgr, b_tail, HIDDEN_KEY, req_id="b"), expected_rows(b_tail)) + + +def test_live_req_reentering_new_reqs_is_not_a_hit(): + """async_chunk continuation: the same id re-enters scheduled_new_reqs with + num_computed_tokens equal to what it computed itself. No hit span is + opened and materialize returns only this step's rows (parity with the + pre-refactor cache; a delivered_upto span is Phase 2). Once the id has + finished, the same id is a new request again and its hit is merged.""" + mgr, view = make_manager() + s1 = run_step(mgr, view, {"a": ([0, 1], 0, 8)}) + mgr.materialize(s1, ["a"]) + + # Next chunk arrives: `a` is scheduled again as a "new" request at 8. + s2 = run_step(mgr, view, {"a": ([0, 1, 2], 8, 4)}, new_hits={"a": 8}) + assert not mgr._step_ctxs[s2].hits + outs = mgr.materialize(s2, ["a"]) + assert torch.equal(outs.hidden_states["a"], expected_rows(view.slots_for("a", 8, 12))) + + # `a` finishes; a later request reusing the id is new and does hit. + s3 = run_step(mgr, view, {"b": ([3], 0, 2)}, finished=["a"]) + mgr.materialize(s3, ["b"]) + s4 = run_step(mgr, view, {"a": ([0, 1, 4], 8, 4)}, new_hits={"a": 8}, finished=["b"]) + assert mgr._step_ctxs[s4].hits == {"a": (8, [0, 1])} + outs = mgr.materialize(s4, ["a"]) + assert outs.hidden_states["a"].shape == (12, HIDDEN) + assert torch.equal(outs.hidden_states["a"][:8], expected_rows(view.slots_for("a", 0, 8))) + + +def test_join_next_step_hit_waits_done_then_reads_pool(): + """CPU stand-in for the non-eager path: dispatch holds the task instead of + writing the pool. A same-step hit must join(done), drain, then read the pool.""" + mgr, view = make_manager() + real_run = mgr._controller._run_eager + real_join = mgr._controller.join + + def join_then_scatter(tids): + # Stand-in committer: the write lands only when someone joins it. + for tid in tids: + task = mgr._controller.get_task(tid) + if task is not None and not task.done.is_set(): + real_run(task) + real_join(tids) + + mgr._controller.dispatch = lambda tasks: None + mgr._controller.join = join_then_scatter + + view.req_blocks["a"] = [0, 1] + sid = run_step( + mgr, + view, + {"a": ([0, 1], 0, 8), "b": ([0, 1, 2], 8, 4)}, + new_hits={"b": 8}, + ) + outs = mgr.materialize(sid, ["a", "b"]) + assert torch.equal(outs.hidden_states["b"][:8], expected_rows(view.slots_for("b", 0, 8))) + slots = view.slots_for("a", 0, 8) + assert torch.equal(mgr._pool.rows(HIDDEN_KEY, slots), expected_rows(slots)) + + +def test_same_step_hit_reads_in_transit(): + mgr, view = make_manager() + view.req_blocks["a"] = [0, 1] + sid = run_step( + mgr, + view, + {"a": ([0, 1], 0, 8), "b": ([0, 1, 2], 8, 4)}, + new_hits={"b": 8}, + ) + outs = mgr.materialize(sid, ["a", "b"]) + assert torch.equal(outs.hidden_states["b"][:8], expected_rows(view.slots_for("b", 0, 8))) + + +def test_absent_hit_fails_fast(): + mgr, view = make_manager() + view.req_blocks["c"] = [5, 6] + sid = run_step(mgr, view, {"c": ([5, 6, 7], 8, 4)}, new_hits={"c": 8}) + d2h = mgr._step_ctxs[sid].d2h + with pytest.raises(OmniPrefixCacheUnmatchError): + mgr.materialize(sid, ["c"]) + if d2h is not None: + assert not mgr._controller._staging_pool._busy[d2h.staging_slot] + + +def test_hit_not_block_aligned_fails_at_register(): + mgr, view = make_manager() + s1 = run_step(mgr, view, {"a": ([0, 1], 0, 8)}) + mgr.materialize(s1, ["a"]) + with pytest.raises(OmniPrefixCacheUnmatchError, match="prefix hit not block aligned"): + run_step(mgr, view, {"b": ([0, 1, 2], 8, 4)}, new_hits={"b": 6}) + + +def test_mm_cached_key_merge(): + mgr, view = make_manager() + feat = 3 + mm1 = {"talker.h": torch.arange(8 * feat, dtype=DTYPE).reshape(8, feat)} + s1 = run_step(mgr, view, {"a": ([0, 1], 0, 8)}, mm=mm1) + mgr.materialize(s1, ["a"]) + mm2 = {"talker.h": torch.full((4, feat), 7.0)} + s2 = run_step(mgr, view, {"b": ([0, 1, 2], 8, 4)}, new_hits={"b": 8}, finished=["a"], mm=mm2) + outs = mgr.materialize(s2, ["b"]) + merged = outs.mm_outputs["talker.h"]["b"] + assert merged.shape == (12, feat) + assert torch.equal(merged[:8], mm1["talker.h"]) + assert torch.equal(merged[8:], mm2["talker.h"]) + + +def test_codes_audio_matches_ordinary_mm_key(): + """Without a deferred policy, codes.audio is an ordinary mm key (dim0 = tokens).""" + mgr, view = make_manager() + step1 = torch.arange(8 * 2, dtype=DTYPE).reshape(8, 2) + s1 = run_step(mgr, view, {"a": ([0, 1], 0, 8)}, mm={"codes.audio": step1}) + assert torch.equal(mgr.materialize(s1, ["a"]).mm_outputs["codes.audio"]["a"], step1) + step2 = torch.full((4, 2), 7.0) + s2 = run_step(mgr, view, {"b": ([0, 1, 2], 8, 4)}, new_hits={"b": 8}, finished=["a"], mm={"codes.audio": step2}) + merged = mgr.materialize(s2, ["b"]).mm_outputs["codes.audio"]["b"] + assert torch.equal(merged, torch.cat([step1, step2], dim=0)) + + +def test_unpadded_mm_registers_on_padded_step(): + """Qwen3-TTS codes.audio is scheduled-length; cudagraph makes padded > n. + + Immediate-only `shape[0] == padded` would skip ensure_key, drop the + rows into leftover, and silently serve hits without a cached prefix. + """ + mgr, view = make_manager() + audio = torch.arange(8 * 2, dtype=DTYPE).reshape(8, 2) + s1 = run_step(mgr, view, {"a": ([0, 1], 0, 8)}, mm={"codes.audio": audio}, num_tokens_padded=16) + assert mgr._pool.has_key("codes.audio") + assert "codes.audio" not in mgr._step_ctxs[s1].mm_cpu_snapshot + assert torch.equal(mgr.materialize(s1, ["a"]).mm_outputs["codes.audio"]["a"], audio) + s2 = run_step( + mgr, + view, + {"b": ([0, 1, 2], 8, 4)}, + new_hits={"b": 8}, + finished=["a"], + mm={"codes.audio": torch.full((4, 2), 3.0)}, + num_tokens_padded=16, + ) + merged = mgr.materialize(s2, ["b"]).mm_outputs["codes.audio"]["b"] + assert torch.equal(merged[:8], audio) + assert torch.equal(merged[8:], torch.full((4, 2), 3.0)) + + +def test_split_step_outputs_routes_immediate_deferred_leftover(): + """One split: immediate device→host, deferred write + leftover read, leftover only.""" + policy = ModelCachePolicy(needs_full_hidden_states=True, deferred_keys=frozenset({"codes.audio"})) + mgr, view = make_manager(policy=policy) + view.order = ["a"] + view.req_blocks["a"] = [0] + view.computed["a"] = 0 + n = 4 + mgr.new_step_starts(FakeSchedOut(new_reqs=[FakeNewReq("a")], num_scheduled={"a": n})) + hidden = torch.ones(n, HIDDEN) + mm = { + "codes.audio": torch.full((n, 2), 5.0), + "codes.ref": torch.zeros(15, 2), + "tok": torch.full((n, 2), 7.0), + } + step_outputs = mgr._split_step_outputs( + hidden, + mm, + n, + n, + slots_cpu=view.step_slots_cpu(["a"], {"a": n}), + req_order=["a"], + num_sched={"a": n}, + query_start={"a": 0}, + ) + assert HIDDEN_KEY in step_outputs.immediate and "tok" in step_outputs.immediate + assert "codes.audio" not in step_outputs.immediate and "codes.ref" not in step_outputs.immediate + assert "codes.audio" in step_outputs.leftover and "codes.ref" in step_outputs.leftover + assert "tok" not in step_outputs.leftover + assert len(step_outputs.deferred_chunks) == 1 + assert step_outputs.deferred_chunks[0][0] == "a" + assert "codes.audio" in step_outputs.deferred_chunks[0][1].tensors + + +def test_leftover_snapshot_preserves_non_token_major_shapes(): + """P1: leftover copy must not use `shape[0] >= n` as a slice predicate. + + Tensors whose first dim is this step's token count (`shape[0] == n`) + may be sliced; every other tensor — including `codes.ref` with + ref_len >> n, and list-held tensors — must keep the input shape. + """ + n = 4 + mm = { + "codes.ref": torch.arange(30, dtype=DTYPE).reshape(15, 2), + "tags": [torch.ones(3, 5), torch.zeros(7)], + "tokenish": torch.arange(n * 2, dtype=DTYPE).reshape(n, 2), + "scalar": torch.tensor(3.0), + } + out, event = _snapshot_leftover_mm_cpu(mm, set(), n) + assert event is None + _assert_leftover_shapes(mm, out, n) + assert torch.equal(out["codes.ref"], mm["codes.ref"]) + assert out["codes.ref"].shape == (15, 2) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA leftover copy") +def test_leftover_snapshot_cuda_is_pinned_and_event_ordered(): + """CUDA leftover lands pinned via a non-blocking copy; values are exact once the event is waited.""" + n = 4 + mm = { + "codes.audio": torch.arange(n * 2, dtype=DTYPE, device="cuda").reshape(n, 2), + "codes.ref": torch.arange(30, dtype=DTYPE, device="cuda").reshape(15, 2), + "tags": [torch.ones(3, 5, device="cuda")], + } + out, event = _snapshot_leftover_mm_cpu(mm, set(), n) + assert event is not None + event.synchronize() + _assert_leftover_shapes(mm, out, n) + for t in (out["codes.audio"], out["codes.ref"], out["tags"][0]): + assert t.device.type == "cpu" and t.is_pinned() + assert torch.equal(out["codes.audio"], mm["codes.audio"].cpu()) + assert torch.equal(out["codes.ref"], mm["codes.ref"].cpu()) + + +def test_codes_ref_matches_old_build_mm_cpu_path(): + """Leftover `codes.ref` equals the old runner path: build_mm_cpu + to_payload_element.""" + pytest.importorskip("vllm") + from vllm_omni.utils.mm_outputs import build_mm_cpu, to_payload_element + + refs = [torch.arange(8, dtype=torch.long), torch.arange(12, dtype=torch.long)] + mgr, view = make_manager() + sid = run_step(mgr, view, {"a": ([0], 0, 2), "b": ([1], 0, 2)}, mm={"codes.ref": refs}) + ctx = mgr._step_ctxs[sid] + assert "codes.ref" in ctx.mm_cpu_snapshot + assert [t.shape for t in ctx.mm_cpu_snapshot["codes.ref"]] == [t.shape for t in refs] + outs = mgr.materialize(sid, ["a", "b"]) + old = build_mm_cpu({"codes.ref": refs}) + total = 4 + for req_id, idx, start, end in (("a", 0, 0, 2), ("b", 1, 2, 4)): + want = to_payload_element(old["codes.ref"], idx, start=start, end=end, pass_lists_through=True, seq_len=total) + got = outs.mm_outputs["codes.ref"][req_id] + assert len(got) == len(want) + for g, w in zip(got, want): + assert torch.equal(g, w) + + +def test_is_step_token_tensor(): + n, padded = 4, 8 + assert _is_step_token_tensor(torch.zeros(n, 2), n, padded) + assert _is_step_token_tensor(torch.zeros(padded, 2), n, padded) + assert not _is_step_token_tensor([torch.zeros(1, 8)], n, padded) + assert not _is_step_token_tensor(torch.zeros(15, 2), n, padded) + assert not _is_step_token_tensor(torch.tensor([1, 2, 3]), n, padded) + + +def test_policy_from_model_shim(): + class M: + requires_full_prefix_cached_hidden_states = False + deferred_prefix_cache_mm_keys = {"codes.audio"} + + p = ModelCachePolicy.from_model(M()) + assert p.needs_full_hidden_states is False + assert p.hidden_key is None + assert p.deferred_keys == frozenset({"codes.audio"}) + assert p.get_hit_keys([HIDDEN_KEY, "codes.audio"]) == ["codes.audio"] + assert p.skip_immediate_mm("codes.audio") + d = ModelCachePolicy.from_model(object()) + assert d.needs_full_hidden_states is True and not d.deferred_keys + assert d.hidden_key == HIDDEN_KEY + assert d.get_hit_keys(["talker.h", HIDDEN_KEY]) == [HIDDEN_KEY, "talker.h"] + + +def test_deferred_key_accumulates_and_flushes_on_finish(): + policy = ModelCachePolicy(needs_full_hidden_states=False, deferred_keys=frozenset({"codes.audio"})) + mgr, view = make_manager(policy=policy) + feat = 2 + for pos in range(2): + sid = run_step( + mgr, + view, + {"a": ([3], pos, 1)}, + mm={"codes.audio": torch.full((1, feat), float(pos + 1))}, + ) + mgr.materialize(sid, ["a"]) + slots = view.slots_for("a", 0, 2) + rows = plan_fetch(mgr, slots, "codes.audio", req_id="a") + assert torch.equal(rows[:, 0], torch.tensor([1.0, 2.0])) + sid = run_step(mgr, view, {"z": ([9], 0, 1)}, finished=["a"]) + mgr.materialize(sid, ["z"]) + mirror = mgr._pool.rows("codes.audio", slots) + assert torch.equal(mirror[:, 0], torch.tensor([1.0, 2.0])) + + +def test_deferred_gpu_bytes_held_until_last_view_drops(): + """Shared deferred clone is one allocation; finishing the first + co-scheduled request must not drop ``_staged_bytes`` while the other + still pins it.""" + policy = ModelCachePolicy(needs_full_hidden_states=False, deferred_keys=frozenset({"k"})) + mgr, view = make_manager(policy=policy) + feat = 2 + n = 4 + s1 = run_step( + mgr, + view, + {"a": ([0], 0, 2), "b": ([1], 0, 2)}, + mm={"k": torch.zeros(n, feat)}, + ) + mgr.materialize(s1, ["a", "b"]) + clone_bytes = n * feat * 4 # float32 + assert mgr._controller._staged_bytes == clone_bytes + s2 = run_step(mgr, view, {"z": ([9], 0, 1)}, finished=["a"]) + mgr.materialize(s2, ["z"]) + assert mgr._controller._staged_bytes == clone_bytes + s3 = run_step(mgr, view, {"z2": ([9], 1, 1)}, finished=["b"]) + mgr.materialize(s3, ["z2"]) + assert mgr._controller._staged_bytes == 0 + + +def test_cap_forces_flush_of_deferred(): + """One long request's deferred clone must not pin the whole budget: + the flush closes its open task and the next step opens a new one.""" + policy = ModelCachePolicy(needs_full_hidden_states=False, deferred_keys=frozenset({"k"})) + mgr, view = make_manager(policy=policy, gpu_staging_bytes=48) # three (1, 4) fp32 chunks + for pos in range(4): + sid = run_step(mgr, view, {"a": ([2, 3], pos, 1)}, mm={"k": torch.full((1, 4), float(pos + 1))}) + mgr.materialize(sid, ["a"]) + # Fourth chunk flushed the first three; only it is still charged. + assert mgr._controller._staged_bytes == 16 + rows = mgr._pool.rows("k", view.slots_for("a", 0, 4))[:, 0].tolist() + assert rows == [1.0, 2.0, 3.0, 0.0] + assert mgr._request_tasks.deferred["a"].write_n == 2 + assert len(mgr._request_tasks.tasks["a"]) == 1 # the flushed write was drained + + +def test_deferred_tenant_succession_no_stale_wins(): + policy = ModelCachePolicy(needs_full_hidden_states=False, deferred_keys=frozenset({"k"})) + mgr, view = make_manager(policy=policy) + s1 = run_step(mgr, view, {"a": ([2], 0, 2)}, mm={"k": torch.full((2, 2), 1.0)}) + mgr.materialize(s1, ["a"]) + s2 = run_step(mgr, view, {"b": ([2], 0, 2)}, mm={"k": torch.full((2, 2), 2.0)}) + mgr.materialize(s2, ["b"]) + slots = view.slots_for("b", 0, 2) + s3 = run_step(mgr, view, {"z1": ([9], 0, 1)}, finished=["b"]) + mgr.materialize(s3, ["z1"]) + s4 = run_step(mgr, view, {"z2": ([9], 1, 1)}, finished=["a"]) + mgr.materialize(s4, ["z2"]) + assert torch.equal(mgr._pool.rows("k", slots), torch.full((2, 2), 2.0)) + + +def test_reacquired_slot_deferred_rewrite_wins_over_interim_tenant(): + """Preempt + resume onto the same block: the rewrite must open a new + write, or the old task's reassigned tombstone drops the fresh rows and + the interim tenant's stale value is served as a's.""" + policy = ModelCachePolicy(needs_full_hidden_states=False, deferred_keys=frozenset({"k"})) + mgr, view = make_manager(policy=policy) + s1 = run_step(mgr, view, {"a": ([2], 0, 1)}, mm={"k": torch.full((1, 2), 11.0)}) + mgr.materialize(s1, ["a"]) + slots = view.slots_for("a", 0, 1) + s2 = run_step(mgr, view, {"b": ([2], 0, 1)}, mm={"k": torch.full((1, 2), 22.0)}) + mgr.materialize(s2, ["b"]) + s3 = run_step(mgr, view, {"z": ([9], 0, 1)}, finished=["b"]) # 22 committed + mgr.materialize(s3, ["z"]) + s4 = run_step(mgr, view, {"a": ([2], 0, 1)}, mm={"k": torch.full((1, 2), 33.0)}) # a re-claims + mgr.materialize(s4, ["a"]) + s5 = run_step(mgr, view, {"z2": ([9], 1, 1)}, finished=["a"]) # flush a's write + mgr.materialize(s5, ["z2"]) + assert torch.equal(mgr._pool.rows("k", slots), torch.full((1, 2), 33.0)) + + +def test_deferred_key_hit_reads_staged_rows_not_mirror(): + policy = ModelCachePolicy(needs_full_hidden_states=True, deferred_keys=frozenset({"k"})) + mgr, view = make_manager(policy=policy) + s1 = run_step(mgr, view, {"a": ([0], 0, 4)}, mm={"k": torch.full((4, 2), 1.0)}) + mgr.materialize(s1, ["a"]) + s2 = run_step(mgr, view, {"b": ([0, 1], 4, 2)}, new_hits={"b": 4}, mm={"k": torch.full((2, 2), 9.0)}) + rows = mgr.materialize(s2, ["b"]).mm_outputs["k"]["b"] + assert torch.equal(rows[:4], torch.full((4, 2), 1.0)), rows + assert torch.equal(rows[4:], torch.full((2, 2), 9.0)) + + +def test_append_to_closed_deferred_entry_opens_new_one(): + """A deferred task closed by the budget flush (or an early escalate) + cannot take more chunks; the next save opens a new write and a hit + reads both.""" + policy = ModelCachePolicy(needs_full_hidden_states=False, deferred_keys=frozenset({"k"})) + mgr, view = make_manager(policy=policy) + s1 = run_step(mgr, view, {"a": ([2], 0, 1)}, mm={"k": torch.full((1, 2), 1.0)}) + mgr.materialize(s1, ["a"]) + first = mgr._request_tasks.deferred["a"] + mgr._controller.escalate([first.tid]) + s2 = run_step(mgr, view, {"a": ([2], 1, 1)}, mm={"k": torch.full((1, 2), 2.0)}) + mgr.materialize(s2, ["a"]) + assert mgr._request_tasks.deferred["a"].tid != first.tid + slots = view.slots_for("a", 0, 2) + rows = plan_fetch(mgr, slots, "k", req_id="a") + assert torch.equal(rows[:, 0], torch.tensor([1.0, 2.0])) + + +def test_deferred_unpadded_registers_on_padded_step(): + """Qwen3-TTS: deferred codes.audio is unpadded; cudagraph pads hidden.""" + policy = ModelCachePolicy(needs_full_hidden_states=False, deferred_keys=frozenset({"codes.audio"})) + mgr, view = make_manager(policy=policy) + audio = torch.arange(8 * 2, dtype=DTYPE).reshape(8, 2) + s1 = run_step(mgr, view, {"a": ([0, 1], 0, 8)}, mm={"codes.audio": audio}, num_tokens_padded=16) + assert mgr._pool.has_key("codes.audio") + assert "codes.audio" in mgr._step_ctxs[s1].mm_cpu_snapshot + rows = plan_fetch(mgr, view.slots_for("a", 0, 8), "codes.audio", req_id="a") + assert torch.equal(rows, audio) + outs = mgr.materialize(s1, ["a"]) + assert torch.equal(outs.mm_outputs["codes.audio"]["a"], audio) + + +def test_check_kv_groups_rejects_empty_or_multi(): + with pytest.raises(OmniPrefixCacheUnmatchError, match="single full-attention"): + check_prefix_cache_kv_groups([]) + with pytest.raises(OmniPrefixCacheUnmatchError, match="single full-attention"): + check_prefix_cache_kv_groups([object(), object()]) + + +def _stage_cfg(*, enable=True, pooling=False, kv_transfer=None, groups=(object(),), spec=None, match_unit=None): + return stage_prefix_cache_config( + kv_cache_config=SimpleNamespace(num_blocks=NUM_BLOCKS, kv_cache_groups=list(groups)), + cache_config=SimpleNamespace(enable_prefix_caching=enable, block_size=BLOCK_SIZE, prefix_match_unit=match_unit), + kv_transfer_config=kv_transfer, + scheduler_config=SimpleNamespace(max_num_batched_tokens=64, max_model_len=128), + model_config=None, + is_pooling_model=pooling, + speculative_config=spec, + ) + + +def test_stage_prefix_cache_config_gate(): + """The runner-side gate the GPU and NPU runners share: off / pooling + stages get no cache; kv_consumer, spec decode, sub-block matching and + hybrid groups refuse loudly.""" + assert _stage_cfg(enable=False) is None + assert _stage_cfg(pooling=True) is None + # kv_consumer / kv_both refuse before the group check, on both platforms. + with pytest.raises(OmniPrefixCacheUnmatchError, match="KV connector"): + _stage_cfg(kv_transfer=SimpleNamespace(is_kv_consumer=True), groups=()) + assert _stage_cfg(kv_transfer=SimpleNamespace(is_kv_consumer=True), enable=False) is None + with pytest.raises(OmniPrefixCacheUnmatchError, match="speculative decoding"): + _stage_cfg(spec=SimpleNamespace(method="ngram"), groups=()) + with pytest.raises(OmniPrefixCacheUnmatchError, match="block-aligned"): + _stage_cfg(match_unit=BLOCK_SIZE // 2, groups=()) + # A unit equal to the block is a whole-block match; only the group check is left. + with pytest.raises(OmniPrefixCacheUnmatchError, match="single full-attention"): + _stage_cfg(match_unit=BLOCK_SIZE, groups=()) + with pytest.raises(OmniPrefixCacheUnmatchError, match="single full-attention"): + _stage_cfg(groups=(object(), object())) + + +def test_npu_runner_uses_shared_prefix_cache_gate(): + """npu_model_runner cannot be imported without vllm_ascend; pin at the + source level that it stages through the shared gate instead of a + hand-rolled subset of the GPU checks.""" + path = Path(__file__).resolve().parents[2] / "vllm_omni/platforms/npu/worker/npu_model_runner.py" + tree = ast.parse(path.read_text()) + calls = {n.func.id for n in ast.walk(tree) if isinstance(n, ast.Call) and isinstance(n.func, ast.Name)} + assert "stage_prefix_cache_config" in calls + assert "check_prefix_cache_kv_groups" not in calls + + +def _table_slots(table, req_idx, token_start, token_end, block_size=BLOCK_SIZE): + if token_end <= token_start: + return torch.empty((0,), dtype=torch.long) + token_positions = torch.arange(token_start, token_end, dtype=torch.long) + block_offsets = token_positions // block_size + max_blocks = int(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 = table[req_idx, block_offsets].to(torch.long) + return block_ids * block_size + (token_positions % block_size) + + +def test_group_view_step_slots_cpu(): + class TensorWrap: + def __init__(self, t): + self.cpu = t + + class Group: + def __init__(self, t): + self.block_table = TensorWrap(t) + + class BT: + def __init__(self, t): + self._g = Group(t) + self.block_tables = [self._g.block_table] + + def __getitem__(self, idx): + assert idx == 0 + return self._g + + class IB: + def __init__(self, table): + self.req_ids = ["r1", "r2"] + self.req_id_to_index = {"r1": 0, "r2": 1} + self.num_computed_tokens_cpu = torch.tensor([8, 0]) + self.block_table = BT(table) + + table = torch.tensor([[2, 5, 7], [1, 0, 0]]) + view = FullAttentionGroupView(IB(table), block_size=BLOCK_SIZE) + assert view.step_slots_cpu(["r1"], {"r1": 2}).tolist() == [7 * 4 + 0, 7 * 4 + 1] + assert view.step_slots_cpu(["r1"], {"r1": 8}).tolist() == [7 * 4 + i for i in range(4)] + assert view.step_slots_cpu(["r1"], {"r1": 0}).numel() == 0 + + +def test_join_next_step_previous_save(): + mgr, view = make_manager() + s1 = run_step(mgr, view, {"a": ([0], 0, 2)}) + assert len(mgr._join_next_step_tids) == 1 + task_id = mgr._join_next_step_tids[0] + mgr.materialize(s1, ["a"]) + s2 = run_step(mgr, view, {"a": ([0], 2, 1)}) + assert mgr._controller.get_task(task_id) is None + assert int(mgr._slot_status.state[HIDDEN_KEY][view.slots_for("a", 0, 2)].min()) == 2 + mgr.materialize(s2, ["a"]) + + +def test_tenant_succession_hit_reads_newest(): + mgr, view = make_manager() + s1 = run_step(mgr, view, {"a": ([0], 0, 4)}) + mgr.materialize(s1, ["a"]) + s2 = run_step(mgr, view, {"b": ([0], 0, 4)}, finished=["a"]) + mgr.materialize(s2, ["b"]) + b_hidden = expected_rows(view.slots_for("b", 0, 4)) + s3 = run_step(mgr, view, {"c": ([0, 1], 4, 2)}, new_hits={"c": 4}, finished=["b"]) + outs = mgr.materialize(s3, ["c"]) + assert torch.equal(outs.hidden_states["c"][:4], b_hidden) + + +def test_tenant_succession_mm_key(): + mgr, view = make_manager() + s1 = run_step(mgr, view, {"a": ([2], 0, 4)}, mm={"k": torch.full((4, 2), 1.0)}) + mgr.materialize(s1, ["a"]) + s2 = run_step(mgr, view, {"b": ([2], 0, 4)}, finished=["a"], mm={"k": torch.full((4, 2), 2.0)}) + mgr.materialize(s2, ["b"]) + s3 = run_step( + mgr, view, {"c": ([2, 3], 4, 2)}, new_hits={"c": 4}, finished=["b"], mm={"k": torch.full((2, 2), 3.0)} + ) + rows = mgr.materialize(s3, ["c"]).mm_outputs["k"]["c"] + assert torch.equal(rows[:4], torch.full((4, 2), 2.0)) + assert torch.equal(rows[4:], torch.full((2, 2), 3.0)) + + +def test_scatter_rows_coalesces_per_key_last_chunk_wins_and_skips_reassigned(): + from vllm_omni.core.prefix_cache.controller import WriteSchedule, WriteTask, _BudgetTicket, _WriteChunk + + t1, t2 = _BudgetTicket(nbytes=8), _BudgetTicket(nbytes=8) + + def chunk(slots, base, ticket): + s = torch.tensor(slots, dtype=torch.int64) + host = {"k": torch.full((len(slots), 2), float(base)) + s[:, None]} + return _WriteChunk(slots_cpu=s, tensors={}, host=host, budget=ticket) + + task = WriteTask( + tid=1, + req_id="a", + write_n=1, + schedule=WriteSchedule.JOIN_ON_FINISH, + chunks=[chunk([0, 1], 10, t1), chunk([2, 3], 20, t1), chunk([1, 4], 30, t2)], + ) + assert task.budget_tickets() == [t1, t2] + task.add_reassigned("k", torch.tensor([3], dtype=torch.int64)) + rows = task.scatter_rows() + assert len(rows) == 1 + key, slots, host = rows[0] + assert key == "k" and slots.tolist() == [0, 2, 1, 4] + # slot 1: the third chunk's row (base 30) replaces the first chunk's; slot 3 is reassigned. + assert host[:, 0].tolist() == [10.0, 22.0, 31.0, 34.0] + + +def test_slot_reuse_pushes_skip_to_old_task(): + """Reassignment = task swap. Keep the old write in progress so the + new write records those rows as no longer owned by it.""" + mgr, view = make_manager() + + def hold(tasks): + for task in tasks: + task.host_ready.set() + task.done.set() # satisfy the next save's done barrier; state stays QUEUED + + mgr._controller.dispatch = hold + policy = ModelCachePolicy(needs_full_hidden_states=False) + mgr.register_policy(policy) + s1 = run_step(mgr, view, {"a": ([2], 0, 1)}, mm={"k": torch.ones(1, 2)}) + mgr.materialize(s1, ["a"]) + old_tid = next(iter(mgr._request_tasks.tasks["a"])) + old_task = mgr._controller.get_task(old_tid) + s2 = run_step(mgr, view, {"b": ([2], 0, 2)}, mm={"k": torch.full((2, 2), 2.0)}) + assert old_task is not None and "k" in old_task.reassigned + reused = view.slots_for("b", 0, 1) + assert bool(torch.isin(reused, old_task.reassigned["k"]).all()) + mgr.materialize(s2, ["b"]) + + +def test_step_slots_cpu_matches_block_table_math(): + class TensorWrap: + def __init__(self, t): + self.cpu = t + + class Group: + def __init__(self, t): + self.block_table = TensorWrap(t) + + class BT: + def __init__(self, t): + self._g = Group(t) + self.block_tables = [self._g.block_table] + + def __getitem__(self, idx): + return self._g + + class IB: + def __init__(self, table, computed): + self.req_ids = ["r1", "r2"] + self.req_id_to_index = {"r1": 0, "r2": 1} + self.num_computed_tokens_cpu = torch.tensor(computed) + self.block_table = BT(table) + + table = torch.tensor([[2, 5, 7], [1, 3, 4]]) + view = FullAttentionGroupView(IB(table, [4, 0]), block_size=BLOCK_SIZE) + num_sched = {"r1": 3, "r2": 5} + got = view.step_slots_cpu(["r1", "r2"], num_sched) + want = torch.cat([_table_slots(table, 0, 4, 7), _table_slots(table, 1, 0, 5)]) + assert torch.equal(got, want), (got, want) + assert torch.equal(view.step_slots_cpu(["r1", "r2"], {"r1": 3, "r2": 0}), _table_slots(table, 0, 4, 7)) + + +def _ib_with_layout(**layout): + """InputBatch fake whose group-0 block table carries vLLM's layout attrs.""" + attrs = dict(kv_cache_block_size=BLOCK_SIZE, blocks_per_kv_block=1, use_hybrid_blocks=False, dcp_world_size=1) + attrs.update(layout) + group = SimpleNamespace(block_table=SimpleNamespace(cpu=torch.zeros((1, 2), dtype=torch.int32)), **attrs) + + class BT: + block_tables = [group.block_table] + + def __getitem__(self, idx): + return group + + return SimpleNamespace(req_ids=[], req_id_to_index={}, num_computed_tokens_cpu=torch.zeros(1), block_table=BT()) + + +def test_group_view_refuses_hybrid_kernel_blocks_and_dcp(): + """step_slots_cpu assumes one allocator block per table column and no + token striping; both are silent wrong-slot writes if not refused.""" + FullAttentionGroupView(_ib_with_layout(), block_size=BLOCK_SIZE) + with pytest.raises(OmniPrefixCacheUnmatchError, match="kernel_block_size"): + FullAttentionGroupView(_ib_with_layout(use_hybrid_blocks=True, blocks_per_kv_block=2), block_size=BLOCK_SIZE) + with pytest.raises(OmniPrefixCacheUnmatchError, match="decode context parallel"): + FullAttentionGroupView(_ib_with_layout(dcp_world_size=2), block_size=BLOCK_SIZE) + with pytest.raises(OmniPrefixCacheUnmatchError, match="does not match"): + FullAttentionGroupView(_ib_with_layout(kv_cache_block_size=BLOCK_SIZE * 2), block_size=BLOCK_SIZE) + + +def test_step_context_consumed_by_id_not_order(): + mgr, view = make_manager() + s1 = run_step(mgr, view, {"a": ([0], 0, 4)}) + s2 = run_step(mgr, view, {"b": ([1], 0, 4)}) + mgr.discard_step(s2) + outs = mgr.materialize(s1, ["a"]) + assert torch.equal(outs.hidden_states["a"], expected_rows(view.slots_for("a", 0, 4))) + assert len(mgr._step_ctxs) == 0 + + +def test_step_context_exactly_once(): + mgr, view = make_manager() + sid = run_step(mgr, view, {"a": ([0], 0, 4)}) + mgr.materialize(sid, ["a"]) + with pytest.raises(OmniPrefixCacheUnmatchError): + mgr.materialize(sid, ["a"]) + with pytest.raises(OmniPrefixCacheUnmatchError): + mgr.discard_step(sid) + + +def test_unconsumed_contexts_timeout(): + mgr, view = make_manager(staging_claim_timeout_s=0.0) + with pytest.raises(OmniPrefixCacheStagingTimeoutError, match="unconsumed step contexts"): + for pos in range(8): + run_step(mgr, view, {"a": ([0, 1, 2, 3], pos, 1)}) + + +def test_save_waits_until_a_slot_is_consumed(): + mgr, view = make_manager(staging_depth=2, staging_claim_timeout_s=2.0) + s0 = run_step(mgr, view, {"a": ([0], 0, 1)}) + s1 = run_step(mgr, view, {"a": ([0], 1, 1)}) + err: list[BaseException] = [] + sid_holder: list[int] = [] + + def _blocked_save() -> None: + try: + sid_holder.append(run_step(mgr, view, {"a": ([0], 2, 1)})) + except BaseException as e: + err.append(e) + + t = threading.Thread(target=_blocked_save, daemon=True) + t.start() + time.sleep(0.1) + assert t.is_alive() + mgr.discard_step(s0) + t.join(timeout=1.0) + assert not t.is_alive() + assert not err + mgr.discard_step(s1) + mgr.discard_step(sid_holder[0]) + + +def test_leftover_only_save_claims_staging_slot(): + policy = ModelCachePolicy(needs_full_hidden_states=False) + mgr, view = make_manager(policy=policy, staging_depth=2, staging_claim_timeout_s=0.0) + leftover = {"codes.ref": torch.zeros(3, 2)} + s0 = run_step(mgr, view, {"a": ([0], 0, 1)}, mm=leftover) + s1 = run_step(mgr, view, {"a": ([0], 1, 1)}, mm=leftover) + assert mgr._step_ctxs[s0].d2h is not None and mgr._step_ctxs[s0].d2h.views == {} + assert mgr._step_ctxs[s1].d2h is not None and mgr._step_ctxs[s1].d2h.views == {} + with pytest.raises(OmniPrefixCacheStagingTimeoutError, match="unconsumed step contexts"): + run_step(mgr, view, {"a": ([0], 2, 1)}, mm=leftover) + + +def test_save_slot_mismatch_fails_fast(): + mgr, view = make_manager() + view.order = ["a"] + view.req_blocks["a"] = [0] + view.computed["a"] = 0 + mgr.new_step_starts(FakeSchedOut(new_reqs=[FakeNewReq("a")], num_scheduled={"a": 2})) + hidden = torch.zeros(4, HIDDEN, dtype=DTYPE) + with pytest.raises(OmniPrefixCacheUnmatchError): + mgr.save_outputs(hidden, {}, num_tokens_unpadded=4, num_tokens_padded=4) + + +def test_materialize_rejects_out_of_snapshot_ids(): + mgr, view = make_manager() + s1 = run_step(mgr, view, {"a": ([0, 1], 0, 8)}) + mgr.materialize(s1, ["a"]) + s2 = run_step(mgr, view, {"b": ([0, 1, 2], 8, 4)}, new_hits={"b": 8}, finished=["a"]) + with pytest.raises(OmniPrefixCacheUnmatchError, match="outside the save snapshot"): + mgr.materialize(s2, ["b", "late_joiner"]) + + +def test_materialize_unknown_step_id_fails_fast(): + mgr, view = make_manager() + sid = run_step(mgr, view, {"a": ([0], 0, 4)}) + with pytest.raises(OmniPrefixCacheUnmatchError): + mgr.materialize(sid + 999, ["ghost"]) + mgr.discard_step(sid) + + +def test_mm_hit_span_never_registered_serves_mirror_baseline(): + """A hit span slot on which a sparse mm key was never registered is + legitimate absence — served from the mirror baseline, not a crash.""" + policy = ModelCachePolicy(needs_full_hidden_states=True, deferred_keys=frozenset({"k"})) + mgr, view = make_manager(policy=policy) + s1 = run_step(mgr, view, {"a": ([0], 0, 4)}, mm={"k": torch.full((4, 2), 5.0)}) + mgr.materialize(s1, ["a"]) + s2 = run_step(mgr, view, {"c": ([1], 0, 4)}) + mgr.materialize(s2, ["c"]) + s3 = run_step(mgr, view, {"b": ([0, 1, 2], 8, 2)}, new_hits={"b": 8}, mm={"k": torch.full((2, 2), 9.0)}) + rows = mgr.materialize(s3, ["b"]).mm_outputs["k"]["b"] + assert torch.equal(rows[:4], torch.full((4, 2), 5.0)) + assert torch.equal(rows[4:8], torch.zeros(4, 2)) + assert torch.equal(rows[8:], torch.full((2, 2), 9.0)) + + +def test_mm_in_transit_unresolvable_fails_fast(): + """Rows registered in-transit whose entry cannot serve them must raise.""" + policy = ModelCachePolicy(needs_full_hidden_states=True, deferred_keys=frozenset({"k"})) + mgr, view = make_manager(policy=policy) + s1 = run_step(mgr, view, {"a": ([0], 0, 4)}, mm={"k": torch.full((4, 2), 1.0)}) + mgr.materialize(s1, ["a"]) + tid = mgr._request_tasks.deferred["a"].tid + mgr._controller._tasks.pop(tid) + s2 = run_step(mgr, view, {"b": ([0, 1], 4, 2)}, new_hits={"b": 4}, mm={"k": torch.full((2, 2), 9.0)}) + with pytest.raises(OmniPrefixCacheUnmatchError): + mgr.materialize(s2, ["b"]) + + +def test_lock_never_covers_fetch_or_join(): + """State lock must not be held across join / join_host_ready / fetch_host.""" + policy = ModelCachePolicy(needs_full_hidden_states=True, deferred_keys=frozenset({"k"})) + mgr, view = make_manager(policy=policy) + calls = [] + + def probe(kind): + on_facade = not threading.current_thread().name.startswith("omni-prefix-cache-prefetch") + calls.append((kind, on_facade and mgr._state_lock.locked())) + + real_fetch = mgr._controller.fetch_host + real_join_hr = mgr._controller.join_host_ready + real_join = mgr._controller.join + mgr._controller.fetch_host = lambda *a, **kw: (probe("fetch"), real_fetch(*a, **kw))[1] + mgr._controller.join_host_ready = lambda ids: (probe("join"), real_join_hr(ids))[1] + mgr._controller.join = lambda ids: (probe("join"), real_join(ids))[1] + + s1 = run_step(mgr, view, {"a": ([0], 0, 4)}, mm={"k": torch.full((4, 2), 1.0)}) + mgr.materialize(s1, ["a"]) + s2 = run_step(mgr, view, {"b": ([0, 1], 4, 2)}, new_hits={"b": 4}, mm={"k": torch.full((2, 2), 9.0)}) + mgr.materialize(s2, ["b"]) + assert any(kind == "fetch" for kind, _ in calls) + assert any(kind == "join" for kind, _ in calls) + assert all(not locked for _, locked in calls), calls + + +def test_failed_write_fails_fast_at_next_facade_entry(): + mgr, view = make_manager() + sid = run_step(mgr, view, {"a": ([0], 0, 4)}) + mgr.materialize(sid, ["a"]) + from vllm_omni.core.prefix_cache.controller import WriteTask, _WriteChunk + + task = WriteTask( + tid=999, + req_id="x", + write_n=1, + schedule=WriteSchedule.JOIN_NEXT_STEP, + chunks=[_WriteChunk(slots_cpu=torch.tensor([0]), tensors={})], + ) + mgr._controller._tasks[999] = task + mgr._controller._fail_task(999) + with pytest.raises(OmniPrefixCacheUnmatchError, match="write failed"): + run_step(mgr, view, {"a": ([0, 1], 4, 1)}) + # Sticky: the record is one-shot and the first raise may land in a + # swallowed path (prefetch Future dropped by discard_step); a later + # entry must keep raising, never serve zeros for the lost rows. + with pytest.raises(OmniPrefixCacheUnmatchError, match="write failed"): + run_step(mgr, view, {"c": ([0, 1], 0, 4)}) + + +def test_join_raises_on_failed_task(): + """A joiner woken by FAILED must not treat the wake as a completed + write; the one-shot failure record may already be drained elsewhere.""" + mgr, view = make_manager() + _hold_writes(mgr) + run_step(mgr, view, {"a": ([0], 0, 4)}) + tid = next(iter(mgr._request_tasks.tasks["a"])) + mgr._controller._fail_task(tid) + with pytest.raises(OmniPrefixCacheUnmatchError, match="failed before done"): + mgr._controller.join([tid]) + + +@pytest.mark.parametrize("fail_at", ["copy", "scatter"]) +def test_eager_dispatch_failure_releases_step_and_all_task_owners(monkeypatch, fail_at): + from vllm_omni.core.prefix_cache.controller import TaskState, WriteTask + + mgr, view = make_manager(staging_claim_timeout_s=0) + ctrl = mgr._controller + preserved_sid = run_step(mgr, view, {"previous": ([3], 0, 4)}) + if fail_at == "copy": + original = WriteTask.mark_host_ready + + def fail_copy(task): + if task.req_id == "b": + raise RuntimeError("injected copy failure") + return original(task) + + monkeypatch.setattr(WriteTask, "mark_host_ready", fail_copy) + else: + original = ctrl._scatter + + def fail_scatter(task): + if task.req_id == "b": + raise RuntimeError("injected scatter failure") + return original(task) + + monkeypatch.setattr(ctrl, "_scatter", fail_scatter) + + try: + with pytest.raises(RuntimeError, match=f"injected {fail_at} failure"): + run_step(mgr, view, {"a": ([0], 0, 4), "b": ([1], 0, 4), "c": ([2], 0, 4)}) + assert set(mgr._step_ctxs) == {preserved_sid} + assert ctrl._staged_bytes == 0 + tasks = [task for task in ctrl._tasks.values() if task.req_id in {"a", "b", "c"}] + assert len(tasks) == 3 + for task in tasks: + expected = TaskState.WRITTEN if task.req_id == "a" else TaskState.FAILED + assert task.state is expected + assert task.done.is_set() and task.host_ready.is_set() + assert all( + holder == StagingBufferHolder.for_step(preserved_sid) + for holders in ctrl._staging_pool._busy + for holder in holders + ) + mgr.discard_step(preserved_sid) + assert all(not holders for holders in ctrl._staging_pool._busy) + with pytest.raises(OmniPrefixCacheUnmatchError, match="write failed"): + mgr.new_step_starts(FakeSchedOut()) + finally: + mgr.shutdown() + + +def test_eager_escalation_failure_releases_deferred_tasks(monkeypatch): + from vllm_omni.core.prefix_cache.controller import TaskState + + mgr, view = make_manager(policy=ModelCachePolicy(needs_full_hidden_states=False, deferred_keys=frozenset({"k"}))) + ctrl = mgr._controller + sid = run_step(mgr, view, {"a": ([0], 0, 4), "b": ([1], 0, 4)}, mm={"k": torch.ones(8, 2)}) + mgr.discard_step(sid) + assert ctrl._staged_bytes > 0 + tasks = list(ctrl._tasks.values()) + + def fail_scatter(task): + raise RuntimeError("injected deferred scatter failure") + + monkeypatch.setattr(ctrl, "_scatter", fail_scatter) + try: + with pytest.raises(RuntimeError, match="injected deferred scatter failure"): + ctrl.escalate([task.tid for task in tasks]) + assert ctrl._staged_bytes == 0 + assert all(task.state is TaskState.FAILED for task in tasks) + assert all(task.done.is_set() and task.host_ready.is_set() for task in tasks) + assert all(not holders for holders in ctrl._staging_pool._busy) + with pytest.raises(OmniPrefixCacheUnmatchError, match="write failed"): + mgr.new_step_starts(FakeSchedOut()) + finally: + mgr.shutdown() + + +def test_per_request_staging_writes(): + mgr, view = make_manager() + sid = run_step(mgr, view, {"p": ([0, 1], 0, 8), "d": ([2], 0, 1)}) + tp = mgr._controller.get_task(next(iter(mgr._request_tasks.tasks["p"]))) + td = mgr._controller.get_task(next(iter(mgr._request_tasks.tasks["d"]))) + assert tp is not None and td is not None + assert (tp.req_id, td.req_id) == ("p", "d") + assert mgr._join_next_step_tids == [tp.tid, td.tid] + outs = mgr.materialize(sid, ["p", "d"]) + assert torch.equal(outs.hidden_states["p"], expected_rows(view.slots_for("p", 0, 8))) + assert torch.equal(outs.hidden_states["d"], expected_rows(view.slots_for("d", 0, 1))) + + +def test_staging_step_prefills_task_host_and_recycles(): + mgr, view = make_manager() + for i in range(6): + sid = run_step(mgr, view, {"a": ([i % 8, (i % 8) + 8], i, 1)}) + ctx = mgr._step_ctxs[sid] + assert ctx.d2h is not None + expect = expected_rows(view.slots_for("a", i, i + 1)) + outs = mgr.materialize(sid, ["a"]) + assert torch.equal(outs.hidden_states["a"], expect), i + + +def test_oversized_step_fails_fast(): + mgr, view = make_manager(staging_capacity_tokens=4) + with pytest.raises(OmniPrefixCacheUnmatchError, match="staging capacity is 4"): + run_step(mgr, view, {"a": ([0, 1], 0, 8)}) + + +def test_staging_slot_eager_save_leaves_only_step_holder(): + """Eager submit writes the pool inline, so the task holder is bound and + released inside save; materialize then empties the slot.""" + mgr, view = make_manager() + sid = run_step(mgr, view, {"p": ([0, 1], 0, 8)}) + ctx = mgr._step_ctxs[sid] + assert ctx.d2h is not None + busy = mgr._controller._staging_pool._busy[ctx.d2h.staging_slot] + assert busy == {StagingBufferHolder.for_step(sid)} + mgr.materialize(sid, ["p"]) + assert not busy + + +def _hold_writes(mgr) -> None: + """Non-eager stand-in: dispatch takes the task to device→host done, but + leaves the pool write to the caller (`mgr._controller._scatter`).""" + + def dispatch(tasks): + for task in tasks: + assert task.claim_copy() + task.mark_host_ready() + + mgr._controller.dispatch = dispatch + + +def test_eager_copy_runs_outside_state_lock_and_key_install_inside(): + """Eager mode (NPU/CPU): the copy + pool write must not run under + _state_lock, or materialize on the builder thread waits a whole D2H; + publishing a new pool key must happen under it.""" + policy = ModelCachePolicy(needs_full_hidden_states=True, deferred_keys=frozenset({"k"})) + mgr, view = make_manager(policy=policy) + seen: dict[str, bool] = {} + + def probe(name): + free = mgr._state_lock.acquire(blocking=False) + if free: + mgr._state_lock.release() + seen[name] = free + + real_run, real_install = mgr._controller._run_eager, mgr._pool.install_key + + def run_eager(task): + probe("run_eager") + real_run(task) + + def install_key(key, storage): + probe("install_key") + real_install(key, storage) + + mgr._controller._run_eager = run_eager + mgr._pool.install_key = install_key + sid = run_step(mgr, view, {"a": ([0], 0, 2)}, mm={"k": torch.ones(2, 2)}) + assert seen == {"run_eager": True, "install_key": False} + mgr.materialize(sid, ["a"]) + assert mgr._pool.has_key("k") and "k" in mgr._slot_status.state + + +def test_eager_finish_escalate_runs_outside_state_lock(): + """Finish/abort escalate must not _run_eager under _state_lock.""" + policy = ModelCachePolicy(needs_full_hidden_states=False, deferred_keys=frozenset({"k"})) + mgr, view = make_manager(policy=policy) + s1 = run_step(mgr, view, {"a": ([2], 0, 1)}, mm={"k": torch.ones(1, 2)}) + mgr.materialize(s1, ["a"]) + seen: list[bool] = [] + real_run = mgr._controller._run_eager + + def run_eager(task): + free = mgr._state_lock.acquire(blocking=False) + if free: + mgr._state_lock.release() + seen.append(free) + real_run(task) + + mgr._controller._run_eager = run_eager + s2 = run_step(mgr, view, {"z": ([9], 0, 1)}, finished=["a"]) + assert seen == [True] + mgr.materialize(s2, ["z"]) + + +def test_staging_task_slot_released_at_pool_write_not_drain(): + mgr, view = make_manager() + _hold_writes(mgr) + sid = run_step(mgr, view, {"p": ([0, 1], 0, 8)}) + ctx = mgr._step_ctxs[sid] + tid = next(iter(mgr._request_tasks.tasks["p"])) + busy = mgr._controller._staging_pool._busy[ctx.d2h.staging_slot] + assert StagingBufferHolder.for_task(tid) in busy and StagingBufferHolder.for_step(sid) in busy + task = mgr._controller.get_task(tid) + mgr._controller._scatter(task) + # Released by the committer's pool write; the manager has not drained yet. + assert busy == {StagingBufferHolder.for_step(sid)} + assert mgr._controller.get_task(tid) is not None + mgr.materialize(sid, ["p"]) + assert not busy + assert mgr._controller.get_task(tid) is None + + +def test_blocked_save_is_released_by_pool_write_alone(): + """All step contexts consumed, the only slot held by a pending write: + the committer finishing that write must unblock the waiting save.""" + mgr, view = make_manager(staging_depth=1, staging_claim_timeout_s=2.0) + _hold_writes(mgr) + s0 = run_step(mgr, view, {"a": ([0], 0, 1)}) + tid = next(iter(mgr._request_tasks.tasks["a"])) + mgr.materialize(s0, ["a"]) + slot = mgr._step_ctxs.get(s0) # consumed + assert slot is None + assert mgr._controller._staging_pool._busy[0] == {StagingBufferHolder.for_task(tid)} + err: list[BaseException] = [] + sid_holder: list[int] = [] + + def _blocked_save() -> None: + try: + sid_holder.append(run_step(mgr, view, {"a": ([0], 1, 1)})) + except BaseException as e: + err.append(e) + + t = threading.Thread(target=_blocked_save, daemon=True) + t.start() + time.sleep(0.1) + assert t.is_alive() + mgr._controller._scatter(mgr._controller.get_task(tid)) + t.join(timeout=1.0) + assert not t.is_alive() + assert not err + mgr.discard_step(sid_holder[0]) + + +def test_hit_prefetch_prebuilds_merged_buffer(): + mgr, view = make_manager() + s1 = run_step(mgr, view, {"a": ([0, 1], 0, 8)}) + mgr.materialize(s1, ["a"]) + s2 = run_step(mgr, view, {"b": ([0, 1, 2], 8, 4)}, new_hits={"b": 8}, finished=["a"]) + ctx = mgr._step_ctxs[s2] + fut = ctx.hit_prefetch["b"][HIDDEN_KEY] + buf = fut.result() + assert buf.shape == (12, HIDDEN) + assert torch.equal(buf[:8], expected_rows(view.slots_for("b", 0, 8))) + merged = mgr.materialize(s2, ["b"]).hidden_states["b"] + assert merged.data_ptr() == buf.data_ptr() + assert torch.equal(merged[8:], expected_rows(view.slots_for("b", 8, 12))) + + +def test_same_step_hit_prefetch_starts_at_save(caplog): + """b hits blocks a computes in the same step: nothing to plan at + new_step_starts (rows ABSENT), planned at publish once a's write is + registered, so the gather runs before the next step can reuse them.""" + mgr, view = make_manager() + view.req_blocks["a"] = [0, 1] + with caplog.at_level(logging.CRITICAL, logger="vllm_omni.core.prefix_cache.manager"): + sid = run_step(mgr, view, {"a": ([0, 1], 0, 8), "b": ([0, 1, 2], 8, 4)}, new_hits={"b": 8}) + assert not any("omni prefix cache unmatch" in r.message for r in caplog.records) + ctx = mgr._step_ctxs[sid] + fut = ctx.hit_prefetch["b"][HIDDEN_KEY] + assert torch.equal(fut.result()[:8], expected_rows(view.slots_for("b", 0, 8))) + outs = mgr.materialize(sid, ["a", "b"]) + assert torch.equal(outs.hidden_states["b"][:8], expected_rows(view.slots_for("b", 0, 8))) + + +def test_delayed_read_of_reassigned_hit_raises(): + """A version check with no COW copy still raises for live and finished. + The production path registers the ref first so remap preserves the rows + (see test_delayed_read_after_committed_reuse_serves_preserved).""" + mgr, view = make_manager() + s1 = run_step(mgr, view, {"a": ([0, 1], 0, 8)}) + mgr.materialize(s1, ["a"]) + slots = view.slots_for("a", 0, 8) + with mgr._state_lock: + planned = mgr._slot_status.get_slot_status(HIDDEN_KEY).slot_version[slots].clone() + mgr._controller.dispatch = lambda tasks: None # new tenant's write stays IN_TRANSIT + s2 = run_step(mgr, view, {"c": ([0, 1], 0, 8)}, finished=["a"]) + for req in ("c", "a"): # live then finished — both fatal without a preserved copy + with pytest.raises(OmniPrefixCacheUnmatchError, match="reassigned"): + mgr._ensure_not_reassigned(slots, HIDDEN_KEY, req_id=req, planned_version=planned) + mgr.discard_step(s2) + + +def test_delayed_read_after_committed_reuse_serves_preserved(): + """Amy's block-0 reuse: plan A's read, then B commits over the same + slots. The reusing write copy-on-writes A's committed rows into the + planned ref; the delayed fetch must return those, not B's.""" + mgr, view = make_manager() + sa = run_step(mgr, view, {"a": ([0], 0, 4)}) + mgr.materialize(sa, ["a"]) # A committed into block 0 + slots = view.slots_for("a", 0, 4) + with mgr._state_lock: + src = mgr._slot_ref(slots, HIDDEN_KEY, "a") # plan A's read + assert src.already_staged and not src.staged_list and not src.join_tids + sb = run_step(mgr, view, {"b": ([0], 0, 4)}, finished=["a"]) + mgr.materialize(sb, ["b"]) # A freed, block 0 reused by B and COMMITTED + # run_step derives rows from slot ids, so A's and B's values coincide. + # Scribble the pool so only the preserved copy can produce A's rows. + assert set(src.preserved) == {int(s) for s in slots.tolist()} + with torch.inference_mode(): + mgr._pool.write(HIDDEN_KEY, slots, torch.full((4, HIDDEN), 99.0, dtype=DTYPE)) + got = mgr._fetch_source(src) + assert torch.equal(got, expected_rows(slots)) + assert all(ref is not src for ref in mgr._pending_reads) + + +def test_done_fires_only_after_completion_record_published(): + """A join(done) waiter (the save barrier) must find the completion record + on its next drain — done-before-append leaves occupancy IN_TRANSIT and the + COW preserve skips the reused rows.""" + mgr, view = make_manager() + checks = [] + orig_scatter = mgr._controller._scatter + + def scatter(task): + orig_set = task.done.set + + def checked_set(): + checks.append(task.tid in mgr._controller._completed) + orig_set() + + task.done.set = checked_set + orig_scatter(task) + + mgr._controller._scatter = scatter + sid = run_step(mgr, view, {"a": ([0], 0, 4)}) + mgr.materialize(sid, ["a"]) + assert checks and all(checks) + + +def test_failed_wakes_only_after_failure_record_published(): + """Same window as the done/_completed one, failure side: a joiner woken + by FAILED must find the record on drain, not read never-written rows.""" + mgr, view = make_manager() + _hold_writes(mgr) # keep the write non-terminal + run_step(mgr, view, {"a": ([0], 0, 4)}) + tid = next(iter(mgr._request_tasks.tasks["a"])) + task = mgr._controller.get_task(tid) + checks = [] + orig_set = task.done.set + + def checked_set(): + checks.append(tid in mgr._controller._failed) + orig_set() + + task.done.set = checked_set + mgr._controller._fail_task(tid) + assert checks and all(checks) + + +def test_save_waits_pool_write_of_previous_step(): + """A save must not reuse slots whose JOIN_NEXT_STEP write has not reached + the pool — no delayed read could recover such a row.""" + mgr, view = make_manager(staging_claim_timeout_s=0.2) + _hold_writes(mgr) # host_ready only; the pool write never happens + run_step(mgr, view, {"a": ([0], 0, 4)}) + with pytest.raises(OmniPrefixCacheUnmatchError, match="did not reach done"): + run_step(mgr, view, {"b": ([1], 0, 4)}) + + +def test_join_next_step_hit_survives_task_already_drained(): + mgr, view = make_manager() + mgr._controller.dispatch = lambda tasks: None # registered, never run + view.req_blocks["a"] = [0, 1] + sid = run_step(mgr, view, {"a": ([0, 1], 0, 8), "b": ([0, 1, 2], 8, 4)}, new_hits={"b": 8}) + with mgr._state_lock: + src = mgr._slot_ref(view.slots_for("b", 0, 8), HIDDEN_KEY, "b") + for tid in src.join_tids: + mgr._controller._run_eager(mgr._controller.get_task(tid)) + with torch.inference_mode(): + with mgr._state_lock: + mgr._commit_drained_writes() + assert mgr._controller.get_task(src.join_tids[0]) is None + assert torch.equal(mgr._fetch_source(src), expected_rows(view.slots_for("b", 0, 8))) + mgr.materialize(sid, ["a", "b"]) + + +def test_leftover_mm_snapshot_survives_live_overwrite(): + """Leftover mm is copied at save; mutating the live buffer + afterwards must not change the snapshot.""" + mgr, view = make_manager() + live = [torch.arange(10, dtype=DTYPE)] + sid = run_step(mgr, view, {"a": ([0], 0, 2)}, mm={"codes.ref": live}) + assert "codes.ref" in mgr._step_ctxs[sid].mm_cpu_snapshot + snap = mgr._step_ctxs[sid].mm_cpu_snapshot["codes.ref"] + live[0].fill_(99.0) + assert torch.equal(snap[0], torch.arange(10, dtype=DTYPE)) + mgr.discard_step(sid) + + +@pytest.mark.parametrize("needs_full_hidden_states", [True, False]) +def test_deferred_leftover_snapshot_survives_live_overwrite(needs_full_hidden_states): + """Deferred tails are leftover-copied at save for materialize; the GPU + freeze is a different product (pool write on finish). ``False`` is the + hidden opt-out + miss path (Qwen3-TTS / Higgs talker).""" + policy = ModelCachePolicy( + needs_full_hidden_states=needs_full_hidden_states, deferred_keys=frozenset({"codes.audio"}) + ) + mgr, view = make_manager(policy=policy) + live = torch.full((2, 2), 1.0) + ref = [torch.arange(4, dtype=DTYPE)] + sid = run_step(mgr, view, {"a": ([0], 0, 2)}, mm={"codes.audio": live, "codes.ref": ref}) + assert "codes.audio" in mgr._step_ctxs[sid].mm_cpu_snapshot + live.fill_(99.0) + ref[0].fill_(99.0) + outs = mgr.materialize(sid, ["a"]) + assert torch.equal(outs.mm_outputs["codes.audio"]["a"], torch.full((2, 2), 1.0)) + assert torch.equal(outs.mm_outputs["codes.ref"]["a"][0], torch.arange(4, dtype=DTYPE)) + assert (outs.hidden_states is None) == (not needs_full_hidden_states) + + +def test_frozen_mm_clone_survives_live_overwrite(): + mgr, view = make_manager() + live = torch.full((2, 2), 1.0) + sid = run_step(mgr, view, {"a": ([0], 0, 2)}, mm={"codes.audio": live}) + assert "codes.audio" not in mgr._step_ctxs[sid].mm_cpu_snapshot + live.fill_(99.0) + outs = mgr.materialize(sid, ["a"]) + assert torch.equal(outs.mm_outputs["codes.audio"]["a"], torch.full((2, 2), 1.0)) + + +def test_hidden_opt_out_miss_returns_mm_and_releases_slot(): + policy = ModelCachePolicy(needs_full_hidden_states=False) + mgr, view = make_manager(policy=policy) + for i in range(6): + sid = run_step(mgr, view, {"a": ([i % 8, (i % 8) + 8], i, 1)}, mm={"k": torch.full((1, 2), float(i))}) + ctx = mgr._step_ctxs[sid] + assert ctx.d2h is not None, i + outs = mgr.materialize(sid, ["a"]) + assert outs.hidden_states is None + assert torch.equal(outs.mm_outputs["k"]["a"], torch.full((1, 2), float(i))) + assert not mgr._controller._staging_pool._busy[ctx.d2h.staging_slot], i + + +def test_save_releases_staging_if_commit_drained_writes_fails(): + mgr, view = make_manager(staging_depth=2) + calls = {"n": 0} + real = mgr._commit_drained_writes + + def wrapped(): + calls["n"] += 1 + if calls["n"] == 2: + raise OmniPrefixCacheUnmatchError("injected fail") + return real() + + mgr._commit_drained_writes = wrapped + with pytest.raises(OmniPrefixCacheUnmatchError, match="injected fail"): + run_step(mgr, view, {"a": ([0], 0, 4)}) + assert all(not busy for busy in mgr._controller._staging_pool._busy) + mgr._commit_drained_writes = real + sid = run_step(mgr, view, {"a": ([0], 0, 4)}) + assert mgr._step_ctxs[sid].d2h is not None + mgr.materialize(sid, ["a"]) + + +def test_fetch_host_maps_slots_across_layouts(): + """Identity, prefix, gapped, and a two-`_WriteChunk` gather.""" + from vllm_omni.core.prefix_cache.block_pool import PrefixBlockPool + from vllm_omni.core.prefix_cache.controller import OmniPrefixCacheController, WriteTask, _WriteChunk + + cfg = PrefixCacheConfig(num_blocks=NUM_BLOCKS, block_size=BLOCK_SIZE) + pool = PrefixBlockPool(cfg) + pool.ensure_key(HIDDEN_KEY, DTYPE, HIDDEN) + ctrl = OmniPrefixCacheController(pool, cfg, eager=True) + + slots = torch.tensor([40, 41, 42, 43, 44], dtype=torch.int64) + rows = torch.arange(slots.numel() * HIDDEN, dtype=DTYPE).reshape(slots.numel(), HIDDEN) + + def _expect(want): + return torch.stack([rows[(slots == s).nonzero()[0, 0]] for s in want.tolist()]) + + one = WriteTask( + tid=1, + req_id="r", + write_n=1, + schedule=WriteSchedule.JOIN_ON_FINISH, + chunks=[_WriteChunk(slots_cpu=slots, tensors={HIDDEN_KEY: rows})], + ) + for want in (slots, slots[:3], slots[[0, 1, 3, 4]]): + assert torch.equal(ctrl.fetch_host(one, want, HIDDEN_KEY), _expect(want)) + + two = WriteTask( + tid=2, + req_id="r", + write_n=2, + schedule=WriteSchedule.JOIN_ON_FINISH, + chunks=[ + _WriteChunk(slots_cpu=slots[:3], tensors={HIDDEN_KEY: rows[:3]}), + _WriteChunk(slots_cpu=slots[3:], tensors={HIDDEN_KEY: rows[3:]}), + ], + ) + want = slots[[0, 3, 1, 4]] + assert torch.equal(ctrl.fetch_host(two, want, HIDDEN_KEY), _expect(want)) + + +def test_fetch_host_waits_staging_step_d2h_event(): + """JOIN_NEXT_STEP stores chunk.host as a staging view; fetch_host waits + step_d2h_event before slicing. Production JOIN_NEXT_STEP hits wait for + the pool write instead; this is the unit-level wait contract for that branch.""" + from vllm_omni.core.prefix_cache.block_pool import PrefixBlockPool + from vllm_omni.core.prefix_cache.controller import OmniPrefixCacheController, WriteTask, _WriteChunk + + cfg = PrefixCacheConfig(num_blocks=NUM_BLOCKS, block_size=BLOCK_SIZE) + pool = PrefixBlockPool(cfg) + ctrl = OmniPrefixCacheController(pool, cfg, eager=True) + + slots = torch.tensor([0, 1, 2], dtype=torch.int64) + src = torch.arange(3 * HIDDEN, dtype=DTYPE).reshape(3, HIDDEN) + landing = torch.zeros_like(src) + + class _HostEvent: + n = 0 + + def synchronize(self): + self.n += 1 + landing.copy_(src) + + event = _HostEvent() + chunk = _WriteChunk(slots_cpu=slots, tensors={HIDDEN_KEY: src}) + chunk.host = {HIDDEN_KEY: landing} + task = WriteTask( + tid=1, + req_id="r", + write_n=1, + schedule=WriteSchedule.JOIN_NEXT_STEP, + chunks=[chunk], + staging_slot=0, + step_d2h_event=event, + ) + rows = ctrl.fetch_host(task, slots, HIDDEN_KEY) + assert event.n == 1 + assert torch.equal(rows, src) + + +def _bare_controller(eager=True): + from vllm_omni.core.prefix_cache.block_pool import PrefixBlockPool + from vllm_omni.core.prefix_cache.controller import OmniPrefixCacheController + + cfg = PrefixCacheConfig(num_blocks=NUM_BLOCKS, block_size=BLOCK_SIZE) + pool = PrefixBlockPool(cfg) + pool.ensure_key(HIDDEN_KEY, DTYPE, HIDDEN) + ctrl = OmniPrefixCacheController(pool, cfg, eager=eager) + return ctrl, pool + + +def _bare_task(tid=7, schedule=WriteSchedule.JOIN_NEXT_STEP, host=None, tensors=None, budget=None, queued=False): + from vllm_omni.core.prefix_cache.controller import TaskState, WriteTask, _WriteChunk + + slots = torch.tensor([0, 1], dtype=torch.int64) + chunk = _WriteChunk(slots_cpu=slots, tensors=tensors or {}, budget=budget) + if host is not None: + chunk.host = {HIDDEN_KEY: host} + task = WriteTask(tid=tid, req_id="r", write_n=1, schedule=schedule, chunks=[chunk]) + if queued: + task.transition(TaskState.QUEUED) + return task, slots + + +def test_task_state_is_a_strict_chain(): + """PENDING -> QUEUED -> COPYING -> HOST_READY -> WRITTEN, one step at a time.""" + from vllm_omni.core.prefix_cache.controller import TaskState + + task, _ = _bare_task() + assert task.state is TaskState.PENDING + assert not task.claim_copy() # PENDING may not skip QUEUED + with pytest.raises(OmniPrefixCacheUnmatchError, match="illegal transition PENDING -> HOST_READY"): + task.transition(TaskState.HOST_READY) + task.transition(TaskState.QUEUED) + assert task.claim_copy() + assert task.state is TaskState.COPYING and not task.host_ready.is_set() + assert not task.claim_copy() # second claimant loses + task.mark_host_ready() + assert task.state is TaskState.HOST_READY and task.host_ready.is_set() and not task.done.is_set() + with pytest.raises(OmniPrefixCacheUnmatchError, match="illegal transition HOST_READY -> COPYING"): + task.transition(TaskState.COPYING) # no going back + task.mark_done() + assert task.state is TaskState.WRITTEN and task.done.is_set() + assert not task.mark_failed() # terminal stays terminal + assert task.state is TaskState.WRITTEN + + +def test_task_failed_sets_both_events_and_closes_append(): + from vllm_omni.core.prefix_cache.controller import TaskState, _WriteChunk + + task, slots = _bare_task() + assert task.mark_failed() + assert task.state is TaskState.FAILED and task.host_ready.is_set() and task.done.is_set() + assert task.append_chunk(_WriteChunk(slots_cpu=slots, tensors={})) is TaskState.FAILED + + +def test_escalate_leaves_claimed_task_alone(): + """Once the worker owns the copy stage (COPYING), escalate must not + re-queue the tid — that re-queue was the double-write race.""" + from vllm_omni.core.prefix_cache.controller import TaskState + + ctrl, _ = _bare_controller(eager=True) + ctrl._eager = False # exercise the queue path without a worker thread + later, _ = _bare_task(tid=1, schedule=WriteSchedule.JOIN_ON_FINISH) + ctrl.submit(later, queued=False) + assert later.state is TaskState.PENDING + ctrl.escalate([1]) + assert later.state is TaskState.QUEUED and list(ctrl._queue_hi) == [1] + ctrl.escalate([1]) # already high-priority: no duplicate + assert list(ctrl._queue_hi) == [1] + + ctrl._queue_hi.clear() + later.transition(TaskState.COPYING) + ctrl.escalate([1]) + assert list(ctrl._queue_hi) == [] and list(ctrl._queue_lo) == [] + + +def test_join_on_stuck_committer_raises_instead_of_hanging(): + from vllm_omni.core.prefix_cache.block_pool import PrefixBlockPool + from vllm_omni.core.prefix_cache.controller import OmniPrefixCacheController, TaskState + + cfg = PrefixCacheConfig(num_blocks=NUM_BLOCKS, block_size=BLOCK_SIZE, staging_claim_timeout_s=0.05) + ctrl = OmniPrefixCacheController(PrefixBlockPool(cfg), cfg, eager=True) + ctrl._eager = False # queue path without a worker thread: nothing will finish the task + task, _ = _bare_task(tid=1, schedule=WriteSchedule.JOIN_ON_FINISH) + ctrl.submit(task, queued=True) + task.transition(TaskState.COPYING) + + with pytest.raises( + OmniPrefixCacheUnmatchError, match=r"task 1 \(join_on_finish\) did not reach host_ready .*state=COPYING" + ): + ctrl.join_host_ready([1]) + with pytest.raises(OmniPrefixCacheUnmatchError, match="did not reach done"): + ctrl.join([1]) + + +def test_scatter_host_ready_writes_only_staged_tasks(): + from vllm_omni.core.prefix_cache.controller import TaskState + + ctrl, pool = _bare_controller() + task, slots = _bare_task(host=torch.ones(2, HIDDEN), queued=True) + ctrl._tasks[task.tid] = task + assert task.claim_copy() + ctrl._blocked.append(task.tid) + ctrl._scatter_host_ready() # COPYING: not ready, stays blocked + assert ctrl._blocked == [task.tid] and list(ctrl._completed) == [] + task.mark_host_ready() + ctrl._scatter_host_ready() + assert ctrl._blocked == [] and list(ctrl._completed) == [task.tid] + assert task.state is TaskState.WRITTEN + assert torch.equal(pool.rows(HIDDEN_KEY, slots), torch.ones(2, HIDDEN)) + + +def test_budget_ticket_uncharges_once_per_clone(): + from vllm_omni.core.prefix_cache.controller import _BudgetTicket + + ctrl, _ = _bare_controller() + ticket = _BudgetTicket(nbytes=100) + ctrl.reserve(100) + a, _ = _bare_task(tid=1, budget=ticket) + b, _ = _bare_task(tid=2, budget=ticket) + ctrl.pin_budget(ticket, 1) + ctrl.pin_budget(ticket, 2) + ctrl._release_staged_bytes(a) + assert ctrl._staged_bytes == 100 # b still holds the clone + ctrl._release_staged_bytes(a) # idempotent per tid + assert ctrl._staged_bytes == 100 + ctrl._release_staged_bytes(b) + assert ctrl._staged_bytes == 0 + ctrl._release_staged_bytes(b) # freed once, never twice + assert ctrl._staged_bytes == 0 + + +def test_fail_task_after_copy_does_not_uncharge_twice(): + from vllm_omni.core.prefix_cache.controller import TaskState, _BudgetTicket + + ctrl, _ = _bare_controller() + ticket = _BudgetTicket(nbytes=64) + ctrl.reserve(64) + task, _ = _bare_task(tid=3, budget=ticket, queued=True) + ctrl._tasks[task.tid] = task + ctrl.pin_budget(ticket, 3) + assert task.claim_copy() + task.mark_host_ready() + ctrl._release_staged_bytes(task) # what _copy_task does on HOST_READY + assert ctrl._staged_bytes == 0 + ctrl._fail_task(3) + assert task.state is TaskState.FAILED and ctrl._staged_bytes == 0 + assert list(ctrl._failed) == [3] + ctrl._fail_task(3) # terminal: no second publish + assert list(ctrl._failed) == [3] + + +def test_from_vllm_config_uses_batched_tokens(): + cfg = PrefixCacheConfig.from_vllm_config( num_blocks=NUM_BLOCKS, block_size=BLOCK_SIZE, - hidden_size=HIDDEN_SIZE, - hs_dtype=DTYPE, + scheduler_config=SimpleNamespace(max_num_batched_tokens=8192, max_model_len=32768, max_num_seqs=64), ) - return cache - - -def get_multimodal_outputs(feat_dims: dict[str, int], seq_len: int) -> dict[str, torch.Tensor]: - fake_mm_inputs = {} - for mm_key, feat_dim in feat_dims.items(): - fake_mm_inputs[mm_key] = torch.rand((seq_len, feat_dim), dtype=DTYPE) - return fake_mm_inputs - - -### Tests for initialization -def test_initialization_simple(): - """Check default initialization only creates the hidden states.""" - cache = get_omni_pcache() - assert isinstance(cache.hidden_states_cache, torch.Tensor) - assert cache.hidden_states_cache.shape == DEFAULT_SHAPE - assert len(cache.mm_outputs_cache) == 0 - assert len(cache.mm_cache_keys) == 0 - - -def test_initialization_with_multimodal(): - """Check initialization + registration of multimodal outputs.""" - cache = get_omni_pcache() - feat_dims = {"foo": 100, "bar": 50, "baz": 10} - mm_outputs = get_multimodal_outputs( - feat_dims, - seq_len=DEFAULT_SEQ_LEN, + assert cfg.staging_capacity_tokens == 8192 + assert cfg.staging_depth == 4 + cfg_fallback = PrefixCacheConfig.from_vllm_config( + num_blocks=NUM_BLOCKS, + block_size=BLOCK_SIZE, + scheduler_config=SimpleNamespace(max_num_batched_tokens=None, max_model_len=4096), ) - # Cast one of the keys to a different dtype; the dtype of the tensor - # that is used to initialize the cache dictates the cache dtype. - mm_outputs["foo"] = mm_outputs["foo"].to(OTHER_DTYPE) - - cache.maybe_init_missing_mm_cache_keys(mm_outputs, DEFAULT_SEQ_LEN) - assert len(cache.mm_cache_keys) == 3 - assert set(cache.mm_cache_keys) == set(feat_dims.keys()) - for mm_key in cache.mm_cache_keys: - cache_tensor = cache.mm_outputs_cache[mm_key] - assert isinstance(cache_tensor, torch.Tensor) - assert cache_tensor.shape[-1] == feat_dims[mm_key] - assert mm_outputs[mm_key].dtype == cache_tensor.dtype - - -def test_init_missing_mm_cache_keys_is_idempotent(): - """Ensure that the cache doesn't reinitialize old keys.""" - cache = get_omni_pcache() - mm_key = "foo" - feat_dims = {mm_key: 100} - mm_outputs = get_multimodal_outputs( - feat_dims, - seq_len=DEFAULT_SEQ_LEN, - ) - cache.maybe_init_missing_mm_cache_keys(mm_outputs, DEFAULT_SEQ_LEN) - assert len(cache.mm_cache_keys) == 1 - assert mm_key in cache.mm_cache_keys - - # Cache is initialized to 0 - fill it with 1s - cache.mm_outputs_cache[mm_key].fill_(1) - - # Ensure that running another initialization - # doesn't zero out our cache values - cache.maybe_init_missing_mm_cache_keys(mm_outputs, DEFAULT_SEQ_LEN) - assert len(cache.mm_cache_keys) == 1 - assert mm_key in cache.mm_cache_keys - assert torch.all(cache.mm_outputs_cache[mm_key] == 1) - - -### Tests for Update -def test_update_no_multimodal(): - """Test that slot mappings act as row indices hidden states.""" - cache = get_omni_pcache() - - num_tokens_unpadded = 8 - slot_offset = 8 - slot_mapping = torch.arange(slot_offset, slot_offset + num_tokens_unpadded) - new_hidden_states = torch.rand((num_tokens_unpadded, HIDDEN_SIZE), dtype=DTYPE) - - cache.update_omni_tensor_prefix_cache( - hidden_states=new_hidden_states, - multimodal_outputs=None, - num_tokens_unpadded=num_tokens_unpadded, - slot_mapping=slot_mapping, - ) - - # Ensure that if we reshape our 3D cache back to 2D, we can use the - # indices in our slot mappings to access the hidden states as expected - hs_rows = cache.hidden_states_cache.view(NUM_BLOCKS * BLOCK_SIZE, HIDDEN_SIZE) - for slot_idx, new_states in zip(slot_mapping, new_hidden_states): - slot_states = hs_rows[slot_idx] - 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", - [ - {"foo": 100, "bar": 100}, - {"foo": 100, "bar": 50, "baz": 10}, - ], -) -def test_update_with_multimodal_outputs(feat_dims): - """Test that slot mappings are correct for multimodal tensors.""" - cache = get_omni_pcache_with_mm_tensors(feat_dims, seq_len=DEFAULT_SEQ_LEN) - - num_tokens_unpadded = 8 - slot_offset = 8 - slot_mapping = torch.arange(slot_offset, slot_offset + num_tokens_unpadded) - feature_dims = {key: val.shape[-1] for key, val in cache.mm_outputs_cache.items()} - mm_outputs = {key: torch.rand((num_tokens_unpadded, feature_dims[key]), dtype=DTYPE) for key in cache.mm_cache_keys} - cache.update_omni_tensor_prefix_cache( - hidden_states=None, - multimodal_outputs=mm_outputs, - num_tokens_unpadded=num_tokens_unpadded, - slot_mapping=slot_mapping, - ) - - for mm_key in feat_dims.keys(): - assert mm_key in cache.mm_outputs_cache - key_feat_dim = feature_dims[mm_key] - mm_state_rows = cache.mm_outputs_cache[mm_key].view(NUM_BLOCKS * BLOCK_SIZE, key_feat_dim) - - # Similar to hidden states, but for each key in the dict; - # Different tensors may have different feature dims - new_mm_outputs = mm_outputs[mm_key] - for slot_idx, new_output in zip(slot_mapping, new_mm_outputs): - slot_states = mm_state_rows[slot_idx] - 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_discard_deferred_mm_outputs_drops_pending_chunks(): - """discard_deferred_mm_outputs is the abort path: a request that - leaves without a cache commit must not leak its staged chunks into - the cache if commit_deferred_mm_outputs is later called for it.""" - 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, - ) - - cache.update_omni_tensor_prefix_cache( - hidden_states=None, - multimodal_outputs={mm_key: torch.arange(8, dtype=torch.long).reshape(4, 2)}, - num_tokens_unpadded=4, - slot_mapping=torch.arange(8, 12), - skip_mm_cache_keys={mm_key}, - ) - - chunk = torch.arange(8, dtype=torch.long).reshape(4, 2) - cache.stage_deferred_mm_outputs( - query_start_loc=torch.tensor([0, 4], dtype=torch.long), - input_batch=input_batch, - multimodal_outputs={mm_key: chunk}, - num_scheduled_tokens={"req1": 4, "req2": 0}, - deferred_mm_cache_keys={mm_key}, - ) - assert "req1" in cache._deferred_mm_outputs - - cache.discard_deferred_mm_outputs("req1") - assert "req1" not in cache._deferred_mm_outputs - - # A later commit call for the now-discarded req_id must be a no-op, - # not a resurrection of the dropped chunk. - cache.commit_deferred_mm_outputs({"req1"}, input_batch) - rows = cache.mm_outputs_cache[mm_key].view(-1, 2) - assert torch.all(rows[8:12] == 0) - - -def test_deferred_mm_outputs_request_id_reused_after_discard(): - """A request ID reused after discard_deferred_mm_outputs must start - clean: the new occupant's committed data must not be mixed with the - previous (discarded) occupant's staged chunks.""" - 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, - ) - - cache.update_omni_tensor_prefix_cache( - hidden_states=None, - multimodal_outputs={mm_key: torch.arange(8, dtype=torch.long).reshape(4, 2)}, - num_tokens_unpadded=4, - slot_mapping=torch.arange(8, 12), - skip_mm_cache_keys={mm_key}, - ) - - stale_chunk = torch.full((4, 2), -1, 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: stale_chunk}, - num_scheduled_tokens={"req1": 4, "req2": 0}, - deferred_mm_cache_keys={mm_key}, - ) - cache.discard_deferred_mm_outputs("req1") - - # "req1" is reused by a brand-new request with different data. - fresh_chunk = torch.arange(8, dtype=torch.long).reshape(4, 2) - cache.stage_deferred_mm_outputs( - query_start_loc=torch.tensor([0, 4], dtype=torch.long), - input_batch=input_batch, - multimodal_outputs={mm_key: fresh_chunk}, - num_scheduled_tokens={"req1": 4, "req2": 0}, - deferred_mm_cache_keys={mm_key}, - ) - cache.commit_deferred_mm_outputs({"req1"}, input_batch) - - rows = cache.mm_outputs_cache[mm_key].view(-1, 2) - assert torch.equal(rows[8:12], fresh_chunk) - - -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. - - Assumption: - req_idx 0 is a cache hit with slots 8, 9, ..., 15 - req_idx 1 is a cache miss - """ - assert req_idx < 2 - if req_idx == 0: - # With the slot offset we provided (8), the corresponding - # blocks IDs are 2 & 3 because the block size is 4. - return torch.tensor([2, 3], dtype=torch.long) - return torch.tensor([], dtype=torch.long) - - -@pytest.mark.parametrize("num_tokens_padded", [None, 16]) -def test_get_merged_hidden_states(num_tokens_padded, mocker): - """Ensure that hidden states are merged correctly.""" - cache = get_omni_pcache() - - orig_num_tokens_unpadded = 8 - slot_offset = 8 # We'll put our states in slots 8, 9, 10, ..., 15 - 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_tokens_padded=num_tokens_padded, - ) - - # Say that we have two requests, but only one of them is a cache hit - 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.rand( - (num_new_toks_req1 + num_new_toks_req2, HIDDEN_SIZE), - dtype=DTYPE, - ) - req1_new_states = new_hidden_states[:num_new_toks_req1] - req2_new_states = new_hidden_states[-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, - num_scheduled_tokens=num_scheduled_tokens, - ) - - assert "req1" in merged_states and "req2" in merged_states - req1_merged_states = merged_states["req1"] - req2_merged_states = merged_states["req2"] - - # First, check the cache hit case - assert req1_merged_states.shape == torch.Size([orig_num_tokens_unpadded + num_new_toks_req1, HIDDEN_SIZE]) - # Ensure that the req1 merged states are the cached states + the new req1 states - assert torch.all(req1_merged_states[:orig_num_tokens_unpadded] == orig_hidden_states) - assert torch.all(req1_merged_states[-num_new_toks_req1:] == req1_new_states) - - # Next, ensure that the cache miss case only has the new states - assert req2_merged_states.shape == torch.Size([num_new_toks_req2, HIDDEN_SIZE]) - 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", - [ - {"foo": 100, "bar": 100}, - {"foo": 100, "bar": 50, "baz": 10}, - ], -) -def test_get_merged_multimodal_outputs(feat_dims, num_tokens_padded, mocker): - cache = get_omni_pcache_with_mm_tensors(feat_dims, seq_len=DEFAULT_SEQ_LEN) - - orig_num_tokens_unpadded = 8 - slot_offset = 8 # We'll put our states in slots 8, 9, 10, ..., 15 - orig_slot_mapping = torch.arange(slot_offset, slot_offset + orig_num_tokens_unpadded) - feature_dims = {key: val.shape[-1] for key, val in cache.mm_outputs_cache.items()} - orig_mm_outputs = { - key: torch.rand((orig_num_tokens_unpadded, feature_dims[key]), dtype=DTYPE) for key in cache.mm_cache_keys - } - - cache.update_omni_tensor_prefix_cache( - hidden_states=None, - multimodal_outputs=orig_mm_outputs, - num_tokens_unpadded=orig_num_tokens_unpadded, - slot_mapping=orig_slot_mapping, - num_tokens_padded=num_tokens_padded, - ) - - # Similar to hs test- say that we have two requests, but only one of them is a cache hit - 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_mm_outputs = {} - for mm_key in cache.mm_cache_keys: - new_mm_outputs[mm_key] = torch.rand( - (num_new_toks_req1 + num_new_toks_req2, feature_dims[mm_key]), - dtype=DTYPE, - ) - # We also want to make sure passthrough data (outside of our keys) isn't dropped - new_mm_outputs["passthrough_data"] = "Something else" - # Lists are a special case because we can't split them yet if we want to match - # the nonprefix cache behavior, because this runs before post process. - new_mm_outputs["passthrough_list"] = ["should", "not", "split"] - - 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_mm_outputs = cache.get_merged_multimodal_states( - query_start_loc=[0, num_new_toks_req1], - input_batch=input_batch, - multimodal_outputs=new_mm_outputs, - num_scheduled_tokens=num_scheduled_tokens, - ) - - # Ensure the passthrough data wasn't dropped - assert "passthrough_data" in merged_mm_outputs - assert "passthrough_list" in merged_mm_outputs - - for mm_key, mm_output in merged_mm_outputs.items(): - # Ensure passthrough data is just forwarded normally and not duplicated - assert isinstance(mm_output, dict) - assert "req1" in mm_output and "req2" in mm_output - if mm_key == "passthrough_data": - assert mm_key not in cache.mm_cache_keys - assert new_mm_outputs[mm_key] == mm_output["req1"] - assert new_mm_outputs[mm_key] == mm_output["req2"] - elif mm_key == "passthrough_list": - assert mm_key not in cache.mm_cache_keys - assert new_mm_outputs[mm_key] == mm_output["req1"] - assert new_mm_outputs[mm_key] == mm_output["req2"] - else: - assert mm_key in cache.mm_cache_keys - curr_feat_dim = feature_dims[mm_key] - # Ensure that req1 (cache hit) merged the mm data - req1_merged_mm_outputs = mm_output["req1"] - req1_new_mm_outputs = new_mm_outputs[mm_key][:num_new_toks_req1] - - assert req1_merged_mm_outputs.shape == torch.Size( - [orig_num_tokens_unpadded + num_new_toks_req1, curr_feat_dim] - ) - # Ensure that the req1 merged mm data are the cached data + the new data - assert torch.all(req1_merged_mm_outputs[:orig_num_tokens_unpadded] == orig_mm_outputs[mm_key]) - assert torch.all(req1_merged_mm_outputs[-num_new_toks_req1:] == req1_new_mm_outputs) - - # Ensure that req2 (cache miss) only has the new mm data - req2_merged_mm_outputs = mm_output["req2"] - req2_new_mm_outputs = new_mm_outputs[mm_key][-num_new_toks_req2:] - - assert req2_merged_mm_outputs.shape == torch.Size([num_new_toks_req2, curr_feat_dim]) - assert torch.all(req2_merged_mm_outputs == req2_new_mm_outputs) - - -@pytest.mark.parametrize( - "passthrough_shape", - [(5,), (5, 3)], - ids=["1d_token_metadata", "2d_feature_tensor"], -) -def test_get_merged_multimodal_outputs_slices_passthrough_tensor_per_request(passthrough_shape, mocker): - """A passthrough key that is itself a real tensor (not the string/list - passthrough already covered above) must be sliced per request through - to_payload_element's tensor branch, under the same mixed cache - hit/miss + uneven scheduled-length batch as test_get_merged_multimodal_outputs. - - Covers both a 1D token-aligned metadata tensor (e.g. ref_code_len, - codec_streaming -- shape (total_tokens,)) and a 2D per-token feature - tensor (shape (total_tokens, feat_dim)), since W2's batching could - easily special-case one and silently mishandle the other. - """ - feat_dims = {"foo": 100, "bar": 100} - cache = get_omni_pcache_with_mm_tensors(feat_dims, seq_len=DEFAULT_SEQ_LEN) - - orig_num_tokens_unpadded = 8 - slot_offset = 8 - orig_slot_mapping = torch.arange(slot_offset, slot_offset + orig_num_tokens_unpadded) - feature_dims = {key: val.shape[-1] for key, val in cache.mm_outputs_cache.items()} - orig_mm_outputs = { - key: torch.rand((orig_num_tokens_unpadded, feature_dims[key]), dtype=DTYPE) for key in cache.mm_cache_keys - } - - cache.update_omni_tensor_prefix_cache( - hidden_states=None, - multimodal_outputs=orig_mm_outputs, - num_tokens_unpadded=orig_num_tokens_unpadded, - slot_mapping=orig_slot_mapping, - num_tokens_padded=None, - ) - - 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_mm_outputs = {} - for mm_key in cache.mm_cache_keys: - new_mm_outputs[mm_key] = torch.rand( - (num_new_toks_req1 + num_new_toks_req2, feature_dims[mm_key]), - dtype=DTYPE, - ) - total_scheduled_tokens = num_new_toks_req1 + num_new_toks_req2 - assert passthrough_shape[0] == total_scheduled_tokens - passthrough_tensor = torch.arange(int(torch.prod(torch.tensor(passthrough_shape))), dtype=DTYPE).reshape( - passthrough_shape - ) - new_mm_outputs["passthrough_tensor"] = passthrough_tensor - - 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_mm_outputs = cache.get_merged_multimodal_states( - query_start_loc=[0, num_new_toks_req1], - input_batch=input_batch, - multimodal_outputs=new_mm_outputs, - num_scheduled_tokens=num_scheduled_tokens, - ) - - assert "passthrough_tensor" not in cache.mm_cache_keys - - passthrough_out = merged_mm_outputs["passthrough_tensor"] - # req1 (cache hit) gets tokens [0:3), req2 (cache miss) gets tokens [3:5) -- - # a passthrough tensor is sliced by scheduled position, not merged with - # any cached history (unlike the mm_cache_keys tensors above). - assert torch.equal(passthrough_out["req1"], passthrough_tensor[0:num_new_toks_req1]) - assert torch.equal(passthrough_out["req2"], passthrough_tensor[num_new_toks_req1:total_scheduled_tokens]) - - -class MergeHarness(NamedTuple): - """A primed cache plus the arguments get_merged_multimodal_states needs.""" - - cache: OmniTensorPrefixCache - input_batch: MockInputBatch - num_scheduled_tokens: dict[str, int] - query_start_loc: list[int] - total_scheduled_tokens: int - new_mm_outputs: dict[str, object] - - -def build_mixed_hit_miss_harness(mocker, num_new_toks_req1: int = 3, num_new_toks_req2: int = 2) -> MergeHarness: - """Two requests over one cache key: req1 hits 8 cached tokens, req2 misses, - and their scheduled lengths differ -- the same batch shape as - test_get_merged_multimodal_outputs, which is what makes a per-request - boundary error observable. Callers add their passthrough keys to - ``new_mm_outputs`` before calling get_merged_multimodal_states. - """ - cache = get_omni_pcache_with_mm_tensors({"foo": 100}, seq_len=DEFAULT_SEQ_LEN) - - orig_num_tokens_unpadded = 8 - feature_dims = {key: val.shape[-1] for key, val in cache.mm_outputs_cache.items()} - cache.update_omni_tensor_prefix_cache( - hidden_states=None, - multimodal_outputs={ - key: torch.rand((orig_num_tokens_unpadded, feature_dims[key]), dtype=DTYPE) for key in cache.mm_cache_keys - }, - num_tokens_unpadded=orig_num_tokens_unpadded, - slot_mapping=torch.arange(8, 8 + orig_num_tokens_unpadded), - num_tokens_padded=None, - ) - cache.add_prefix_cached_new_req_id("req1") - - total_scheduled_tokens = num_new_toks_req1 + num_new_toks_req2 - mocker.patch( - "vllm_omni.core.prefix_cache.OmniTensorPrefixCache._get_cached_block_ids", - new=fake_get_cached_block_ids, - ) - return MergeHarness( - cache=cache, - input_batch=MockInputBatch(num_computed_tokens_cpu=torch.Tensor([orig_num_tokens_unpadded, 0])), - num_scheduled_tokens={"req1": num_new_toks_req1, "req2": num_new_toks_req2}, - query_start_loc=[0, num_new_toks_req1], - total_scheduled_tokens=total_scheduled_tokens, - new_mm_outputs={ - key: torch.rand((total_scheduled_tokens, feature_dims[key]), dtype=DTYPE) for key in cache.mm_cache_keys - }, - ) - - -def test_get_merged_multimodal_outputs_passes_per_request_tensor_list_through(mocker): - """The passthrough shape qwen3-tts actually emits: a nested ``codes``/``meta`` - dict whose ``codes.ref`` holds one tensor per request - (qwen3_tts_talker.make_omni_output), carried through to_payload_element's - pass_lists_through branch beside 1D int metadata that IS sliced per request. - - Each request must receive the list *whole*. Per-request indexing happens - later, in the runner's _unwrap_lists; a W2 that split the list here would - make that second index select a different request's reference codes. The - existing list coverage passes a list of strings, which cannot catch a - tensor-specific regression, and no other test covers an int dtype. - """ - harness = build_mixed_hit_miss_harness(mocker) - - ref_codes = [torch.tensor([11, 12], dtype=torch.long), torch.tensor([21], dtype=torch.long)] - ref_code_len = torch.arange(harness.total_scheduled_tokens, dtype=torch.int32) - harness.new_mm_outputs["codes"] = {"ref": ref_codes} - harness.new_mm_outputs["meta"] = {"ref_code_len": ref_code_len} - - merged = harness.cache.get_merged_multimodal_states( - query_start_loc=harness.query_start_loc, - input_batch=harness.input_batch, - multimodal_outputs=harness.new_mm_outputs, - num_scheduled_tokens=harness.num_scheduled_tokens, - ) - - assert "codes" not in harness.cache.mm_cache_keys - for req_id in ("req1", "req2"): - req_refs = merged["codes"][req_id]["ref"] - assert len(req_refs) == len(ref_codes) - for got, want in zip(req_refs, ref_codes): - assert torch.equal(got, want) - assert got.dtype == torch.long - - # Per-request clones: the branch exists so one request's downstream mutation - # cannot reach another request's payload or the source list. - merged["codes"]["req1"]["ref"][0] += 100 - assert torch.equal(merged["codes"]["req2"]["ref"][0], torch.tensor([11, 12], dtype=torch.long)) - assert torch.equal(ref_codes[0], torch.tensor([11, 12], dtype=torch.long)) - - # Token-aligned 1D metadata is still split per request, unlike the list. - req1_len = harness.num_scheduled_tokens["req1"] - assert torch.equal(merged["meta"]["req1"]["ref_code_len"], ref_code_len[:req1_len]) - assert torch.equal(merged["meta"]["req2"]["ref_code_len"], ref_code_len[req1_len:]) - assert merged["meta"]["req1"]["ref_code_len"].dtype == torch.int32 - - -@pytest.mark.parametrize( - "passthrough_shape", - [(3, 4), (1, 4), (7,)], - ids=["shorter_2d", "single_row_2d", "longer_1d"], -) -def test_get_merged_multimodal_outputs_broadcasts_non_token_aligned_passthrough(passthrough_shape, mocker): - """A passthrough tensor whose first dimension is not the scheduled-token - count is request-invariant: to_payload_element hands every request the whole - tensor, cloned, rather than slicing it. - - W2 batching most plausibly breaks by assuming every passthrough tensor is - token-aligned; that assumption would slice these and give each request a - truncated view, or an empty one once start exceeds the first dimension. - """ - harness = build_mixed_hit_miss_harness(mocker) - assert passthrough_shape[0] != harness.total_scheduled_tokens - - passthrough = torch.arange(math.prod(passthrough_shape), dtype=DTYPE).reshape(passthrough_shape) - expected = passthrough.clone() - harness.new_mm_outputs["speaker_embedding"] = passthrough - - merged = harness.cache.get_merged_multimodal_states( - query_start_loc=harness.query_start_loc, - input_batch=harness.input_batch, - multimodal_outputs=harness.new_mm_outputs, - num_scheduled_tokens=harness.num_scheduled_tokens, - ) - - assert "speaker_embedding" not in harness.cache.mm_cache_keys - out = merged["speaker_embedding"] - for req_id in ("req1", "req2"): - assert out[req_id].shape == passthrough_shape - assert torch.equal(out[req_id], expected) - - out["req1"] += 100 - assert torch.equal(out["req2"], expected) - assert torch.equal(passthrough, expected) + assert cfg_fallback.staging_capacity_tokens == 4096 diff --git a/tests/core/test_prefix_cache_async_write.py b/tests/core/test_prefix_cache_async_write.py deleted file mode 100644 index fa1980227..000000000 --- a/tests/core/test_prefix_cache_async_write.py +++ /dev/null @@ -1,119 +0,0 @@ -import pytest -import torch - -from vllm_omni.core.prefix_cache import OmniTensorPrefixCache - -# `cpu` is the lane selector, not a hardware claim: the only Buildkite steps that -# reach tests/ select `core_model and cpu`, and they run on GPU hardware (l4_1). -# Marking this `cuda` collects it nowhere. Real hardware need is the skipif below. -pytestmark = [ - pytest.mark.core_model, - pytest.mark.cpu, - pytest.mark.skipif(not torch.cuda.is_available(), reason="requires a CUDA device"), -] - -NUM_BLOCKS = 10 -BLOCK_SIZE = 4 -HIDDEN_SIZE = 2 -DTYPE = torch.float32 - - -def get_omni_pcache() -> OmniTensorPrefixCache: - """Build an OmniTensorPrefixCache, but don't init mm tensors.""" - return OmniTensorPrefixCache( - num_blocks=NUM_BLOCKS, - block_size=BLOCK_SIZE, - hidden_size=HIDDEN_SIZE, - hs_dtype=DTYPE, - ) - - -def test_drain_ready_async_writes_is_a_noop_before_the_event_is_ready(monkeypatch): - """drain_ready_async_writes() must not touch the cache while the - scheduled D2H's event has not completed, and must consume it exactly - once after it has. - - Real GPU timing (e.g. via torch.cuda._sleep to stall the copy stream) - was tried and found unreliable on this box: cross-stream wait_stream/ - wait_event did not reliably delay event.query() from reporting ready - immediately, likely a virtualization/driver quirk of this cloud GPU. - Forcing the event's query() result directly tests the branching logic - in drain_ready_async_writes/_consume_pending_write deterministically, - independent of real CUDA scheduling behavior. - """ - cache = get_omni_pcache() - mm_key = "codes.audio" - cache.maybe_init_missing_mm_cache_keys({mm_key: torch.zeros(4, HIDDEN_SIZE)}, seq_len=4) - - num_tokens = 4 - slot_mapping_gpu = torch.arange(8, 8 + num_tokens, dtype=torch.int64, device="cuda") - mm_gpu = { - mm_key: torch.arange(num_tokens * HIDDEN_SIZE, dtype=torch.float32, device="cuda").reshape( - num_tokens, HIDDEN_SIZE - ) - } - - cache.schedule_async_write( - hidden_states_gpu=None, - multimodal_outputs_gpu=mm_gpu, - slot_mapping_gpu=slot_mapping_gpu, - num_tokens_unpadded=num_tokens, - num_tokens_padded=num_tokens, - ) - - pending_event = cache._pending_write.event - rows = cache.mm_outputs_cache[mm_key].view(-1, HIDDEN_SIZE) - - monkeypatch.setattr(pending_event, "query", lambda: False) - assert cache.drain_ready_async_writes() == 0 - assert torch.all(rows[8 : 8 + num_tokens] == 0) - assert cache._pending_write is not None - - monkeypatch.setattr(pending_event, "query", lambda: True) - assert cache.drain_ready_async_writes() == 1 - assert torch.equal(rows[8 : 8 + num_tokens], mm_gpu[mm_key].cpu()) - assert cache._pending_write is None - - # Nothing pending anymore: a second drain is a no-op, not a re-scatter. - assert cache.drain_ready_async_writes() == 0 - - -def test_schedule_async_write_consumes_previous_pending_write_first(): - """Scheduling step N+1's write must first consume (scatter) step N's - pending write, even if the caller never explicitly called - drain_ready_async_writes() in between -- otherwise step N's data would - be silently dropped when overwritten by step N+1's pending write.""" - cache = get_omni_pcache() - mm_key = "codes.audio" - cache.maybe_init_missing_mm_cache_keys({mm_key: torch.zeros(4, HIDDEN_SIZE)}, seq_len=4) - - num_tokens = 4 - first_slots = torch.arange(8, 8 + num_tokens, dtype=torch.int64, device="cuda") - first_gpu = {mm_key: torch.full((num_tokens, HIDDEN_SIZE), 1.0, device="cuda")} - second_slots = torch.arange(12, 12 + num_tokens, dtype=torch.int64, device="cuda") - second_gpu = {mm_key: torch.full((num_tokens, HIDDEN_SIZE), 2.0, device="cuda")} - - cache.schedule_async_write( - hidden_states_gpu=None, - multimodal_outputs_gpu=first_gpu, - slot_mapping_gpu=first_slots, - num_tokens_unpadded=num_tokens, - num_tokens_padded=num_tokens, - ) - torch.accelerator.synchronize() - - # This call's internal _consume_pending_write() must flush the FIRST - # write before the second one is scheduled; the caller never drained. - cache.schedule_async_write( - hidden_states_gpu=None, - multimodal_outputs_gpu=second_gpu, - slot_mapping_gpu=second_slots, - num_tokens_unpadded=num_tokens, - num_tokens_padded=num_tokens, - ) - torch.accelerator.synchronize() - cache.drain_ready_async_writes() - - rows = cache.mm_outputs_cache[mm_key].view(-1, HIDDEN_SIZE) - assert torch.all(rows[8:12] == 1.0) - assert torch.all(rows[12:16] == 2.0) diff --git a/tests/core/test_prefix_cache_runner_mixin.py b/tests/core/test_prefix_cache_runner_mixin.py new file mode 100644 index 000000000..3809c20b8 --- /dev/null +++ b/tests/core/test_prefix_cache_runner_mixin.py @@ -0,0 +1,182 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""CPU tests for PrefixCacheRunnerMixin: one-shot construction, policy +snapshot identity/isolation, and the shared save/materialize gates. +No vLLM and no GPU required (same shim as test_prefix_cache.py).""" + +import sys +from pathlib import Path +from types import ModuleType, SimpleNamespace +from typing import Any + +import pytest +import torch + +try: # pragma: no cover - shim only matters on vllm-less dev machines + import vllm # noqa: F401 +except ModuleNotFoundError: + _root = Path(__file__).resolve().parents[2] + for _pkg in ("vllm_omni", "vllm_omni.core"): + if _pkg not in sys.modules: + _m = ModuleType(_pkg) + _m.__path__ = [str(_root / _pkg.replace(".", "/"))] + sys.modules[_pkg] = _m + import logging + + _vllm = ModuleType("vllm") + _vllm_logger = ModuleType("vllm.logger") + setattr(_vllm_logger, "init_logger", logging.getLogger) + setattr(_vllm, "logger", _vllm_logger) + sys.modules["vllm"] = _vllm + sys.modules["vllm.logger"] = _vllm_logger + +from vllm_omni.core.prefix_cache.interface import PrefixCacheConfig, StageCacheOutputs +from vllm_omni.core.prefix_cache.runner_mixin import PrefixCacheRunnerMixin + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +BLOCK_SIZE = 4 + + +def _patch_pp(monkeypatch, is_last: bool) -> None: + """Route the mixin's call-time get_pp_group import to a fake rank.""" + mod = ModuleType("vllm.distributed.parallel_state") + setattr(mod, "get_pp_group", lambda: SimpleNamespace(is_last_rank=is_last)) + if "vllm.distributed" not in sys.modules: + monkeypatch.setitem(sys.modules, "vllm.distributed", ModuleType("vllm.distributed")) + monkeypatch.setitem(sys.modules, "vllm.distributed.parallel_state", mod) + + +def _full_attention_groups(): + """A single-group config that passes check_prefix_cache_kv_groups. + + Real FullAttentionSpec when vllm is installed; otherwise a stand-in + registered under the import path the check resolves at call time. + """ + try: + from vllm.v1.kv_cache_interface import FullAttentionSpec + except ModuleNotFoundError: + iface = sys.modules.get("vllm.v1.kv_cache_interface") + if iface is None: + iface = ModuleType("vllm.v1.kv_cache_interface") + + class _FullAttentionSpecStub: + pass + + setattr(iface, "FullAttentionSpec", _FullAttentionSpecStub) + sys.modules["vllm.v1"] = ModuleType("vllm.v1") + sys.modules["vllm.v1.kv_cache_interface"] = iface + FullAttentionSpec = getattr(iface, "FullAttentionSpec") + return [SimpleNamespace(kv_cache_spec=FullAttentionSpec.__new__(FullAttentionSpec))] + + +class _FakeBlockTables: + """Group-0 table whose layout attrs pass check_prefix_cache_block_layout.""" + + def __init__(self): + self.block_tables = [object()] + + def __getitem__(self, idx): + return SimpleNamespace() + + +class _Runner(PrefixCacheRunnerMixin): + def __init__(self, cfg: PrefixCacheConfig | None = None): + # Any: tests swap in _CacheStub for the manager. + self.omni_prefix_cache: Any = None + self._omni_prefix_cache_cfg = cfg + self.input_batch = SimpleNamespace(block_table=_FakeBlockTables(), req_ids=[]) + self.kv_cache_config = SimpleNamespace(kv_cache_groups=_full_attention_groups()) + self.is_pooling_model = False + + +def _cfg() -> PrefixCacheConfig: + return PrefixCacheConfig(num_blocks=8, block_size=BLOCK_SIZE) + + +def _sched(): + return SimpleNamespace(scheduled_new_reqs=(), finished_req_ids=set(), num_scheduled_tokens={}) + + +class _CacheStub: + def __init__(self): + self.save_calls = [] + self.materialize_calls = [] + self.outs = StageCacheOutputs(hidden_states=None, mm_outputs={}) + + def save_outputs(self, hidden, mm, *, num_tokens_unpadded, num_tokens_padded): + self.save_calls.append((hidden, mm, num_tokens_unpadded, num_tokens_padded)) + return 7 + + def materialize(self, step_id, req_ids): + self.materialize_calls.append((step_id, req_ids)) + return self.outs + + +def test_step_begin_builds_once_and_registers_snapshot_policy(monkeypatch): + _patch_pp(monkeypatch, is_last=True) + r = _Runner(_cfg()) + r._snapshot_prefix_cache_model_policy( + SimpleNamespace(requires_full_prefix_cached_hidden_states=True, deferred_prefix_cache_mm_keys=("k",)) + ) + r._prefix_cache_step_begin(_sched()) + mgr = r.omni_prefix_cache + assert mgr is not None and r._omni_prefix_cache_cfg is None + assert mgr._policy is r._omni_cache_policy # the load-time snapshot, same object + r._prefix_cache_step_begin(_sched()) + assert r.omni_prefix_cache is mgr # one-shot + + +def test_default_policy_isolated_per_instance(): + a, b = _Runner(), _Runner() + a._snapshot_prefix_cache_model_policy(SimpleNamespace(requires_full_prefix_cached_hidden_states=False)) + assert a._omni_cache_policy.needs_full_hidden_states is False + assert a._model_needs_full_prefix_hidden_states() is False + assert b._omni_cache_policy.needs_full_hidden_states is True # class default untouched + assert PrefixCacheRunnerMixin._omni_cache_policy.needs_full_hidden_states is True + + +def test_non_last_rank_consumes_cfg_without_building(monkeypatch): + _patch_pp(monkeypatch, is_last=False) + r = _Runner(_cfg()) + r._prefix_cache_step_begin(_sched()) + assert r.omni_prefix_cache is None and r._omni_prefix_cache_cfg is None + + +def test_disabled_cache_step_begin_is_a_noop(): + r = _Runner(cfg=None) # enable_prefix_caching off: nothing staged + r._prefix_cache_step_begin(_sched()) + assert r.omni_prefix_cache is None + + +def test_save_step_gates_and_passthrough(monkeypatch): + r = _Runner() + stub = _CacheStub() + hidden = torch.zeros(2, 2) + + def save(): + return r._prefix_cache_save_step(hidden, None, num_tokens_unpadded=2, num_tokens_padded=2) + + _patch_pp(monkeypatch, is_last=True) + assert save() is None # cache off + r.omni_prefix_cache = stub + r.is_pooling_model = True + assert save() is None # pooling stage never writes + r.is_pooling_model = False + _patch_pp(monkeypatch, is_last=False) + assert save() is None # not the last PP rank + _patch_pp(monkeypatch, is_last=True) + assert save() == 7 + assert stub.save_calls == [(hidden, {}, 2, 2)] # empty mm stays {} + + +def test_materialize_requires_explicit_step_and_snapshot_req_ids(): + r = _Runner() + stub = _CacheStub() + r.omni_prefix_cache = stub + assert r._prefix_cache_materialize(None, ["a"]) == (None, None) + assert stub.materialize_calls == [] # step_id None: nothing to consume + hidden = {"a": torch.zeros(1, 2)} + stub.outs = StageCacheOutputs(hidden_states=hidden, mm_outputs={}) + assert r._prefix_cache_materialize(7, ["a"]) == (hidden, None) # empty mm -> None + assert stub.materialize_calls == [(7, ["a"])] diff --git a/tests/e2e/online_serving/test_qwen3_omni.py b/tests/e2e/online_serving/test_qwen3_omni.py index f6ede8944..5b4bbc51e 100644 --- a/tests/e2e/online_serving/test_qwen3_omni.py +++ b/tests/e2e/online_serving/test_qwen3_omni.py @@ -24,29 +24,40 @@ _MODEL = "Qwen/Qwen3-Omni-30B-A3B-Instruct" _CI_DEPLOY = get_deploy_config_path("ci/qwen3_omni_moe.yaml") -# For prefix caching checks against we enable it on the thinker and talker via CLI override -# and enable prompt token details so that we can determine if any tokens were cached. -# We also explicitly set block size so that we can make sure the cached token counts are a -# multiple of the block size. +# Prefix cache is opt-in on the deploy YAML. CI overlay also pins +# ``async_chunk: False``; the existing server keeps ``--no-async-chunk``. +# Production Qwen YAML is ``async_chunk: true``, so a sibling server boots +# with ``--async-chunk`` and the same cache overrides (async output builder). BLOCK_SIZE = 16 +_PREFIX_CACHE_SERVER_ARGS = [ + "--block-size", + str(BLOCK_SIZE), + "--stage-overrides", + '{"0": {"enable_prefix_caching": true}, "1": {"enable_prefix_caching": true}}', + "--enable-prompt-tokens-details", +] test_params = [ pytest.param( OmniServerParams( model=_MODEL, stage_config_path=_CI_DEPLOY, use_stage_cli=True, - server_args=[ - "--no-async-chunk", - "--block-size", - str(BLOCK_SIZE), - "--stage-overrides", - '{"0": {"enable_prefix_caching": true}, "1": {"enable_prefix_caching": true}}', - "--enable-prompt-tokens-details", - ], + server_args=["--no-async-chunk", *_PREFIX_CACHE_SERVER_ARGS], ), id="default", ) ] +prefix_cache_async_chunk_params = [ + pytest.param( + OmniServerParams( + model=_MODEL, + stage_config_path=_CI_DEPLOY, + use_stage_cli=True, + server_args=["--async-chunk", *_PREFIX_CACHE_SERVER_ARGS], + ), + id="prefix_cache_async_chunk", + ) +] def get_system_prompt(): @@ -245,6 +256,83 @@ def test_thinker_prefix_caching_audio_output(omni_server, online_client) -> None _run_prefix_cache_check(online_client, request_config) +def _assert_omni_payload_complete(resp) -> None: + """Async-chunk + prefix-cache must still deliver a full downstream payload.""" + assert resp.success + assert resp.text_content + has_audio = bool(resp.audio_bytes) or bool(resp.audio_data) + assert has_audio + + +@pytest.mark.advanced_model +@pytest.mark.core_model +@pytest.mark.omni +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) +@pytest.mark.parametrize("omni_server", prefix_cache_async_chunk_params, indirect=True) +def test_thinker_prefix_caching_async_chunk_identical(omni_server, online_client) -> None: + """Production path: ``async_chunk: true`` plus thinker/talker prefix cache. + + Two identical streaming requests. The second must report a block-aligned + cache hit and still produce text + audio (async output builder). + """ + messages = dummy_messages_from_mix_data( + system_prompt=get_system_prompt(), + content_text=get_prompt(), + ) + request_config = { + "model": omni_server.model, + "messages": messages, + "stream": True, + "stream_options": {"include_usage": True}, + } + first_response, cached_response = _run_prefix_cache_check(online_client, request_config) + _assert_omni_payload_complete(first_response) + _assert_omni_payload_complete(cached_response) + + +@pytest.mark.advanced_model +@pytest.mark.core_model +@pytest.mark.omni +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) +@pytest.mark.parametrize("omni_server", prefix_cache_async_chunk_params, indirect=True) +def test_thinker_prefix_caching_async_chunk_shared_image_prefix(omni_server, online_client) -> None: + """Shared image / system prefix, different user text — not a full-prompt replay. + + The second request must still hit cached tokens (block-aligned, less than + its prompt). A merge that only works when ``hit_upto == prompt_len`` would + fail this. + """ + image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(224, 224)['base64']}" + first_messages = dummy_messages_from_mix_data( + system_prompt=get_system_prompt(), + image_data_url=image_data_url, + content_text=get_prompt("text_image"), + ) + second_messages = dummy_messages_from_mix_data( + system_prompt=get_system_prompt(), + image_data_url=image_data_url, + content_text="How many squares are in this image? Answer in 20 words.", + ) + first_cfg = { + "model": omni_server.model, + "messages": first_messages, + "stream": True, + "stream_options": {"include_usage": True}, + "modalities": ["text"], + } + second_cfg = {**first_cfg, "messages": second_messages} + first_response = online_client.send_omni_request(first_cfg, request_num=1)[0] + cached_response = online_client.send_omni_request(second_cfg, request_num=1)[0] + assert first_response.success and cached_response.success + assert first_response.text_content and cached_response.text_content + num_cached = cached_response.cached_tokens + num_prompt = cached_response.prompt_tokens + assert num_cached is not None and num_prompt is not None + assert num_cached > 0 + assert num_cached % BLOCK_SIZE == 0 + assert num_cached < num_prompt + + @pytest.mark.advanced_model @pytest.mark.core_model @pytest.mark.omni diff --git a/tests/model_executor/models/voxcpm2/test_runner_assisted_unified_graph_static.py b/tests/model_executor/models/voxcpm2/test_runner_assisted_unified_graph_static.py index 0007f021a..0c2d9ee67 100644 --- a/tests/model_executor/models/voxcpm2/test_runner_assisted_unified_graph_static.py +++ b/tests/model_executor/models/voxcpm2/test_runner_assisted_unified_graph_static.py @@ -74,7 +74,9 @@ def test_ar_runner_without_model_hook_stays_on_normal_path(): assert "exceptException" not in compact_request_source set_context_source = source[source.index("def _set_runner_assisted_full_attention_metadata_context") :] - set_context_source = set_context_source[: set_context_source.index("def _deferred_prefix_cache_mm_keys")] + set_context_source = set_context_source[ + : set_context_source.index("def _prepare_prefix_cache_pooler_payload_sources") + ] assert "except Exception" not in set_context_source diff --git a/tests/platforms/npu/test_npu_prefix_cache_wiring.py b/tests/platforms/npu/test_npu_prefix_cache_wiring.py new file mode 100644 index 000000000..12cded5ba --- /dev/null +++ b/tests/platforms/npu/test_npu_prefix_cache_wiring.py @@ -0,0 +1,158 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""NPU prefix-cache wiring: save_outputs -> sid -> materialize on CPU. + +The NPU runner consumes the prefix cache in eager mode (no CUDA streams) and +hands the step id across execute_model()/sample_tokens() as the LAST field of +its positionally-packed ``ExecuteModelState``. Neither path runs in CUDA CI, +so a hit-read change can silently break them (`vllm_ascend` imports keep the +runner modules out of reach of a plain unit test). + +Two layers of protection, both pure CPU: + +1. AST characterization of the ``ExecuteModelState`` field order in the NPU + and GPU runners — the pack/unpack sites pass fields positionally, so a + reorder or drop of ``prefix_cache_step_id`` corrupts the sid without any + import-time error. +2. The eager-mode facade contract the NPU runner relies on: controller + auto-selects eager when CUDA is unavailable, and save -> materialize + returns byte-identical rows for both the miss and the hit-merge paths. +""" + +import ast +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch + +from vllm_omni.core.prefix_cache.interface import PrefixCacheConfig +from vllm_omni.core.prefix_cache.manager import OmniPrefixCacheManager + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +_REPO_ROOT = Path(__file__).resolve().parents[3] +_NPU_RUNNER = _REPO_ROOT / "vllm_omni/platforms/npu/worker/npu_ar_model_runner.py" +_GPU_RUNNER = _REPO_ROOT / "vllm_omni/worker/gpu_ar_model_runner.py" + +NUM_BLOCKS = 16 +BLOCK_SIZE = 4 +HIDDEN = 8 +DTYPE = torch.float32 + + +def _named_tuple_fields(path: Path, class_name: str) -> list[str]: + tree = ast.parse(path.read_text()) + for node in ast.walk(tree): + if isinstance(node, ast.ClassDef) and node.name == class_name: + return [stmt.target.id for stmt in node.body if isinstance(stmt, ast.AnnAssign)] + raise AssertionError(f"{class_name} not found in {path}") + + +def test_execute_model_state_keeps_prefix_cache_sid_last(): + """Both runners pack ExecuteModelState positionally: the sid must stay + the LAST field, and the NPU field order is pinned exactly — a reorder + passes import and every CUDA test, then ships the wrong value as the + sid on hardware.""" + npu_fields = _named_tuple_fields(_NPU_RUNNER, "ExecuteModelState") + assert npu_fields == [ + "scheduler_output", + "logits", + "spec_decode_metadata", + "spec_decode_common_attn_metadata", + "hidden_states", + "sample_hidden_states", + "aux_hidden_states", + "attn_metadata", + "positions", + "ec_connector_output", + "cudagraph_stats", + "batch_desc", + "multimodal_outputs", + "prefix_cache_step_id", + ] + gpu_fields = _named_tuple_fields(_GPU_RUNNER, "ExecuteModelState") + assert gpu_fields[-1] == "prefix_cache_step_id" + + +class _FakeView: + """Minimal group-view double (mirrors tests/core/test_prefix_cache).""" + + def __init__(self): + self.block_size = BLOCK_SIZE + self.num_blocks = NUM_BLOCKS + self.req_blocks: dict[str, list[int]] = {} + self.order: list[str] = [] + self.computed: dict[str, int] = {} + + def slots_for(self, req_id, token_start, token_end): + blocks = self.req_blocks[req_id] + return torch.tensor( + [blocks[p // BLOCK_SIZE] * BLOCK_SIZE + p % BLOCK_SIZE for p in range(token_start, token_end)], + dtype=torch.long, + ) + + def batch_req_ids(self) -> list[str]: + return list(self.order) + + def step_slots_cpu(self, req_ids, num_scheduled) -> torch.Tensor: + parts = [] + for r in req_ids: + n = int(num_scheduled.get(r, 0)) + if n > 0: + start = self.computed.get(r, 0) + parts.append(self.slots_for(r, start, start + n)) + return torch.cat(parts) if parts else torch.empty((0,), dtype=torch.long) + + +def _run_step(mgr, view, req_id, blocks, start_pos, sched, *, hit=0, finished=()): + """One NPU-shaped step: new_step_starts -> save_outputs -> sid.""" + view.order = [req_id] + view.req_blocks[req_id] = blocks + view.computed[req_id] = start_pos + slots = view.slots_for(req_id, start_pos, start_pos + sched) + hidden = slots.to(DTYPE).unsqueeze(1).expand(sched, HIDDEN).clone() + sched_out = SimpleNamespace( + scheduled_new_reqs=[ + SimpleNamespace(req_id=req_id, num_computed_tokens=hit, block_ids=[list(blocks)]), + ], + finished_req_ids=set(finished), + num_scheduled_tokens={req_id: sched}, + ) + mgr.new_step_starts(sched_out) + return mgr.save_outputs(hidden, {}, num_tokens_unpadded=sched, num_tokens_padded=sched) + + +def _expected_rows(slots: torch.Tensor) -> torch.Tensor: + return slots.to(DTYPE).unsqueeze(1).expand(slots.numel(), HIDDEN) + + +def _make_npu_mode_manager(monkeypatch) -> tuple[OmniPrefixCacheManager, _FakeView]: + # The NPU condition: no CUDA -> the controller must auto-select eager + # (submit() completes copy+scatter synchronously), exactly how + # npu_model_runner builds the manager with the default eager=None. + monkeypatch.setattr(torch.cuda, "is_available", lambda: False) + view = _FakeView() + config = PrefixCacheConfig(num_blocks=NUM_BLOCKS, block_size=BLOCK_SIZE) + return OmniPrefixCacheManager(config, view), view + + +def test_npu_mode_auto_selects_eager_and_roundtrips(monkeypatch): + mgr, view = _make_npu_mode_manager(monkeypatch) + assert mgr._controller._eager is True + sid = _run_step(mgr, view, "a", [0, 1], 0, 8) + outs = mgr.materialize(sid, ["a"]) + assert torch.equal(outs.hidden_states["a"], _expected_rows(view.slots_for("a", 0, 8))) + + +def test_npu_mode_hit_merges_cached_prefix(monkeypatch): + """The hit-read path in eager mode: consumer's merged output must be + [cached prefix rows + this step's rows], byte-identical.""" + mgr, view = _make_npu_mode_manager(monkeypatch) + s1 = _run_step(mgr, view, "a", [0, 1], 0, 8) + mgr.materialize(s1, ["a"]) + s2 = _run_step(mgr, view, "b", [0, 1, 2], 8, 4, hit=8, finished=["a"]) + merged = mgr.materialize(s2, ["b"]).hidden_states["b"] + assert merged.shape == (12, HIDDEN) + assert torch.equal(merged[:8], _expected_rows(view.slots_for("b", 0, 8))) + assert torch.equal(merged[8:], _expected_rows(view.slots_for("b", 8, 12))) diff --git a/tests/worker/test_fa3_scheduler_metadata_size.py b/tests/worker/test_fa3_scheduler_metadata_size.py index 5e308bed7..ac7d1a752 100644 --- a/tests/worker/test_fa3_scheduler_metadata_size.py +++ b/tests/worker/test_fa3_scheduler_metadata_size.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Characterization of ``OmniGPUModelRunner.initialize_metadata_builders``. Upstream ``FlashAttentionMetadataBuilder`` pre-allocates ``scheduler_metadata`` @@ -19,14 +19,12 @@ from types import SimpleNamespace import pytest import torch +from vllm.v1.worker.gpu_model_runner import GPUModelRunner from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner pytestmark = [pytest.mark.core_model, pytest.mark.cpu] -# The class ``super().initialize_metadata_builders`` dispatches to. -_UPSTREAM_RUNNER = OmniGPUModelRunner.__mro__[1] - def _builder(*, prealloc: int | None, max_num_splits: int) -> SimpleNamespace: sm = None if prealloc is None else torch.zeros(prealloc, dtype=torch.int32) @@ -36,10 +34,14 @@ def _builder(*, prealloc: int | None, max_num_splits: int) -> SimpleNamespace: def _make_runner(monkeypatch, *, max_num_seqs: int, builders: list) -> OmniGPUModelRunner: # Stub the upstream initialize_metadata_builders; our fake runner already # carries the attn_groups the Omni override iterates over. - monkeypatch.setattr(_UPSTREAM_RUNNER, "initialize_metadata_builders", lambda self, *a, **k: None) + monkeypatch.setattr(GPUModelRunner, "initialize_metadata_builders", lambda self, *a, **k: None) runner = object.__new__(OmniGPUModelRunner) runner.scheduler_config = SimpleNamespace(max_num_seqs=max_num_seqs) runner.cache_config = SimpleNamespace(enable_prefix_caching=False) # skip omni_prefix_cache branch + runner.vllm_config = SimpleNamespace(kv_transfer_config=None) + runner.model_config = SimpleNamespace() + runner.is_pooling_model = False + runner.speculative_config = None runner.attn_groups = [[SimpleNamespace(metadata_builders=builders)]] return runner diff --git a/tests/worker/test_gpu_ar_model_runner.py b/tests/worker/test_gpu_ar_model_runner.py index 33abbb618..d09282035 100644 --- a/tests/worker/test_gpu_ar_model_runner.py +++ b/tests/worker/test_gpu_ar_model_runner.py @@ -18,6 +18,7 @@ from vllm.sampling_params import SamplingParams from vllm.v1.worker import gpu_input_batch from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch +from vllm_omni.core.prefix_cache import ModelCachePolicy from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech from vllm_omni.entrypoints.openai.tts_adapters.base import PreparedRequest @@ -37,6 +38,8 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu] def _make_runner(engine_output_type: str | None, downstream_req_ids: set[str]) -> GPUARModelRunner: runner = object.__new__(GPUARModelRunner) + runner._omni_cache_policy = ModelCachePolicy(needs_full_hidden_states=True) + runner._pooler_payload_include_hidden_flag = True runner.vllm_config = SimpleNamespace( model_config=SimpleNamespace(engine_output_type=engine_output_type), ) @@ -57,6 +60,8 @@ def _mtp_runner(*, async_scheduling: bool, buffers: dict[str, dict]) -> tuple[GP return torch.tensor([[11, 12, 13]], dtype=torch.long) runner = object.__new__(GPUARModelRunner) + runner._omni_cache_policy = ModelCachePolicy(needs_full_hidden_states=True) + runner._pooler_payload_include_hidden_flag = True runner.use_async_scheduling = async_scheduling runner.model = SimpleNamespace(post_sample_talker_mtp=post_sample_talker_mtp) runner.requests = {req_id: SimpleNamespace() for req_id in buffers} @@ -143,6 +148,8 @@ def test_post_sample_talker_mtp_uses_gpu_token_with_async_scheduling() -> None: def test_post_sample_talker_mtp_skips_shape_validation_without_duplex_rows() -> None: runner = object.__new__(GPUARModelRunner) + runner._omni_cache_policy = ModelCachePolicy(needs_full_hidden_states=True) + runner._pooler_payload_include_hidden_flag = True runner.use_async_scheduling = False runner.model = SimpleNamespace( post_sample_talker_mtp=lambda **_: pytest.fail("hook must not run"), @@ -178,6 +185,8 @@ def test_post_sample_talker_mtp_rejects_invalid_selected_token_shape( error_match: str, ) -> None: runner = object.__new__(GPUARModelRunner) + runner._omni_cache_policy = ModelCachePolicy(needs_full_hidden_states=True) + runner._pooler_payload_include_hidden_flag = True runner.use_async_scheduling = False runner.model = SimpleNamespace( post_sample_talker_mtp=lambda **_: pytest.fail("hook must not run"), @@ -265,6 +274,8 @@ def test_speech_extra_params_reach_model_sampler_as_sampling_metadata(monkeypatc received = [] runner = object.__new__(GPUARModelRunner) + runner._omni_cache_policy = ModelCachePolicy(needs_full_hidden_states=True) + runner._pooler_payload_include_hidden_flag = True runner.input_batch = input_batch runner.model = SimpleNamespace( prefer_model_sampler=True, @@ -323,6 +334,8 @@ def test_sparse_mm_req_ids_requires_sparse_audio_marker(): def test_runner_assisted_full_attention_metadata_request_is_opt_in(): runner = object.__new__(GPUARModelRunner) + runner._omni_cache_policy = ModelCachePolicy(needs_full_hidden_states=True) + runner._pooler_payload_include_hidden_flag = True runner.model = object() runner.scheduler_config = SimpleNamespace(max_num_seqs=16) @@ -377,6 +390,8 @@ def test_runner_assisted_full_attention_metadata_request_and_context_hooks(): calls.append(("context", {"enabled": enabled, "num_reqs": num_reqs})) runner = object.__new__(GPUARModelRunner) + runner._omni_cache_policy = ModelCachePolicy(needs_full_hidden_states=True) + runner._pooler_payload_include_hidden_flag = True runner.model = Model() runner.scheduler_config = SimpleNamespace(max_num_seqs=8) @@ -466,6 +481,8 @@ def test_omni_async_gpu_model_runner_output_reraises_background_exception(): def _make_async_output_runner(engine_output_type: str = "audio"): runner = object.__new__(GPUARModelRunner) + runner._omni_cache_policy = ModelCachePolicy(needs_full_hidden_states=True) + runner._pooler_payload_include_hidden_flag = True model_config = SimpleNamespace( engine_output_type=engine_output_type, async_chunk=True, @@ -642,10 +659,11 @@ def test_async_omni_output_guard_requires_safe_conditions(): assert GPUARModelRunner._should_use_async_omni_output(runner) + # Prefix cache no longer forces the sync output path; leftover mm is + # snapshotted at save so the builder can materialize a step late. runner.omni_prefix_cache = object() - assert not GPUARModelRunner._should_use_async_omni_output(runner) + assert GPUARModelRunner._should_use_async_omni_output(runner) - runner.omni_prefix_cache = None runner.model.has_postprocess = True assert not GPUARModelRunner._should_use_async_omni_output(runner) @@ -656,6 +674,7 @@ def test_async_omni_output_guard_requires_safe_conditions(): def test_build_omni_output_skips_hidden_when_model_opts_out(monkeypatch): runner = _make_async_output_runner(engine_output_type="latent") runner.model.omni_pooler_payload_include_hidden = False + runner._pooler_payload_include_hidden_flag = False monkeypatch.setattr( GPUARModelRunner, @@ -699,6 +718,7 @@ def test_build_omni_output_skips_hidden_when_model_opts_out(monkeypatch): def test_build_omni_output_splits_mm_by_scheduled_tokens_when_hidden_is_tail_only(monkeypatch): runner = _make_async_output_runner(engine_output_type="latent") runner.model.omni_pooler_payload_include_hidden = False + runner._pooler_payload_include_hidden_flag = False runner._async_chunk = False runner.requests = {"r1": object(), "r2": object(), "r3": object()} @@ -745,6 +765,7 @@ def test_build_omni_output_splits_mm_by_hidden_len_when_scheduled_is_padded(monk """Thinker mm rows align to hidden_states.shape[0], not padded scheduled count.""" runner = _make_async_output_runner(engine_output_type="latent") runner.model.omni_pooler_payload_include_hidden = False + runner._pooler_payload_include_hidden_flag = False runner._async_chunk = False runner.requests = {"r1": object(), "r2": object(), "r3": object()} @@ -790,6 +811,7 @@ def test_build_omni_output_splits_mm_by_hidden_len_when_scheduled_is_padded(monk def test_async_snapshot_payload_omits_hidden_when_model_opts_out(): runner = _make_async_output_runner() runner.model.omni_pooler_payload_include_hidden = False + runner._pooler_payload_include_hidden_flag = False payload = GPUARModelRunner._build_omni_async_snapshot_payload( runner, @@ -819,6 +841,8 @@ def test_runner_assisted_full_attention_metadata_refresh_pads_buffers(): self.commits.append(num_reqs_padded) runner = object.__new__(GPUARModelRunner) + runner._omni_cache_policy = ModelCachePolicy(needs_full_hidden_states=True) + runner._pooler_payload_include_hidden_flag = True block_table = BlockTable() runner.input_batch = SimpleNamespace( num_computed_tokens_cpu=torch.tensor([10, 20, 99, 99], dtype=torch.int32), @@ -846,6 +870,8 @@ def test_runner_assisted_full_attention_metadata_refresh_pads_buffers(): @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._omni_cache_policy = ModelCachePolicy(needs_full_hidden_states=True) + runner._pooler_payload_include_hidden_flag = True runner.execute_model_state = ExecuteModelState( SimpleNamespace( total_num_scheduled_tokens=3, @@ -908,12 +934,11 @@ def test_sample_tokens_tail_only_prefix_cache_uses_staged_cpu_hidden_states(monk ) 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), + "_prepare_prefix_cache_pooler_payload_sources", + lambda self, **kwargs: (kwargs["staged_hidden_states_cpu"], None, None), ) monkeypatch.setattr(GPUARModelRunner, "_process_additional_information_updates", lambda *args, **kwargs: None) monkeypatch.setattr(GPUARModelRunner, "_should_accumulate_full_payload_output", lambda self: False) @@ -946,8 +971,6 @@ def test_build_omni_output_falls_back_to_mm_cpu_without_prefix_merge(monkeypatch lambda self, req_ids: ("latent", req_ids), ) monkeypatch.setattr(GPUARModelRunner, "_model_needs_full_prefix_hidden_states", lambda self: False) - monkeypatch.setattr(GPUARModelRunner, "_deferred_prefix_cache_mm_keys", lambda self: {"codes.audio"}) - monkeypatch.setattr(GPUARModelRunner, "_stage_deferred_prefix_cache_mm_outputs", lambda *args, **kwargs: None) monkeypatch.setattr( GPUARModelRunner, "_prepare_prefix_cache_pooler_payload_sources", @@ -1095,8 +1118,6 @@ def test_build_omni_output_filters_multimodal_by_partial_downstream_batch(monkey lambda self, req_ids: ("latent", ["r1", "r3"]), ) monkeypatch.setattr(GPUARModelRunner, "_model_needs_full_prefix_hidden_states", lambda self: False) - monkeypatch.setattr(GPUARModelRunner, "_deferred_prefix_cache_mm_keys", lambda self: {"codes.audio"}) - monkeypatch.setattr(GPUARModelRunner, "_stage_deferred_prefix_cache_mm_outputs", lambda *args, **kwargs: None) monkeypatch.setattr( GPUARModelRunner, "_prepare_prefix_cache_pooler_payload_sources", @@ -1172,8 +1193,6 @@ def test_build_omni_output_uses_combined_prefix_cache_mm_payload_for_partial_dow lambda self, req_ids: ("latent", ["r1", "r3"]), ) monkeypatch.setattr(GPUARModelRunner, "_model_needs_full_prefix_hidden_states", lambda self: False) - monkeypatch.setattr(GPUARModelRunner, "_deferred_prefix_cache_mm_keys", lambda self: {"codes.audio"}) - monkeypatch.setattr(GPUARModelRunner, "_stage_deferred_prefix_cache_mm_outputs", lambda *args, **kwargs: None) monkeypatch.setattr( GPUARModelRunner, "_prepare_prefix_cache_pooler_payload_sources", diff --git a/vllm_omni/core/prefix_cache.py b/vllm_omni/core/prefix_cache.py deleted file mode 100644 index 17a981140..000000000 --- a/vllm_omni/core/prefix_cache.py +++ /dev/null @@ -1,737 +0,0 @@ -""" -Utilities for Prefix Caching in Omni models. -""" - -from dataclasses import dataclass, field - -import torch -from vllm.logger import init_logger -from vllm.v1.worker.gpu_input_batch import InputBatch - -from vllm_omni.utils.mm_outputs import build_mm_cpu, to_payload_element - -logger = init_logger(__name__) - - -@dataclass -class _PendingAsyncWrite: - """One in-flight async write into the prefix cache. - - Holds the CPU tensors that PyTorch implicitly allocated via - ``.to("cpu", non_blocking=True)`` together with a CUDA event recorded - on the copy stream. The consumer calls ``event.synchronize()`` once - (blocking) to guarantee the D2H has completed before reading. - """ - - event: "torch.cuda.Event" - num_tokens: int - slots_cpu: torch.Tensor - hidden_cpu: torch.Tensor | None - mm_cpu: dict[str, torch.Tensor] = field(default_factory=dict) - - -class OmniTensorPrefixCache: - """Prefix cache for hidden states (model outputs) and model specific - multimodal outputs. - - This class implements prefix caching in a non-invasive way on top of - vLLM by leveraging the same slot mappings that the vLLM scheduler uses - for the KV Cache. - - Conceptually, this means we are mapping vLLM's cache mapping: - (num_blocks, block_size) - - to 3D tensors of shape: - (num_blocks, block_size, feature_size) - - Note that feature_size may vary across multimodal_outputs. - """ - - def __init__( - self, - num_blocks: int, - block_size: int, - hidden_size: int, - hs_dtype: torch.dtype, - ): - self.num_blocks = num_blocks - self.block_size = block_size - self.default_hidden_size = hidden_size - - # Initialize the hidden states cache immediately - self.hidden_states_cache = self._get_cache_tensor(dtype=hs_dtype) - - # Defer initialization of the mm_outputs_cache until we - # actually see mm output tensors dependent on num tokens. - 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]]]] = {} - - # Async write pipeline state: - # - one dedicated copy stream - # - one in-flight pending write at a time (the producer is exactly - # one forward step ahead of the consumer) - # - schedule_async_write() consumes the *previous* pending write - # (via event.synchronize() + scatter) before kicking off the new - # D2H. By the time we hit synchronize, the previous step's GPU - # work has been done for ~one forward step, so the sync is - # near-zero wall-clock. - self._async_initialized: bool = False - self._async_copy_stream: torch.cuda.Stream | None = None - self._pending_write: _PendingAsyncWrite | None = None - - def maybe_init_missing_mm_cache_keys(self, multimodal_outputs: dict, seq_len: int): - """Given multimodal outputs from executing the model, dynamically - determine which multimodal outputs are tensors depending on sequence - length and should be cached, and initialize the cache tensors - accordingly. - - NOTE: This is done to avoid the need for explicit specification of - cache keys for every model/stage and aligns with the current way - that we slice the multimodal outputs based on the first dimension. - - This will usually be called by the first forward pass, i.e., - determined by the warmup. - """ - for key, val in multimodal_outputs.items(): - # 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, - hidden_size=feat_dim, - ) - self.mm_cache_keys.add(key) - new_tensor_shape = self.mm_outputs_cache[key].shape - logger.info("Initializing multimodal output cache of size %s for key: %s", list(new_tensor_shape), key) - - def _get_cache_tensor(self, dtype: torch.dtype, hidden_size: int | None = None) -> torch.Tensor: - """Allocate a CPU cache tensor for a specific key. - - When CUDA is available the tensor is pinned: this is what lets the - async-write pipeline ride the GPU->CPU copy on the dedicated copy - stream as a true async ``cudaMemcpyAsync`` and lets the post-copy CPU - scatter avoid page-resolution stalls. On CPU-only builds pinning is - unsupported (``torch.zeros(pin_memory=True)`` raises), and the async - pipeline is disabled anyway (``schedule_async_write`` early-returns - when ``torch.cuda.is_available()`` is False), so we fall back to a - plain pageable allocation with no behavioral change. - """ - actual_hidden_size = hidden_size if hidden_size is not None else self.default_hidden_size - return torch.zeros( - (self.num_blocks, self.block_size, actual_hidden_size), - dtype=dtype, - device="cpu", - pin_memory=torch.cuda.is_available(), - ) - - # ===================================================================== - # Async write pipeline: - # schedule_async_write : called per forward step, enqueues a - # non-blocking D2H + records a CUDA event - # drain_ready_async_writes : non-blocking poll, applies any writes - # whose D2H has completed - # ===================================================================== - - def _ensure_async_resources(self) -> None: - """Lazily allocate the dedicated copy stream. - - The per-step CPU staging tensors are allocated by PyTorch on demand - inside ``.to("cpu", non_blocking=True)``, which gives them - implicitly-pinned memory and lets PyTorch's allocator amortize / - reuse them. - """ - if self._async_initialized: - return - if not torch.cuda.is_available(): - return - self._async_copy_stream = torch.cuda.Stream() - self._async_initialized = True - logger.info("Omni prefix cache async-write copy stream initialized.") - - def schedule_async_write( - self, - hidden_states_gpu: torch.Tensor | None, - multimodal_outputs_gpu: dict[str, torch.Tensor] | None, - slot_mapping_gpu: torch.Tensor, - num_tokens_unpadded: int, - num_tokens_padded: int, - skip_mm_cache_keys: set[str] | None = None, - ) -> None: - """Kick off a non-blocking GPU->CPU copy for this step's writes. - - 1. **Consume the previous step's write** (if any): block on its - event and scatter the CPU tensors into the cache. By the time we - reach this point, the previous step's GPU work has been done for - ~one forward step, so the synchronize is near-zero wall-clock. - 2. **Schedule THIS step's D2H** on a dedicated copy stream using - ``.to("cpu", non_blocking=True)``. PyTorch allocates the pinned - destination implicitly. - 3. **Record an event** and stash everything for the next step's - consume. - - Holding exactly one in-flight write at a time keeps the design - simple: no ring, no event-polling, no driver-mutex pressure from - repeated event.query() calls. - """ - if not torch.cuda.is_available(): - return - - self._ensure_async_resources() - if not self._async_initialized: - return - - # Decide which mm keys we'll cache. Init the dest cache tensors - # before scheduling the copies so the apply step has somewhere - # to scatter into. - skip = skip_mm_cache_keys or set() - cacheable_mm_keys: list[str] = [] - if multimodal_outputs_gpu: - self.maybe_init_missing_mm_cache_keys(multimodal_outputs_gpu, seq_len=num_tokens_padded) - for k in self.mm_cache_keys: - if k in skip: - continue - v = multimodal_outputs_gpu.get(k) - if not isinstance(v, torch.Tensor) or v.ndim < 2: - continue - if v.shape[0] < num_tokens_unpadded: - continue - cacheable_mm_keys.append(k) - - # Step 1: consume the previous step's write FIRST. The previous - # step's event was recorded ~one forward step ago, so - # synchronize() is effectively a no-op cost-wise; the scatter - # into pinned CPU caches is also fast. - self._consume_pending_write() - - if hidden_states_gpu is None and not cacheable_mm_keys: - return # nothing to schedule this step - - # Step 2: schedule the D2H on the copy stream. PyTorch implicitly - # allocates pinned host destinations for the ``.to("cpu", - # non_blocking=True)`` results. - copy_stream = self._async_copy_stream - default_stream = torch.cuda.current_stream() - with torch.cuda.stream(copy_stream): - copy_stream.wait_stream(default_stream) - - # slot_mapping: copy as-is; the int64 cast required by - # ``index_copy_`` is deferred to ``_consume_pending_write`` - # so we avoid launching an extra GPU kernel here that would - # contend the driver mutex with the compute stream. - slots_cpu = slot_mapping_gpu[:num_tokens_unpadded].to("cpu", non_blocking=True) - hidden_cpu: torch.Tensor | None = None - if hidden_states_gpu is not None: - hidden_cpu = hidden_states_gpu[:num_tokens_unpadded].to("cpu", non_blocking=True) - mm_cpu: dict[str, torch.Tensor] = {} - for k in cacheable_mm_keys: - mm_cpu[k] = multimodal_outputs_gpu[k][:num_tokens_unpadded].to("cpu", non_blocking=True) - - event = torch.cuda.Event() - event.record() - - # Step 3: stash for next step's consume. - self._pending_write = _PendingAsyncWrite( - event=event, - num_tokens=num_tokens_unpadded, - slots_cpu=slots_cpu, - hidden_cpu=hidden_cpu, - mm_cpu=mm_cpu, - ) - - def drain_ready_async_writes(self) -> int: - """Consume one in-flight write if it has finished (non-blocking). - - Exposed so the model runner can drain at the top of - ``execute_model`` regardless of whether a new write will be - scheduled this step. - """ - if self._pending_write is None: - return 0 - if not self._pending_write.event.query(): - return 0 - self._consume_pending_write() - return 1 - - def _consume_pending_write(self) -> None: - """Synchronize on the pending event and scatter the CPU tensors - into the destination caches. Drops state on completion. - """ - entry = self._pending_write - if entry is None: - return - self._pending_write = None - entry.event.synchronize() - n = entry.num_tokens - if n <= 0: - return - - slots = entry.slots_cpu - if slots.dtype != torch.int64: - slots = slots.to(torch.int64) - - # NOTE: ``index_copy_(0, idx, src)`` is the row-scatter primitive in - # PyTorch. It's semantically identical to ``flat[idx] = src`` but - # dispatches to a much simpler CPU implementation (single dim, - # no broadcasting, no accumulate), which on pinned memory at large - # row sizes is observed to be much faster than the generic - # ``aten::index_put_`` path used by advanced-indexing assignment. - if entry.hidden_cpu is not None: - flat = self.hidden_states_cache.view(-1, self.hidden_states_cache.shape[-1]) - flat.index_copy_(0, slots, entry.hidden_cpu) - for k, src_cpu in entry.mm_cpu.items(): - mm_cache = self.mm_outputs_cache.get(k) - if mm_cache is None: - continue - flat = mm_cache.view(-1, mm_cache.shape[-1]) - flat.index_copy_(0, slots, src_cpu) - - def add_prefix_cached_new_req_id(self, req_id: str): - """Adds a new request ID to the set of prefix cache hits on the batch.""" - self._new_req_cache_hit_ids.add(req_id) - - def reset_prefix_cached_new_req_ids(self): - """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.""" - 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, - multimodal_outputs: dict[str, torch.Tensor] | None, - num_tokens_unpadded: int, - 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. - - Args: - hidden_states: Hidden states tensor to cache (if any) - multimodal_outputs: Multimodal dict whose tensors may be cached - 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 - 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: - 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 - 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. - # ``index_copy_`` is the row-scatter primitive — it's semantically - # ``cache[idx] = src`` but dispatches to a faster CPU path than the - # generic advanced-indexing assignment (which goes through - # ``aten::index_put_`` and is dominated by per-row dispatch - # overhead on large feature widths). - flat_cache = self.hidden_states_cache.view(-1, self.hidden_states_cache.shape[-1]) - slot_idx = unpadded_slot_mapping - if slot_idx.dtype != torch.int64: - slot_idx = slot_idx.to(torch.int64) - flat_cache.index_copy_(0, slot_idx, hidden_states) - logger.debug("Writing to hidden states for %s tokens", num_tokens_unpadded) - - # Do the same for the stage's cached multimodal outputs - if multimodal_outputs is not None: - # If we haven't initialized the keys already, do it now - # We check against the padded token count since we haven't sliced yet - self.maybe_init_missing_mm_cache_keys( - multimodal_outputs, - seq_len=num_tokens_padded, - ) - - 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) - flat_cache = mm_cache.view(-1, mm_cache.shape[-1]) - slot_idx = unpadded_slot_mapping - if slot_idx.dtype != torch.int64: - slot_idx = slot_idx.to(torch.int64) - flat_cache.index_copy_(0, slot_idx, 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]) - if slots.dtype != torch.int64: - slots = slots.to(torch.int64) - flat_cache.index_copy_(0, 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, - query_start_loc: torch.Tensor, - input_batch: InputBatch, - num_scheduled_tokens: dict[str, int], - ) -> dict[str, object]: - """Build the multimodal passthrough data per request for - 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=total_scheduled_tokens, - ) - return elem_dict - - def get_merged_multimodal_states( - self, - query_start_loc: torch.Tensor, - input_batch: InputBatch, - multimodal_outputs: dict | None, - num_scheduled_tokens: dict[str, int], - ): - """Get the merged multimodal states if hidden state prefix caching is enabled.""" - combined_multimodal_outputs = {} - # Talkers that produce multimodal outputs only at decode (e.g. Higgs - # Audio v3's audio codes) have ``multimodal_outputs is None`` at the - # initial prefill step. Treat it as an empty mapping so the merger - # short-circuits cleanly and the cache merge still runs once codes - # start arriving. - mm_outputs = multimodal_outputs if multimodal_outputs is not None else {} - # First get the prefix cached tensors that are present in the mm data - for mm_key in self.mm_cache_keys: - if mm_key in mm_outputs: - combined_multimodal_outputs[mm_key] = self._get_merged_tensors( - query_start_loc=query_start_loc, - input_batch=input_batch, - cache=self.mm_outputs_cache[mm_key], - hidden_states=mm_outputs[mm_key], - num_scheduled_tokens=num_scheduled_tokens, - ) - - # Then, get everything else (passthrough data); first, convert to CPU - # tensors similarly to the non prefix cached path, and then populate - # the subdicts mapping request IDs -> payload objects - passthrough_keys = set(mm_outputs.keys()) - self.mm_cache_keys - passthrough_mm_data = {k: v for k, v in mm_outputs.items() if k in passthrough_keys} - mm_cpu = build_mm_cpu(multimodal_outputs=passthrough_mm_data) - - for mm_key, mm_val in mm_cpu.items(): - combined_multimodal_outputs[mm_key] = self._coerce_to_payload_dict( - element=mm_val, - query_start_loc=query_start_loc, - input_batch=input_batch, - num_scheduled_tokens=num_scheduled_tokens, - ) - return combined_multimodal_outputs - - 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( - 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, - ) - - def _get_merged_tensors( - self, - query_start_loc: torch.Tensor, - input_batch: InputBatch, - 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. 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: - logger.warning_once( - "Omni prefix caching is enabled, but the batch block table appears to" - " have multiple kv groups; only the first group will be used!" - ) - - combined_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] - - if req_id in self._new_req_cache_hit_ids: - block_ids = self._get_cached_block_ids(req_idx, input_batch) - cached_hs = cache[block_ids].reshape(-1, cache.shape[-1]) - - # Slice the hidden states corresponding to this request; - # we do this by using the query start - start = query_start_loc[req_idx] - new_hs = hidden_states[start : start + num_scheduled_tokens[req_id]] - combined_hidden_states[req_id] = torch.cat([cached_hs, new_hs], dim=0) - else: - # cache miss for this request, pass through normally - start = query_start_loc[req_idx] - new_hs = hidden_states[start : start + num_scheduled_tokens[req_id]] - combined_hidden_states[req_id] = new_hs - - return combined_hidden_states - - def _get_cached_block_ids(self, req_idx: int, input_batch: InputBatch) -> torch.Tensor: - """Given an input batch and request index in the batch (not ID), get the - block IDs corresponding to the cache hit. - """ - num_computed = input_batch.num_computed_tokens_cpu[req_idx] - # NOTE: vLLM only caches full blocks - num_cached_blocks = num_computed // self.block_size - # Get the block IDs attached to this cache hit and reindex into - # the flattened cached hidden states (i.e., 1 row per token). - return input_batch.block_table[0].block_table.cpu[req_idx, :num_cached_blocks] diff --git a/vllm_omni/core/prefix_cache/__init__.py b/vllm_omni/core/prefix_cache/__init__.py new file mode 100644 index 000000000..4d14da6fa --- /dev/null +++ b/vllm_omni/core/prefix_cache/__init__.py @@ -0,0 +1,58 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Omni prefix cache. + +Structure mirrors vllm/v1/core. + +Import invariant: this package must not import ``vllm`` at module +scope (``TYPE_CHECKING`` and function-body imports are fine). +``tests/core/test_prefix_cache.py`` loads these modules without a +vllm install; a top-level ``from vllm...`` turns that suite red. + +That is also why these files use ``logging.getLogger`` instead of +the repo-wide ``init_logger`` — ``init_logger`` comes from +``vllm.logger``. The deviation is required, not an oversight. + +``_merge_uncached_mm`` is the exception at *call* time: it lazily +imports ``vllm_omni.utils.mm_outputs``, which pulls in vllm. Import +of this package still succeeds without vllm; materialize of leftover +mm does not. +""" + +from vllm_omni.core.prefix_cache.block_pool import PrefixBlockPool +from vllm_omni.core.prefix_cache.controller import OmniPrefixCacheController +from vllm_omni.core.prefix_cache.group_view import ( + FullAttentionGroupView, + check_prefix_cache_kv_groups, + check_prefix_cache_kv_transfer, + get_prefix_cache_group_view, + stage_prefix_cache_config, +) +from vllm_omni.core.prefix_cache.interface import ( + HIDDEN_KEY, + ModelCachePolicy, + OmniPrefixCacheStagingTimeoutError, + OmniPrefixCacheUnmatchError, + PrefixCacheConfig, + StageCacheOutputs, + WriteSchedule, +) +from vllm_omni.core.prefix_cache.manager import OmniPrefixCacheManager + +__all__ = [ + "HIDDEN_KEY", + "FullAttentionGroupView", + "ModelCachePolicy", + "OmniPrefixCacheController", + "OmniPrefixCacheManager", + "OmniPrefixCacheStagingTimeoutError", + "OmniPrefixCacheUnmatchError", + "StageCacheOutputs", + "PrefixBlockPool", + "PrefixCacheConfig", + "WriteSchedule", + "check_prefix_cache_kv_groups", + "check_prefix_cache_kv_transfer", + "get_prefix_cache_group_view", + "stage_prefix_cache_config", +] diff --git a/vllm_omni/core/prefix_cache/block_pool.py b/vllm_omni/core/prefix_cache/block_pool.py new file mode 100644 index 000000000..9cf34cb2d --- /dev/null +++ b/vllm_omni/core/prefix_cache/block_pool.py @@ -0,0 +1,84 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Pinned CPU block mirror (mirrors vLLM BlockPool).""" + +import logging + +import torch + +from vllm_omni.core.prefix_cache.interface import PrefixCacheConfig, TensorName + +logger = logging.getLogger(__name__) + + +class PrefixBlockPool: + """Durable CPU mirror of vLLM KV slots, one tensor per name. + + Storage is ``(num_blocks, block_size, feat)`` per key, viewed flat as + ``(num_slots, feat)`` so vLLM slot ids index rows directly. Write + access is single-writer (the controller committer thread); readers + take row views. + + ``_caches`` is not an unbounded store and has no LRU. Dict keys are + tensor names (``__hidden_states__`` plus mm whose first dim is tokens), + opened once by ``ensure_key`` and never dropped — a model emits a + handful of those names, not a per-request set. Each value is a + fixed tensor sized to the same ``num_blocks`` as the upstream KV + pool. Rows *are* those kv slots: when vLLM recycles a block, the + next pool write overwrites that row. Slot reuse is the eviction. + """ + + def __init__(self, config: PrefixCacheConfig): + self._config = config + self._caches: dict[TensorName, torch.Tensor] = {} + + def _alloc(self, dtype: torch.dtype, feat: int) -> torch.Tensor: + return torch.zeros( + (self._config.num_blocks, self._config.block_size, feat), + dtype=dtype, + device="cpu", + # Pinning lets device→host overlap compute; unsupported on + # CPU-only builds where that overlap is off anyway. + pin_memory=torch.cuda.is_available(), + ) + + def alloc_key(self, key: str, dtype: torch.dtype, feat: int) -> torch.Tensor | None: + """Allocate storage for a new ``key`` without publishing it (None if + already open). Pinned allocation is slow; run this unlocked and + hand the tensor to ``install_key`` under the manager's state lock.""" + if key in self._caches: + return None + return self._alloc(dtype, feat) + + def install_key(self, key: str, storage: torch.Tensor) -> None: + if key in self._caches: + return + self._caches[key] = storage + logger.info("prefix_cache: initialized mirror %s for key %s", list(storage.shape), key) + + def ensure_key(self, key: str, dtype: torch.dtype, feat: int) -> None: + storage = self.alloc_key(key, dtype, feat) + if storage is not None: + self.install_key(key, storage) + + def has_key(self, key: str) -> bool: + return key in self._caches + + def keys(self) -> set[TensorName]: + return set(self._caches.keys()) + + def _flat(self, key: str) -> torch.Tensor: + cache = self._caches[key] + return cache.view(-1, cache.shape[-1]) + + def rows(self, key: str, slots: torch.Tensor) -> torch.Tensor: + """Gather rows for non-contiguous slots (returns a copy).""" + return self._flat(key).index_select(0, slots) + + def write(self, key: str, slots: torch.Tensor, src_cpu: torch.Tensor) -> None: + """Write rows into the pool; caller (committer thread) is the single writer.""" + if slots.dtype != torch.int64: + slots = slots.to(torch.int64) + # index_copy_ dispatches to a faster single-dim CPU path than + # advanced-indexing assignment (aten::index_put_). + self._flat(key).index_copy_(0, slots, src_cpu) diff --git a/vllm_omni/core/prefix_cache/controller.py b/vllm_omni/core/prefix_cache/controller.py new file mode 100644 index 000000000..4c26ac055 --- /dev/null +++ b/vllm_omni/core/prefix_cache/controller.py @@ -0,0 +1,976 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Runs WriteTasks and the step device→host staging pool. + +The manager owns request/slot identity and when to submit. This class +owns the staging pool, the GPU-byte budget, the copy queues, and the +single committer that writes into the CPU block pool. + +Two device→host paths — the step path does not write `chunk.host` in +the committer: + + JOIN_NEXT_STEP save already launched a whole-step device→host into + a staging slot and set `chunk.host` as views. + Committer waits that `step_d2h_event`, then writes + the pool. + JOIN_ON_FINISH committer copies the device clone → owned host + tensors, then writes the pool. + +Async: high-priority then low-priority queues, then pool write. +Eager: submit() does wait+write inline (CPU tests / no CUDA). +""" + +from __future__ import annotations + +import logging +import threading +import time +from collections import deque +from collections.abc import Callable +from dataclasses import dataclass, field +from enum import Enum, auto +from typing import Literal, NamedTuple + +import torch + +from vllm_omni.core.prefix_cache.block_pool import PrefixBlockPool +from vllm_omni.core.prefix_cache.interface import ( + OmniPrefixCacheStagingTimeoutError, + OmniPrefixCacheUnmatchError, + PrefixCacheConfig, + ReqId, + TensorName, + Tid, + WriteSchedule, +) + +logger = logging.getLogger(__name__) + + +class StagingBufferHolder(NamedTuple): + """One owner of a staging-buffer slot. The slot is free when none remain. + + Not a buffer state — concurrent owners share the same slot: + - for_step: claimed at save, released when materialize/discard consumes the ctx + - for_task: bound before WriteTask submit, released when that task completes + """ + + kind: Literal["step", "task"] + owner_id: int + + @classmethod + def for_step(cls, step_id: int) -> StagingBufferHolder: + return cls("step", step_id) + + @classmethod + def for_task(cls, tid: int) -> StagingBufferHolder: + return cls("task", tid) + + +@dataclass +class _BudgetTicket: + """GPU-byte-budget charge for one device clone shared by several writes. + + Charged once at ``reserve()``; freed when the last pinned tid releases. + Immediate: the whole-step clone, pinned by that step's JOIN_NEXT_STEP + tasks. Deferred: the shared deferred clone, pinned by every task a + chunk of it was appended to. Does not store tensors. + """ + + nbytes: int + tids: set[Tid] = field(default_factory=set) + _freed: bool = False + + def release(self, tid: Tid) -> int: + """Drop `tid`; return the bytes to uncharge (nbytes once, when the + last holder leaves; 0 otherwise). Idempotent per tid.""" + self.tids.discard(tid) + if self.tids or self._freed: + return 0 + self._freed = True + return self.nbytes + + +@dataclass +class _WriteChunk: + """One save's slots + tensors appended onto a WriteTask. + + Controller identity is the parent ``tid``, not a request. Deferred + writes grow by appending one chunk per save. + """ + + slots_cpu: torch.Tensor # int64 flat row ids, in token order + tensors: dict[TensorName, torch.Tensor] # freeze (device or eager CPU) + host: dict[TensorName, torch.Tensor] = field(default_factory=dict) # host view of rows [0:n) + # GPU-byte charge of the clone these tensors view; shared across the + # chunks cut from one save. None for chunks with no device clone. + budget: _BudgetTicket | None = None + + +class TaskState(Enum): + """WriteTask lifecycle: a strict chain, plus FAILED from any non-terminal. + + PENDING -> QUEUED -> COPYING -> HOST_READY -> WRITTEN + + PENDING registered, not on a copy queue (deferred: waiting for + finish / budget pressure) + QUEUED on a copy queue + COPYING one thread owns the copy stage (staging: waits the step + device→host; deferred: copies the device clone) + HOST_READY host rows ready (`host_ready` set); device clone dropped + WRITTEN in the CPU pool (`done` set) + FAILED committer could not finish (`host_ready` + `done` set; + manager raises from every later entry) + + Moves only through `WriteTask.transition`; skipping a step is illegal. + """ + + PENDING = auto() + QUEUED = auto() + COPYING = auto() + HOST_READY = auto() + WRITTEN = auto() + FAILED = auto() + + @property + def is_terminal(self) -> bool: + return self in (TaskState.WRITTEN, TaskState.FAILED) + + +# The one legal forward move from each non-terminal state. +_CHAIN: dict[TaskState, TaskState] = { + TaskState.PENDING: TaskState.QUEUED, + TaskState.QUEUED: TaskState.COPYING, + TaskState.COPYING: TaskState.HOST_READY, + TaskState.HOST_READY: TaskState.WRITTEN, +} +_NOT_YET_COPYING = frozenset({TaskState.PENDING, TaskState.QUEUED}) + + +def _is_legal_move(frm: TaskState, to: TaskState) -> bool: + return (to is TaskState.FAILED and not frm.is_terminal) or _CHAIN.get(frm) is to + + +@dataclass +class WriteTask: + """One write of (slot, key) rows for a single request. + + Identity: `tid` is the handle in the manager's (slot, key) tables. + `req_id` + `write_n` mark whose write this is and the nth time that + request opened a write. One write may cover several keys. + + `state` is the single source of truth for where the write is (see + `TaskState`). `host_ready` / `done` are wait primitives set by the + transitions into HOST_READY / WRITTEN / FAILED, never directly. + + How HOST_READY is reached: + - JOIN_NEXT_STEP: `chunk.host` is a staging view set at save; the + copy stage only waits `step_d2h_event`. + - JOIN_ON_FINISH: the copy stage copies device→host into `chunk.host`. + + `JOIN_NEXT_STEP` is queued at submit. `JOIN_ON_FINISH` stays + PENDING until finish or GPU-byte-budget pressure escalates it; + budget flush takes the unfinished task with the oldest `enqueued_time`. + + Concurrent readers/writers: + - Staging readers (materialize clone, committer pool write) all wait + the same `step_d2h_event` before touching the view. + - A later task taking the same (slot, key) records those rows in + `reassigned`; the old pool write skips them. + - `append_chunk` is refused once the copy stage is claimed; the caller + opens a fresh task rather than mutating a closed one. + - `lock` covers `state` / `reassigned` / `append_chunk` / host↔freeze / + `slot_to_row`. `scatter_rows` snapshots `reassigned`. + """ + + tid: Tid + req_id: ReqId + write_n: int # 1-based: nth write opened by this request + schedule: WriteSchedule + chunks: list[_WriteChunk] + # On-device clone finished on the compute stream. Copy/read streams + # must wait this before touching `chunks[].tensors`, or they can read + # the next CUDA-graph static-buffer overwrite. + freeze_event: torch.cuda.Event | None = None + state: TaskState = TaskState.PENDING + # Slots this write no longer owns (a newer write took them over). + reassigned: dict[TensorName, torch.Tensor] = field(default_factory=dict) + # Set on entering HOST_READY (or FAILED): this write's host rows are ready. + host_ready: threading.Event = field(default_factory=threading.Event) + # Set on entering WRITTEN (or FAILED): the CPU-pool write is over. + done: threading.Event = field(default_factory=threading.Event) + # Guards state / reassigned / append / host↔freeze. + lock: threading.Lock = field(default_factory=threading.Lock) + # time.monotonic() at submit; GPU-byte flush picks the oldest of these. + enqueued_time: float = field(default_factory=time.monotonic) + # Immediate path: this write's view into the shared step staging page. + staging_slot: int | None = None + # Immediate path: this step's CUDA device→host event (shared). None if deferred. + step_d2h_event: torch.cuda.Event | None = None + # slot -> (which chunk, row in that tensor). Built on demand + # when a write has more than one `_WriteChunk`; pool write uses slot, the tensor uses row. + _slot_to_row: dict[int, tuple[int, int]] | None = None + + def add_reassigned(self, key: str, slots: torch.Tensor) -> None: + with self.lock: + prev = self.reassigned.get(key) + self.reassigned[key] = slots.clone() if prev is None else torch.cat([prev, slots]) + + def reassigned_intersects(self, chunk: _WriteChunk) -> bool: + """True if this write already lost ownership of any (key, slot) the + chunk would add; the caller opens a fresh WriteTask instead.""" + with self.lock: + if not self.reassigned: + return False + for key in chunk.tensors: + taken = self.reassigned.get(key) + if taken is not None and taken.numel() and bool(torch.isin(chunk.slots_cpu, taken).any()): + return True + return False + + # ------------------------------------------------------------ lifecycle + + def transition(self, to: TaskState) -> None: + """Move to `to`; raises unless it is the next chain step or FAILED.""" + with self.lock: + self._transition_locked(to) + + def try_transition(self, to: TaskState) -> bool: + """Move to `to` if legal from the current state. True if moved.""" + with self.lock: + if not _is_legal_move(self.state, to): + return False + self._transition_locked(to) + return True + + def _transition_locked(self, to: TaskState, wake: bool = True) -> None: + if not _is_legal_move(self.state, to): + raise OmniPrefixCacheUnmatchError(f"task {self.tid}: illegal transition {self.state.name} -> {to.name}") + self.state = to + if not wake: + return + if to is TaskState.HOST_READY: + self.host_ready.set() + elif to is TaskState.WRITTEN: + self.done.set() + elif to is TaskState.FAILED: + self.host_ready.set() + self.done.set() + + def claim_copy(self) -> bool: + """QUEUED -> COPYING. Only one thread may run the copy stage; True if this caller won.""" + return self.try_transition(TaskState.COPYING) + + @property + def is_terminal(self) -> bool: + return self.state.is_terminal + + def append_chunk(self, chunk: _WriteChunk, freeze_event: torch.cuda.Event | None = None) -> TaskState | None: + """Grow this write with one save's rows. Returns None when appended, + else the state that closed the task (COPYING or later). + + `freeze_event` is stored in the same snapshot: events on one compute + stream are ordered, so the newest also covers every earlier clone. + """ + with self.lock: + if self.state not in _NOT_YET_COPYING: + return self.state + self.chunks.append(chunk) + self._slot_to_row = None + if freeze_event is not None: + self.freeze_event = freeze_event + return None + + def get_host_tensor(self, si: int, key: str) -> torch.Tensor | None: + """`chunks[si]` host if written, else device freeze. One snapshot.""" + with self.lock: + chunk = self.chunks[si] + src = chunk.host.get(key) + if src is None: + src = chunk.tensors.get(key) + return src + + def set_host_tensor(self, rows: list[tuple[_WriteChunk, str, torch.Tensor]]) -> None: + """COPYING -> HOST_READY: write these host tensors, drop the device freeze.""" + with self.lock: + for chunk, key, tensor in rows: + chunk.host[key] = tensor + self._clear_tensors() + self._transition_locked(TaskState.HOST_READY) + + def mark_host_ready(self) -> None: + """COPYING -> HOST_READY for staging views: wait the step device→host, drop freeze.""" + if self.step_d2h_event is not None: + self.step_d2h_event.synchronize() + with self.lock: + self._clear_tensors() + self._transition_locked(TaskState.HOST_READY) + + def mark_failed(self, *, defer_wake: bool = False) -> bool: + """-> FAILED from any non-terminal state; unblocks joiners. False if + already terminal. ``defer_wake``: the caller publishes the failure + record first, then calls ``wake_failed_waiters``.""" + with self.lock: + if self.state.is_terminal: + return False + self._clear_tensors() + self._transition_locked(TaskState.FAILED, wake=not defer_wake) + return True + + def wake_failed_waiters(self) -> None: + """Set the events after the failure record is published (see + ``mark_failed(defer_wake=True)``): a woken joiner must find it.""" + self.host_ready.set() + self.done.set() + + def mark_done(self) -> None: + """HOST_READY -> WRITTEN.""" + self.transition(TaskState.WRITTEN) + + def scatter_rows(self) -> list[tuple[TensorName, torch.Tensor, torch.Tensor]]: + """`(key, slots, host)` to write, one entry per key. Omits slots in + `reassigned`. A slot written twice by this task (preempt + resume + onto the same block) keeps the later chunk's row. + """ + with self.lock: + reassigned = {k: s.clone() for k, s in self.reassigned.items()} + by_key: dict[TensorName, tuple[list[torch.Tensor], list[torch.Tensor]]] = {} + for chunk in self.chunks: + for k, host in chunk.host.items(): + s, h = by_key.setdefault(k, ([], [])) + s.append(chunk.slots_cpu) + h.append(host) + out: list[tuple[TensorName, torch.Tensor, torch.Tensor]] = [] + for k, (s, h) in by_key.items(): + slots = s[0] if len(s) == 1 else torch.cat(s) + host = h[0] if len(h) == 1 else torch.cat(h, dim=0) + if len(s) > 1 and torch.unique(slots).numel() != slots.numel(): + last_pos = {int(slot): i for i, slot in enumerate(slots.tolist())} + keep_idx = torch.tensor(sorted(last_pos.values()), dtype=torch.int64) + slots, host = slots[keep_idx], host[keep_idx] + taken = reassigned.get(k) + if taken is not None and taken.numel(): + keep = ~torch.isin(slots, taken) + if not bool(keep.any()): + continue + slots, host = slots[keep], host[keep] + out.append((k, slots, host)) + return out + + def clear_tensors(self) -> None: + """Drop the device freeze. Host is unchanged (staging wait / fail).""" + with self.lock: + self._clear_tensors() + + def _clear_tensors(self) -> None: + for chunk in self.chunks: + chunk.tensors = {} + + def budget_tickets(self) -> list[_BudgetTicket]: + """Distinct tickets across this task's chunks (by identity; one per save).""" + out: dict[int, _BudgetTicket] = {} + for chunk in self.chunks: + if chunk.budget is not None: + out.setdefault(id(chunk.budget), chunk.budget) + return list(out.values()) + + def slot_to_row(self) -> dict[int, tuple[int, int]]: + with self.lock: + if self._slot_to_row is None: + m: dict[int, tuple[int, int]] = {} + for si, chunk in enumerate(self.chunks): + for ri, slot in enumerate(chunk.slots_cpu.tolist()): + m[slot] = (si, ri) + self._slot_to_row = m + return self._slot_to_row + + +class StagingBufferPool: + """Reusable pinned landing zone for ONE whole-step device→host at save. + + Per-task `chunk.host` is a row-range view into a slot, so the committer + skips a per-task device→host. Slots recycle; this is not the CPU block pool. + + A slot stays busy while anyone still holds it: the step (until + materialize/discard) and each immediate write that views the page + (until its pool write). Prefix hits do not hold a slot — + they wait for the pool write and read the durable pool. + + Saves with only leftover mm still claim a slot (empty views) so + every step id shares this bound. A full pool waits; timeout then errors. + """ + + def __init__(self, depth: int, capacity: int): + self.depth = depth + self.capacity = capacity # rows per slot; a larger step raises + self._bufs: dict[TensorName, torch.Tensor] = {} # [depth*capacity, width] + self._busy: list[set[StagingBufferHolder]] = [set() for _ in range(depth)] + self._slot_free_condition = threading.Condition() + self._closed = False + + def _buf(self, key: str, width: int, dtype: torch.dtype, pin: bool) -> torch.Tensor: + buf = self._bufs.get(key) + if buf is None or buf.shape[-1] != width or buf.dtype != dtype: + buf = torch.empty((self.depth * self.capacity, width), dtype=dtype, pin_memory=pin) + self._bufs[key] = buf + return buf + + def claim(self, holder: StagingBufferHolder, timeout: float) -> int: + """Grab a free slot for `holder`. Waits until one is free, then + times out. ``timeout<=0`` fails immediately if none are free. + """ + deadline = time.monotonic() + max(0.0, timeout) + with self._slot_free_condition: + while True: + if self._closed: + raise OmniPrefixCacheUnmatchError("staging pool shut down while waiting for a slot") + for slot in range(self.depth): + if not self._busy[slot]: + self._busy[slot].add(holder) + return slot + remaining = deadline - time.monotonic() + if remaining <= 0: + raise OmniPrefixCacheStagingTimeoutError(f"timed out waiting for a staging slot after {timeout}s") + self._slot_free_condition.wait(timeout=remaining) + + def bind(self, slot: int, holder: StagingBufferHolder) -> None: + with self._slot_free_condition: + self._busy[slot].add(holder) + + def release(self, slot: int, holder: StagingBufferHolder) -> None: + with self._slot_free_condition: + self._busy[slot].discard(holder) + if not self._busy[slot]: + self._slot_free_condition.notify() + + def close(self) -> None: + with self._slot_free_condition: + self._closed = True + self._slot_free_condition.notify_all() + + def views(self, slot: int, key: str, n: int, width: int, dtype: torch.dtype, pin: bool) -> torch.Tensor: + base = slot * self.capacity + return self._buf(key, width, dtype, pin)[base : base + n] + + +@dataclass +class StepD2HClaim: + """One whole-step landing in StagingBufferPool. + + Return of ``stage_step_host``. The manager stores this on + ``_StepContext`` until materialize/discard releases the step holder. + """ + + staging_slot: int # StagingBufferPool index + views: dict[TensorName, torch.Tensor] # host rows [0:n) + event: torch.cuda.Event | None = None # torch.cuda.Event; None on eager/CPU + + +class OmniPrefixCacheController: + """Staging pool + committer. Step device→host is launched at save; + this thread waits that event (JOIN_NEXT_STEP) or copies deferred + rows (JOIN_ON_FINISH), then writes into the CPU pool. + """ + + def __init__(self, pool: PrefixBlockPool, config: PrefixCacheConfig, eager: bool | None = None): + self._pool = pool + self._config = config + self._eager = (not torch.cuda.is_available()) if eager is None else eager + self._tasks: dict[Tid, WriteTask] = {} + self._completed: deque[Tid] = deque() # pool write done, awaiting manager drain + self._failed: deque[Tid] = deque() # write failed; manager raises from every later entry (sticky) + self._staged_bytes = 0 + self._lock = threading.Lock() + self._wake = threading.Condition(self._lock) + self._queue_hi: deque[Tid] = deque() # JOIN_NEXT_STEP + forced JOIN_ON_FINISH + self._queue_lo: deque[Tid] = deque() # JOIN_ON_FINISH waiting for finish/budget + self._blocked: list[Tid] = [] # host copy done, awaiting pool write + self._shutdown = False + self._copy_stream: torch.cuda.Stream | None = None + self._read_stream: torch.cuda.Stream | None = None + self._worker: threading.Thread | None = None + self._staging_pool = StagingBufferPool(config.staging_depth, config.staging_capacity_tokens) + if not self._eager: + self._copy_stream = torch.cuda.Stream() + self._read_stream = torch.cuda.Stream() + self._worker = threading.Thread(target=self._worker_loop, name="omni-prefix-cache-committer", daemon=True) + self._worker.start() + + # --------------------------------------------------------- step device→host staging + + def _d2h_on_stream( + self, + stream: torch.cuda.Stream, + freeze_event: torch.cuda.Event | None, + copy: Callable[[], None], + ) -> torch.cuda.Event: + """Issue device→host on `stream` after freeze; return the done event.""" + with torch.cuda.stream(stream): + if freeze_event is not None: + stream.wait_event(freeze_event) + copy() + ev = torch.cuda.Event() + ev.record() + return ev + + def stage_step_host( + self, + tensors: dict[str, torch.Tensor], + n: int, + freeze_event: torch.cuda.Event | None, + step_holder: StagingBufferHolder, + ) -> StepD2HClaim: + """Claim a staging slot and, when `tensors` is non-empty, launch + ONE whole-step device→host into it. + + Leftover-only saves pass empty `tensors` and still take a slot + (empty views) so every step id shares this bound. A full pool + waits for materialize/discard; timeout then errors. A step + larger than the page overflows the next slot — that is a config + break. The caller binds tasks after submit; `step_holder` is + released by materialize/discard via staging_release. + """ + if tensors and n > self._staging_pool.capacity: + raise OmniPrefixCacheUnmatchError( + f"step has {n} tokens; staging capacity is {self._staging_pool.capacity} " + "(size staging_capacity_tokens to max_num_batched_tokens)" + ) + slot = self._staging_pool.claim(step_holder, self._config.staging_claim_timeout_s) + try: + pin = not self._eager + views: dict[str, torch.Tensor] = {} + event: torch.cuda.Event | None = None + if self._eager or all(t.device.type == "cpu" for t in tensors.values()): + for key, src in tensors.items(): + v = self._staging_pool.views(slot, key, n, int(src.shape[-1]), src.dtype, pin) + v.copy_(src) + views[key] = v + else: + + def _copy_to_staging() -> None: + for key, src in tensors.items(): + v = self._staging_pool.views(slot, key, n, int(src.shape[-1]), src.dtype, pin) + v.copy_(src, non_blocking=True) + views[key] = v + + event = self._d2h_on_stream(self._copy_stream, freeze_event, _copy_to_staging) + return StepD2HClaim(staging_slot=slot, views=views, event=event) + except Exception: + self._staging_pool.release(slot, step_holder) + raise + + def staging_bind(self, slot: int, holder: StagingBufferHolder) -> None: + self._staging_pool.bind(slot, holder) + + def staging_release(self, slot: int, holder: StagingBufferHolder) -> None: + self._staging_pool.release(slot, holder) + + def _release_task_slot(self, task: WriteTask) -> None: + """Drop the task's staging hold at a terminal state (pool write or failure).""" + if task.staging_slot is not None: + self._staging_pool.release(task.staging_slot, StagingBufferHolder.for_task(task.tid)) + + def in_flight_tasks(self) -> int: + """Registered tasks not yet drained (diagnostics only).""" + return len(self._tasks) + + # ------------------------------------------------------------------ submit + + def register(self, task: WriteTask, queued: bool = True) -> None: + """Make a task visible (registry + QUEUED) without running anything; + safe under the manager's state lock. queued=False (deferred tasks) + stays PENDING on the GPU clone until finish/abort or the GPU-byte + budget forces a copy. Queued tasks must then go through + ``dispatch``. + + Caller must reserve() the task bytes first (budget flush can + block; the manager does that outside the state lock) and pin the + task on its budget ticket(s) before register. + """ + task.enqueued_time = time.monotonic() + with self._lock: + self._tasks[task.tid] = task + if queued: + task.transition(TaskState.QUEUED) + + def dispatch(self, tasks: list[WriteTask]) -> None: + """Hand registered, queued tasks to the copy path. Threaded: enqueue. + Eager: the copy + pool write run here, inline — never call this + under the manager's state lock.""" + if not tasks: + return + if self._eager: + try: + for task in tasks: + self._run_eager(task) + except Exception: + # Later tasks were registered but will never be dispatched. + # Fail them too so shared budget owners and waiters release. + for task in tasks: + self._fail_task(task.tid) + raise + return + with self._wake: + for task in tasks: + (self._queue_hi if task.schedule is WriteSchedule.JOIN_NEXT_STEP else self._queue_lo).append(task.tid) + self._wake.notify_all() + + def submit(self, task: WriteTask, queued: bool = True) -> None: + """register + dispatch in one call (callers not holding the state lock).""" + self.register(task, queued) + if queued: + self.dispatch([task]) + + def append_chunk( + self, task: WriteTask, chunk: _WriteChunk, freeze_event: torch.cuda.Event | None = None + ) -> TaskState | None: + """Append to a pending task. None when appended, else the closing state.""" + return task.append_chunk(chunk, freeze_event) + + def pin_budget(self, ticket: _BudgetTicket, tid: int) -> None: + """Record that ``tid`` holds a view of the clone ``ticket`` charges.""" + with self._lock: + ticket.tids.add(tid) + + def reserve(self, nbytes: int) -> None: + """Reserve GPU-clone bytes; blocking flush happens here, so callers + must not hold the manager's state lock.""" + self._reserve_bytes(nbytes) + + def _release_staged_bytes(self, task: WriteTask) -> None: + """Drop this task's hold on its budget ticket(s). Idempotent; a + ticket uncharges once, when its last holder leaves.""" + tickets = task.budget_tickets() + with self._wake: + for ticket in tickets: + self._staged_bytes -= ticket.release(task.tid) + + def _reserve_bytes(self, nbytes: int) -> None: + # GPU-byte budget: force-copy oldest pending tasks until under + # budget. Bounded wait: their device→host has usually long completed. + while True: + with self._lock: + pending = [tid for tid, t in self._tasks.items() if not t.is_terminal] + if self._staged_bytes + nbytes <= self._config.gpu_staging_bytes or not pending: + # Under budget or no pending tasks; admit reservation. + self._staged_bytes += nbytes + return + oldest = min(pending, key=lambda tid: self._tasks[tid].enqueued_time) + logger.warning("omni prefix cache: staging cap hit, force-flushing task %d", oldest) + self.escalate([oldest]) + self.join([oldest]) + + # ------------------------------------------------------------- lifecycle + + def escalate(self, tids: list[int]) -> None: + """Move pending tasks to the front of the high-priority queue. + + PENDING -> QUEUED at the head; QUEUED on the low-priority queue + moves up; QUEUED already high-priority is a no-op. COPYING and later + belong to the worker and are untouched: the worker claims under + `_wake` too, so a task cannot be popped and re-queued behind its back. + """ + if self._eager: + self.dispatch([task for tid in tids if (task := self._tasks.get(tid)) is not None and not task.is_terminal]) + return + with self._wake: + for tid in tids: + task = self._tasks.get(tid) + if task is None: + continue + if task.try_transition(TaskState.QUEUED): + self._queue_hi.appendleft(tid) + elif task.state is TaskState.QUEUED and tid in self._queue_lo: + self._queue_lo.remove(tid) + self._queue_hi.appendleft(tid) + self._wake.notify_all() + + def join(self, tids: list[int]) -> None: + """Block until each task has finished the CPU-pool write (or failed).""" + self._wait_tasks(tids, "done") + + def join_host_ready(self, tids: list[int]) -> None: + """Block until each task's device→host is complete (`host_ready`). + + Staging: committer has waited `step_d2h_event`. Deferred: committer + has written `chunk.host`. Does not wait for the CPU-pool write. + """ + self._wait_tasks(tids, "host_ready") + + def _wait_tasks(self, tids: list[int], event: Literal["done", "host_ready"]) -> None: + timeout = self._config.staging_claim_timeout_s + deadline = time.monotonic() + timeout + for tid in tids: + task = self._tasks.get(tid) + if task is None: + continue + if not getattr(task, event).wait(max(0.0, deadline - time.monotonic())): + raise OmniPrefixCacheUnmatchError( + f"task {tid} ({task.schedule.value}) did not reach {event} within {timeout:g}s " + f"(state={task.state.name}, in_flight_tasks={len(self._tasks)})" + ) + if task.state is TaskState.FAILED: + # FAILED sets both events; success here would read rows the + # write never produced. The one-shot failure record may have + # been drained by another thread — do not rely on it. + raise OmniPrefixCacheUnmatchError(f"task {tid} ({task.schedule.value}) write failed before {event}") + + def drain_completed(self) -> list[int]: + """Pop pool-written tasks from `_completed` and drop them from `_tasks`. + Staging holders were already released at the pool write.""" + out: list[int] = [] + with self._lock: + while self._completed: + out.append(self._completed.popleft()) + for tid in out: + self._tasks.pop(tid, None) + return out + + def drain_failed(self) -> list[int]: + """Pop failed task ids from `_failed`. Does not drop `_tasks`.""" + out: list[int] = [] + with self._lock: + while self._failed: + out.append(self._failed.popleft()) + return out + + def get_task(self, tid: int) -> WriteTask | None: + return self._tasks.get(tid) + + def shutdown(self) -> None: + self._staging_pool.close() + with self._wake: + self._shutdown = True + self._wake.notify_all() + if self._worker is not None: + self._worker.join(timeout=5.0) + + # ------------------------------------------------------------ fetch_host + + @torch.inference_mode() + def fetch_host(self, task: WriteTask, slots: torch.Tensor, key: str) -> torch.Tensor: + """Rows for `slots` of one not-yet-done JOIN_ON_FINISH task. + + `_slot_ref` puts JOIN_NEXT_STEP tids in `join_tids` (wait then + pool). This path reads committer-written `chunk.host`, or the + device clone if that device→host has not landed. + """ + if task.step_d2h_event is not None: + task.step_d2h_event.synchronize() + return self._rows_from(task, slots, key) + + def _rows_from(self, task: WriteTask, slots: torch.Tensor, key: str) -> torch.Tensor: + """Map `slots` to rows across one or more `_WriteChunk`s; preserve caller order.""" + s2r = task.slot_to_row() + idx = [s2r[int(s)] for s in slots.tolist()] + parts: list[torch.Tensor] = [] + order: list[int] = [] + pos = 0 + rows_groups: dict[int, list[tuple[int, int]]] = {} + for si, ri in idx: + rows_groups.setdefault(si, []).append((pos, ri)) + pos += 1 + for si, items in rows_groups.items(): + src = task.get_host_tensor(si, key) + if src is None: + continue + rows_idx = torch.tensor([ri for _, ri in items], dtype=torch.long) + picked = self._slice_rows(task, src, rows_idx, host=(src.device.type == "cpu")) + parts.append(picked) + order.extend(p for p, _ in items) + if not parts: + raise KeyError(f"key {key} not present in task {task.tid}") + cat = torch.cat(parts, dim=0) + out = torch.empty_like(cat) + out[torch.tensor(order, dtype=torch.long)] = cat + return out + + def _slice_rows(self, task: WriteTask, src: torch.Tensor, rows_idx: torch.Tensor, host: bool) -> torch.Tensor: + n = int(rows_idx.numel()) + # Ascending-run check without materializing an arange: endpoints plus + # a monotonic diff are enough, and the common case is one long run. + contiguous = ( + n > 0 and int(rows_idx[-1]) - int(rows_idx[0]) == n - 1 and (n < 2 or bool((rows_idx.diff() == 1).all())) + ) + + def _pick() -> torch.Tensor: + return src[rows_idx[0] : rows_idx[0] + n] if contiguous else src.index_select(0, rows_idx) + + if host or src.device.type == "cpu": + return _pick() + if self._read_stream is None: + return _pick().detach().cpu() + copied: list[torch.Tensor] = [] + + def _copy_to_cpu() -> None: + copied.append(_pick().to("cpu", non_blocking=True)) + + ev = self._d2h_on_stream(self._read_stream, task.freeze_event, _copy_to_cpu) + ev.synchronize() + return copied[0] + + # ------------------------------------------------------------ eager mode + + @torch.inference_mode() + def _run_eager(self, task: WriteTask) -> None: + """Eager stand-in for queue + worker: walk the whole chain inline.""" + task.try_transition(TaskState.QUEUED) # no-op if submit already queued it + if not task.claim_copy(): + # Copy stage already ran (a held stub, or a prior escalate); finish the write. + if task.state is TaskState.HOST_READY: + self._scatter(task) + return + if task.schedule is WriteSchedule.JOIN_NEXT_STEP: + # Host is already a staging view (copied at save). Drop freeze. + # `step_d2h_event` is None on CPU; `mark_host_ready` skips wait. + task.mark_host_ready() + else: + # JOIN_ON_FINISH: no copy stream; freeze → owned host inline. + task.set_host_tensor( + [ + (chunk, k, t.detach().cpu() if t.device.type != "cpu" else t.clone()) + for chunk in task.chunks + for k, t in chunk.tensors.items() + ] + ) + self._release_staged_bytes(task) + self._scatter(task) + + # ---------------------------------------------------------- worker loop + + def _worker_loop(self) -> None: + # A dying committer would strand every join() forever, so the loop + # never propagates: it fails the offending task and keeps serving. + while True: + tid = None + try: + task: WriteTask | None = None + with self._wake: + while not self._shutdown and not self._queue_hi and not self._queue_lo: + if self._blocked: + break + # submit / escalate / shutdown all notify. + self._wake.wait() + if self._shutdown and not self._queue_hi and not self._queue_lo and not self._blocked: + return + if self._queue_hi: + tid = self._queue_hi.popleft() + elif self._queue_lo: + tid = self._queue_lo.popleft() + if tid is not None: + task = self._tasks.get(tid) + if task is not None: + # Claim in the same critical section as the pop: + # escalate holds `_wake` and sees COPYING. + task.transition(TaskState.COPYING) + if task is not None: + self._copy_task(task) + with self._wake: + assert task.tid not in self._blocked, f"task {task.tid} reached HOST_READY twice" + self._blocked.append(task.tid) + self._scatter_host_ready() + except BaseException: + logger.exception("omni prefix cache committer failed on task %s; releasing waiters", tid) + self._fail_task(tid) + + @torch.inference_mode() + def _copy_task(self, task: WriteTask) -> None: + """COPYING -> HOST_READY. Staging: wait the save-time device→host event. + Deferred: this is the device→host into owned `chunk.host` tensors. + """ + if task.schedule is WriteSchedule.JOIN_NEXT_STEP: + # `chunk.host` is already a staging view; device→host ran at save. + # `mark_host_ready` waits `step_d2h_event` if one was recorded. + # No per-task copy. Pool write is `_scatter_host_ready`. + task.mark_host_ready() + self._release_staged_bytes(task) + return + # One device cat + one device→host per key, not one per chunk (a + # long request has one chunk per step). + chunk_bytes = self._config.copy_chunk_bytes + by_key: dict[str, list[_WriteChunk]] = {} + for chunk in task.chunks: + for k in chunk.tensors: + by_key.setdefault(k, []).append(chunk) + pending: list[tuple[str, list[_WriteChunk], list[torch.Tensor]]] = [] + if self._copy_stream is None: + raise OmniPrefixCacheUnmatchError("deferred prefix cache copy requires a CUDA copy stream") + with torch.cuda.stream(self._copy_stream): + if task.freeze_event is not None: + self._copy_stream.wait_event(task.freeze_event) + for k, chunks in by_key.items(): + srcs = [c.tensors[k] for c in chunks] + src = srcs[0] if len(srcs) == 1 else torch.cat(srcs, dim=0) + rows = int(src.shape[0]) + row_bytes = max(1, src[:1].numel() * src.element_size()) + step = max(1, rows if rows * row_bytes <= chunk_bytes else chunk_bytes // row_bytes) + parts = [src[s : s + step].to("cpu", non_blocking=True) for s in range(0, max(rows, 1), step)] + pending.append((k, chunks, parts)) + ev = torch.cuda.Event() + ev.record() + ev.synchronize() + rows_out: list[tuple[_WriteChunk, str, torch.Tensor]] = [] + for k, chunks, parts in pending: + host = parts[0] if len(parts) == 1 else torch.cat(parts, dim=0) + if len(chunks) == 1: + rows_out.append((chunks[0], k, host)) + continue + for chunk, view in zip(chunks, torch.split(host, [int(c.tensors[k].shape[0]) for c in chunks], dim=0)): + rows_out.append((chunk, k, view)) + task.set_host_tensor(rows_out) + self._release_staged_bytes(task) + + def _fail_task(self, tid: int | None) -> None: + """-> FAILED: release waiters for a task the committer could not complete. + + No-op once the task is terminal. Budget release is idempotent, so a + raise after a successful copy stage does not uncharge twice. + """ + task = self._tasks.get(tid) if tid is not None else None + if task is None or not task.mark_failed(defer_wake=True): + return + self._release_staged_bytes(task) + with self._wake: + if tid in self._blocked: + self._blocked.remove(tid) + self._release_task_slot(task) + with self._lock: + # Publish the failure: rows behind already-published block hashes + # never landed, which the manager must raise on (hiding it would + # crash on every future hit that touches these slots). + self._failed.append(task.tid) + # Events last: a joiner woken by them must find the record on drain, + # not read never-written pool rows in a done-but-unpublished window. + task.wake_failed_waiters() + + @torch.inference_mode() + def _scatter_host_ready(self) -> None: + """HOST_READY -> WRITTEN for every `_blocked` task that has reached HOST_READY.""" + with self._wake: + ready = [tid for tid in self._blocked if (t := self._tasks.get(tid)) and t.state is TaskState.HOST_READY] + for tid in ready: + self._blocked.remove(tid) + for tid in ready: + task = self._tasks[tid] + try: + self._scatter(task) + except BaseException: + # Attribute the failure to THIS task: letting it propagate + # would fail whichever entry the worker loop happened to be + # copying, double-release its bytes, and strand this one's + # join() forever. + logger.exception("omni prefix cache scatter failed on task %s; releasing waiters", tid) + self._fail_task(tid) + + @torch.inference_mode() + def _scatter(self, task: WriteTask) -> None: + for key, slots, host in task.scatter_rows(): + self._pool.write(key, slots, host) + # `done` is set LAST: a join(done) waiter (the save barrier) must find + # the completion record on its next drain and the staging slot free — + # a done-but-undrained window would leave occupancy IN_TRANSIT and + # make the COW preserve skip the reused rows. + self._release_task_slot(task) + with self._lock: + self._completed.append(task.tid) + task.mark_done() diff --git a/vllm_omni/core/prefix_cache/group_view.py b/vllm_omni/core/prefix_cache/group_view.py new file mode 100644 index 000000000..c935f72b6 --- /dev/null +++ b/vllm_omni/core/prefix_cache/group_view.py @@ -0,0 +1,208 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""KV-cache group access for the omni prefix cache. + +The sole path through which the prefix cache touches vLLM scheduler +internals (block tables, slot mappings). Group-spec rejection happens +at kv-cache init via ``check_prefix_cache_kv_groups``; the factory only +returns None when the input batch has no block table. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch + +from vllm_omni.core.prefix_cache.interface import OmniPrefixCacheUnmatchError, PrefixCacheConfig + +if TYPE_CHECKING: + from vllm.v1.worker.gpu_input_batch import InputBatch + + +class FullAttentionGroupView: + """View over the first (full-attention) KV-cache group. + + Step slots come from the CPU block table (`step_slots_cpu`), not the + device slot_mapping. + """ + + def __init__(self, input_batch: InputBatch, block_size: int): + self._input_batch = input_batch + self.block_size = block_size + check_prefix_cache_block_layout(input_batch.block_table[0], block_size) + + def _block_table_cpu(self) -> torch.Tensor: + return self._input_batch.block_table[0].block_table.cpu + + def batch_req_ids(self) -> list[str]: + return list(self._input_batch.req_ids) + + def step_slots_cpu(self, req_ids: list[str], num_scheduled: dict[str, int]) -> torch.Tensor: + """This step's slot mapping, computed on CPU from the block table. + + The device slot_mapping would need a stream sync to read back, which + stalls the whole forward; the CPU block table carries the same + information (positions are num_computed .. +num_scheduled per request). + """ + block_table = self._block_table_cpu() + bs = self.block_size + max_blocks = int(block_table.shape[1]) + computed = self._input_batch.num_computed_tokens_cpu + parts: list[torch.Tensor] = [] + for req_id in req_ids: + n = int(num_scheduled.get(req_id, 0)) + if n <= 0: + continue + req_idx = self._input_batch.req_id_to_index[req_id] + start = int(computed[req_idx]) + pos = torch.arange(start, start + n, dtype=torch.long) + offs = pos // bs + if int(offs[-1]) >= max_blocks: + keep = offs < max_blocks + pos, offs = pos[keep], offs[keep] + if pos.numel() == 0: + continue + parts.append(block_table[req_idx, offs].to(torch.long) * bs + (pos % bs)) + return torch.cat(parts) if parts else torch.empty((0,), dtype=torch.long) + + +def check_prefix_cache_block_layout(block_table: object, block_size: int) -> None: + """Reject block-table layouts ``step_slots_cpu`` cannot mirror. + + The CPU slot math is ``table[req, pos // block_size] * block_size + + pos % block_size`` over allocator block ids. vLLM breaks that in two + cases: hybrid kernel blocks (the row holds ``blocks_per_kv_block`` + kernel ids per allocator block, e.g. FlashInfer/FlashMLA with + ``--block-size 128``) and decode context parallel (tokens are striped + across ranks and hashed at ``block_size * dcp_world_size``). + """ + if getattr(block_table, "use_hybrid_blocks", False) or int(getattr(block_table, "blocks_per_kv_block", 1)) != 1: + raise OmniPrefixCacheUnmatchError( + "omni prefix caching requires kernel_block_size == block_size; the attention backend " + f"splits each block into {getattr(block_table, 'blocks_per_kv_block', '?')} kernel blocks. " + "Pick a --block-size the backend supports natively or disable enable_prefix_caching" + ) + if int(getattr(block_table, "dcp_world_size", 1)) != 1: + raise OmniPrefixCacheUnmatchError( + "omni prefix caching does not support decode context parallel " + f"(dcp_world_size={getattr(block_table, 'dcp_world_size', '?')}); disable enable_prefix_caching" + ) + kv_bs = getattr(block_table, "kv_cache_block_size", block_size) + if int(kv_bs) != int(block_size): + raise OmniPrefixCacheUnmatchError( + f"omni prefix caching block_size {block_size} does not match the kv-cache block table ({kv_bs})" + ) + + +def check_prefix_cache_kv_groups(kv_cache_groups: object) -> None: + """Reject hybrid / multi-group models at kv-cache init, not first step. + + Only needs ``kv_cache_config.kv_cache_groups``. ``FullAttentionSpec`` + is imported here so ``tests/core`` can load this module without vllm. + """ + groups = list(kv_cache_groups or ()) + if len(groups) != 1: + raise OmniPrefixCacheUnmatchError( + "omni prefix caching requires a single full-attention kv-cache group; " + f"found {len(groups)}. disable enable_prefix_caching for this model" + ) + from vllm.v1.kv_cache_interface import FullAttentionSpec + + spec = getattr(groups[0], "kv_cache_spec", None) + if not isinstance(spec, FullAttentionSpec): + raise OmniPrefixCacheUnmatchError( + "omni prefix caching requires a single full-attention kv-cache group; " + f"found {type(spec).__name__}. disable enable_prefix_caching for this model" + ) + + +def check_prefix_cache_kv_transfer(kv_transfer_config: object) -> None: + """Reject kv_consumer / kv_both stages. + + KV loaded from a producer also shows up as ``num_computed_tokens``; the + manager would read it as a local hit with no rows behind it. + """ + if kv_transfer_config is not None and getattr(kv_transfer_config, "is_kv_consumer", False): + raise OmniPrefixCacheUnmatchError( + "omni prefix caching cannot tell locally cached tokens from KV received " + "through a KV connector; disable enable_prefix_caching on kv_consumer / " + "kv_both stages" + ) + + +def check_prefix_cache_token_accounting(cache_config: object, speculative_config: object) -> None: + """Reject configs where ``num_computed_tokens`` stops meaning "first + scheduled position, block aligned". + + Speculative decoding: with async scheduling vLLM keeps + ``num_computed_tokens_cpu`` optimistic (all drafts accepted) during the + forward and corrects it after; ``step_slots_cpu`` would read the + uncorrected value and mirror rows at the wrong slots. Refused as a + whole until that path is verified. ``prefix_match_unit`` smaller than + the block enables sub-block hits, which the hit registry rejects as + unaligned on the first hit; refuse at init instead. + """ + if speculative_config is not None: + raise OmniPrefixCacheUnmatchError( + "omni prefix caching is not supported together with speculative decoding; " + "disable enable_prefix_caching on this stage" + ) + unit = getattr(cache_config, "prefix_match_unit", None) + block_size = getattr(cache_config, "block_size", None) + if unit is not None and block_size is not None and int(unit) != int(block_size): + raise OmniPrefixCacheUnmatchError( + f"omni prefix caching requires block-aligned hits; prefix_match_unit={unit} with block_size={block_size} " + "enables sub-block hits. Unset prefix_match_unit or disable enable_prefix_caching" + ) + + +def stage_prefix_cache_config( + *, + kv_cache_config: object, + cache_config: object, + kv_transfer_config: object, + scheduler_config: object, + model_config: object, + is_pooling_model: bool, + speculative_config: object = None, +) -> PrefixCacheConfig | None: + """Runner-side gate shared by the GPU and NPU model runners. + + Returns None when the stage does not run an omni prefix cache + (``enable_prefix_caching`` off, or a pooling stage that never saves). + Otherwise refuses kv_consumer / kv_both, speculative decoding, + sub-block matching and hybrid kv groups, then sizes the config from + the scheduler. One place so a platform runner cannot silently skip a + refusal the other one has. + """ + if not getattr(cache_config, "enable_prefix_caching", False) or is_pooling_model: + return None + check_prefix_cache_kv_transfer(kv_transfer_config) + check_prefix_cache_token_accounting(cache_config, speculative_config) + check_prefix_cache_kv_groups(getattr(kv_cache_config, "kv_cache_groups", None)) + return PrefixCacheConfig.from_vllm_config( + num_blocks=kv_cache_config.num_blocks, # type: ignore[attr-defined] + block_size=cache_config.block_size, # type: ignore[attr-defined] + scheduler_config=scheduler_config, + model_config=model_config, + ) + + +def get_prefix_cache_group_view( + input_batch: InputBatch, + block_size: int, + kv_cache_groups: object = None, +) -> FullAttentionGroupView | None: + """Build the group view; None only if the batch has no block table. + + Group spec is checked first (raises). Selection is by spec, not by + counting block tables: a hybrid model's group 0 is not necessarily + full attention, and a narrower per-group table would make + step_slots_cpu silently clamp. + """ + check_prefix_cache_kv_groups(kv_cache_groups) + block_tables = getattr(input_batch.block_table, "block_tables", None) + if not block_tables: + return None + return FullAttentionGroupView(input_batch, block_size) diff --git a/vllm_omni/core/prefix_cache/interface.py b/vllm_omni/core/prefix_cache/interface.py new file mode 100644 index 000000000..274706f75 --- /dev/null +++ b/vllm_omni/core/prefix_cache/interface.py @@ -0,0 +1,171 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Interface types for the omni prefix cache. + +Naming aligns with vLLM's v1/core KV-cache design. +""" + +from collections.abc import Iterable +from dataclasses import dataclass +from enum import Enum +from typing import Any, NamedTuple, TypeAlias + +import torch + +# Four identities (still str/int at runtime, not interchangeable): +# TensorName which cached tensor (hidden / mm), not a prefix hash +# ReqId vLLM request id +# Tid WriteTask handle +# StepId save_outputs id; materialize or discard exactly once +TensorName: TypeAlias = str +ReqId: TypeAlias = str +Tid: TypeAlias = int +StepId: TypeAlias = int + +# Reserved pool key for hidden states (mm keys are flat dotted names). +# Identity only — whether a model caches it lives on ModelCachePolicy. +HIDDEN_KEY: TensorName = "__hidden_states__" + + +def is_hidden_key(key: TensorName) -> bool: + return key == HIDDEN_KEY + + +def without_hidden(keys: Iterable[TensorName]) -> set[TensorName]: + """Drop the reserved hidden identity; the rest are mm names.""" + return {k for k in keys if k != HIDDEN_KEY} + + +class WriteSchedule(Enum): + """Write scheduling policy for one WriteTask.""" + + # Immediately-cached keys: device→host launched at save into the + # staging pool; committer waits that event and writes the CPU pool. + # The next save waits done (a reused slot must be pool-complete). + JOIN_NEXT_STEP = "join_next_step" + # Deferred mm: stays on the GPU clone until finish/abort (GPU-byte + # budget may force a copy earlier). One WriteTask per request. + JOIN_ON_FINISH = "join_on_finish" + + +@dataclass(frozen=True) +class PrefixCacheConfig: + """Sizing and flow-control knobs (mirrors KVCacheConfig).""" + + num_blocks: int + block_size: int + # GPU-clone byte budget for JOIN_ON_FINISH; exceeding it forces a copy. + gpu_staging_bytes: int = 512 * 1024 * 1024 + # Device→host staging: circular slots, one whole step each (not per request). + # Host memory per key ≈ staging_depth * staging_capacity_tokens * width * dtype. + # Prefer from_vllm_config so staging_capacity_tokens tracks max_num_batched_tokens. + staging_depth: int = 4 + staging_capacity_tokens: int = 1024 + # How long save waits for a free staging slot, and join/join_host_ready + # for the committer; exceeding it raises. + staging_claim_timeout_s: float = 30.0 + # Device→host chunk size for JOIN_ON_FINISH (copied a piece at a time). + copy_chunk_bytes: int = 16 * 1024 * 1024 + + @classmethod + def from_vllm_config( + cls, + *, + num_blocks: int, + block_size: int, + scheduler_config: Any = None, + model_config: Any = None, + ) -> "PrefixCacheConfig": + """Size device→host staging from the running scheduler. + + A slot holds one *step* (the whole batch), not one request: + ``staging_capacity_tokens`` is ``max_num_batched_tokens`` (falls back + to ``max_model_len``); ``staging_depth`` is how many unconsumed + step ids may exist at once, not ``max_num_seqs``. + + Pinned staging is allocated lazily per key at + ``depth * capacity_tokens * width * dtype``. There is no clamp: a + step larger than capacity raises. A 16k-token thinking batch at + hidden=2048 bf16 is ~256 MiB for hidden alone; each mm key adds + another pool tensor. + + ``staging_depth`` is the dataclass default (4). There is no CLI or + deploy YAML knob — changing it is a code change. Every save that + issues a step id claims one slot, including saves with only leftover + mm. A full pool waits for materialize/discard; + ``staging_claim_timeout_s`` then errors. + """ + batched = getattr(scheduler_config, "max_num_batched_tokens", None) + model_len = getattr(scheduler_config, "max_model_len", None) + if not model_len and model_config is not None: + model_len = getattr(model_config, "max_model_len", None) + try: + capacity = int(batched or model_len or 1024) + except (TypeError, ValueError): + capacity = 1024 + return cls( + num_blocks=num_blocks, + block_size=block_size, + staging_capacity_tokens=max(1, capacity), + ) + + +class StageCacheOutputs(NamedTuple): + """Plain value object: per-request merged stage outputs.""" + + # req -> full-prompt hidden states (None when policy skips them) + hidden_states: dict[ReqId, torch.Tensor] | None + # tensor name -> req -> payload element + mm_outputs: dict[TensorName, dict[ReqId, Any]] + + +class OmniPrefixCacheUnmatchError(RuntimeError): + """Contract, config, or KV-occupancy error that must raise. + + Includes hit spans that resolve to absent slots (omni cache diverged + from vLLM KV), a step id consumed twice or never saved, a step larger + than the staging page, and a failed write. Do not pretend these + were a miss. + """ + + +class OmniPrefixCacheStagingTimeoutError(OmniPrefixCacheUnmatchError): + """Save waited for a free staging slot and timed out.""" + + +@dataclass(frozen=True) +class ModelCachePolicy: + """Replaces getattr probing on models for cache behavior decisions. + + Hidden's *name* is ``HIDDEN_KEY`` (shared identity). This object + answers whether this model caches it, and which mm keys are deferred. + """ + + needs_full_hidden_states: bool = True + # Mm whose first dim is this step's token count; stays on the GPU + # clone until finish/abort (JOIN_ON_FINISH). Also skipped by the + # immediate on-device clone / device→host path. + deferred_keys: frozenset[TensorName] = frozenset() + + @property + def hidden_key(self) -> TensorName | None: + """Pool key for hidden, or None when this model opts out.""" + return HIDDEN_KEY if self.needs_full_hidden_states else None + + def get_hit_keys(self, keys: Iterable[TensorName]) -> list[TensorName]: + """Keys to plan/prefetch for a hit: hidden first (if cached), then mm.""" + mm = sorted(without_hidden(keys)) + return [HIDDEN_KEY, *mm] if self.needs_full_hidden_states else mm + + def skip_immediate_mm(self, key: TensorName) -> bool: + """Immediate on-device clone must not take hidden or deferred keys from mm.""" + return is_hidden_key(key) or key in self.deferred_keys + + @classmethod + def from_model(cls, model: Any) -> "ModelCachePolicy": + """Shim over legacy per-model attributes (deprecation window).""" + deferred = frozenset(getattr(model, "deferred_prefix_cache_mm_keys", ()) or ()) + return cls( + needs_full_hidden_states=bool(getattr(model, "requires_full_prefix_cached_hidden_states", True)), + deferred_keys=deferred, + ) diff --git a/vllm_omni/core/prefix_cache/manager.py b/vllm_omni/core/prefix_cache/manager.py new file mode 100644 index 000000000..5294a25d4 --- /dev/null +++ b/vllm_omni/core/prefix_cache/manager.py @@ -0,0 +1,1475 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Omni prefix cache, manager side. + +Owns slot occupancy, the request-task table, the hit/span registry, +per-step snapshots, and the merge. The controller owns the staging +pool, copy queues, and writing rows into the CPU block pool. The +state lock covers those tables only — never a wait-for-copy, a +GPU-byte-budget flush, or a memcpy. + +Helper docstrings mark ``_state_lock`` (non-reentrant): + + Caller holds already inside a critical section; do not acquire + Takes acquires here (``@_locked`` or ``with``) + +Two host stores: + + StagingBufferPool reusable step-sized pages. save copies this + step's immediately-cached keys (hidden + + non-deferred mm) device→host into one page. + Per-task `chunk.host` is a view into that page, + not a second copy. + PrefixBlockPool durable (kv_slot, key) prefix cache. The + committer only writes into it. + +Two write paths (which keys, not how many tokens): + + JOIN_NEXT_STEP immediately-cached keys. Device→host is already + in flight at submit; the committer waits + `step_d2h_event` then copies host→pool. The next + save waits `done`: reused slots must not leave a + pending pool write behind. + JOIN_ON_FINISH deferred mm. Stays on the device clone; the + committer does that device→host, then writes the + pool. Forced onto the high-priority queue on + finish/abort or GPU-byte-budget pressure. + +Per real scheduler_output, engine-thread order: + + new_step_starts before _update_states drops finished requests + (register hits, start prefix prefetch) + forward + save_outputs clone off live buffers + launch staging copy; + returns step id + materialize or discard_step exactly one of the two, once + +materialize may run on the async output builder while the engine is +already in the next step. Warmup/dummy runs are never fed. +""" + +from __future__ import annotations + +import logging +import threading +from collections import deque +from collections.abc import Iterable, Mapping +from concurrent.futures import Future, ThreadPoolExecutor +from dataclasses import dataclass, field +from enum import IntEnum +from typing import TYPE_CHECKING, Any, NamedTuple, NoReturn + +import torch + +from vllm_omni.core.prefix_cache.block_pool import PrefixBlockPool +from vllm_omni.core.prefix_cache.controller import ( + OmniPrefixCacheController, + StagingBufferHolder, + StepD2HClaim, + WriteTask, + _BudgetTicket, + _WriteChunk, +) +from vllm_omni.core.prefix_cache.interface import ( + ModelCachePolicy, + OmniPrefixCacheStagingTimeoutError, + OmniPrefixCacheUnmatchError, + PrefixCacheConfig, + ReqId, + StageCacheOutputs, + StepId, + TensorName, + Tid, + WriteSchedule, + is_hidden_key, + without_hidden, +) + +if TYPE_CHECKING: + from vllm.v1.core.sched.output import SchedulerOutput + + from vllm_omni.core.prefix_cache.group_view import FullAttentionGroupView + +logger = logging.getLogger(__name__) + + +class _Occupancy(IntEnum): + ABSENT = 0 + IN_TRANSIT = 1 + COMMITTED = 2 + + +def _is_step_token_tensor(val: Any, n: int, padded: int) -> bool: + """2D+ tensor whose first dim is this step's token count (``n`` or padded). + + True means callers may take ``val[:n]``. Leftover tensors (``codes.ref``), + lists, and other shapes are False. + """ + return isinstance(val, torch.Tensor) and val.ndim >= 2 and int(val.shape[0]) in (n, padded) + + +def _raise_unreadable_hit(req_id: str, key: str, why: str) -> NoReturn: + raise OmniPrefixCacheUnmatchError(f"hit span for req {req_id} key={key} is not readable ({why})") + + +def _snapshot_leftover_mm_cpu( + mm_outputs: dict[str, Any], + device_snapshot_keys: set[str], + num_tokens_unpadded: int, + num_tokens_padded: int | None = None, +) -> tuple[dict[str, Any], torch.cuda.Event | None]: + """CPU copy of mm that did not land on the staging page. + + Skip ``device_snapshot_keys`` (those already have a device→host page). + Copy the rest — deferred mm, lists, ``codes.ref`` — so materialize can + run after the next forward overwrites graph buffers. Slice ``[:n]`` + only when ``shape[0] == n``; ``>= n`` would clip ``codes.ref``. + + CUDA tensors land in pinned memory through a non-blocking copy on the + current stream: stream order keeps it ahead of the next forward, and + the engine thread does not wait for this forward to finish. The + returned event (None when nothing was on CUDA) must be synchronized + before the snapshot is read. + """ + n = num_tokens_unpadded + padded = n if num_tokens_padded is None else int(num_tokens_padded) + on_cuda = False + + def _copy(val: Any) -> Any: + nonlocal on_cuda + if isinstance(val, torch.Tensor): + t = val[:n] if _is_step_token_tensor(val, n, padded) and int(val.shape[0]) == n else val + t = t.detach() + if t.is_cuda: + host = torch.empty(t.shape, dtype=t.dtype, device="cpu", pin_memory=True) + host.copy_(t, non_blocking=True) + on_cuda = True + return host + # .cpu() copies other device tensors; CPU/pinned views still share storage. + copied = t.cpu() if t.device.type != "cpu" else t.clone() + return copied.contiguous() + if isinstance(val, Mapping): + return {k: _copy(v) for k, v in val.items()} + if isinstance(val, list): + return [_copy(v) for v in val] + if isinstance(val, tuple): + return tuple(_copy(v) for v in val) + return val + + leftover = { + key: _copy(val) for key, val in mm_outputs.items() if key not in device_snapshot_keys and not is_hidden_key(key) + } + event = None + if on_cuda: + event = torch.cuda.Event() + event.record() + return leftover, event + + +def _unpin_leftover(val: Any) -> Any: + """Clone pinned tensors out so the payload does not hold pinned pages.""" + if isinstance(val, torch.Tensor): + return val.clone() if val.is_pinned() else val + if isinstance(val, Mapping): + return {k: _unpin_leftover(v) for k, v in val.items()} + if isinstance(val, list): + return [_unpin_leftover(v) for v in val] + if isinstance(val, tuple): + return tuple(_unpin_leftover(v) for v in val) + return val + + +@dataclass +class _StepOutputs: + """This step's outputs split by consumer. + + A tensor whose first dim equals this step's token count (or the + CUDA-graph padded count) can be sliced ``[:n]`` into per-token rows. + ``codes.ref`` and lists are not that shape. + + ``immediate``: those per-token rows copied device→host this step + (hidden + non-deferred mm). ``deferred_chunks``: per-token deferred + mm, packed per request for JOIN_ON_FINISH. ``leftover``: CPU replica + for this step's materialize of everything that did not get a staging + page. + A deferred per-token key is in both ``deferred_chunks`` (later cache + write) and ``leftover`` (this-step read) — two consumers, not a + duplicate store. + + ``immediate_budget`` charges the immediate clones once; every + JOIN_NEXT_STEP task of the step pins it. Deferred chunks carry their + own shared ticket. + """ + + immediate: dict[str, torch.Tensor] + deferred_chunks: list[tuple[str, _WriteChunk]] + leftover: dict[str, Any] + # Completion of the leftover device→host copies; None when none ran on CUDA. + leftover_event: torch.cuda.Event | None = None + immediate_budget: _BudgetTicket | None = None + # Pool storage allocated (unlocked) for keys first seen this step; + # published into the pool / occupancy tables under the state lock. + new_key_storage: dict[str, torch.Tensor] = field(default_factory=dict) + + @property + def deferred_budget(self) -> _BudgetTicket | None: + return self.deferred_chunks[0][1].budget if self.deferred_chunks else None + + def freeze_targets(self) -> list[torch.Tensor]: + """Device clones the freeze event must cover.""" + return list(self.immediate.values()) + [t for _, c in self.deferred_chunks for t in c.tensors.values()] + + def budget_bytes(self) -> int: + return sum(t.nbytes for t in (self.immediate_budget, self.deferred_budget) if t is not None) + + +def _locked(fn): + """Serialize public entry points: the async output builder calls + materialize() while the engine thread is in the next step.""" + + def wrapper(self, *args, **kwargs): + with self._state_lock: + return fn(self, *args, **kwargs) + + wrapper.__name__ = fn.__name__ + wrapper.__doc__ = fn.__doc__ + return wrapper + + +@dataclass +class _SlotRef: + """Where one (req, key) span's slots live: planned under the lock, fetched outside it. + + Schedule split, not a single read tier: + - JOIN_NEXT_STEP in-transit: ``join_tids``. Fetch waits ``done``, + drains, then reads the pool. Staging views are never sliced. + - JOIN_ON_FINISH in-transit: ``staged_list`` task refs. Fetch uses + ``fetch_host`` (device freeze / committer host). + - Already in the CPU pool: ``already_staged`` → pool. + + A JOIN_NEXT_STEP task may disappear between plan and join (another + entry already published it into the pool); ``join`` no-ops and the + pool rows persist. + """ + + slots: torch.Tensor # KV slot ids (PrefixBlockPool rows) + key: TensorName + req_id: ReqId + already_staged: bool # this key is already in the CPU pool + staged_list: list[tuple[WriteTask, torch.Tensor]] # JOIN_ON_FINISH only + join_tids: list[Tid] = field(default_factory=list) # JOIN_NEXT_STEP in-transit + # Slot version snapshot at plan time, checked on fetch. + reserved_version: torch.Tensor | None = None + # Copy-on-write rescue: slot's committed rows for pending reads. + preserved: dict[int, torch.Tensor] = field(default_factory=dict) + + +@dataclass(kw_only=True) +class _StepContext: + """Save-time snapshot of one step; consumed exactly once. + + Built on the engine thread in save_outputs so materialize (possibly on + the async builder) never reads the live batch. Host rows live on `d2h` + (empty views when this save only had leftover mm). materialize or + discard_step pops it and frees the staging slot. A later save waits + if all slots are still held. + """ + + # Packed layout in batch order: req -> [start, end) of this step's rows. + spans: dict[ReqId, tuple[int, int]] + num_tokens_unpadded: int = 0 + + # Hits snapshotted at new_step_starts. Prefetch fills [hit | empty tail] + # during forward; materialize writes the tail. + hits: dict[ReqId, tuple[int, list[int]]] # (hit_upto, blocks) + hit_prefetch: dict[ReqId, dict[TensorName, Future]] = field(default_factory=dict) + + # Key split frozen at save: recompute at materialize races ensure_key. + cached_keys: set[TensorName] = field(default_factory=set) + # Leftover mm copied to CPU at save (this-step deferred rows + mm + # that is not written to the pool). + mm_cpu_snapshot: dict[TensorName, Any] = field(default_factory=dict) + # Wait this before reading mm_cpu_snapshot (None when nothing was on CUDA). + mm_cpu_snapshot_event: torch.cuda.Event | None = None + + # Staging slot for this step id (empty views when only leftover mm). + d2h: StepD2HClaim | None = None + + +class _SlotStatus(NamedTuple): + """Occupancy row for one tensor name (views into the table, not copies).""" + + state: torch.Tensor # int8[num_slots] + tids: torch.Tensor # Tid per kv slot; 0 = none + slot_version: torch.Tensor # int64[num_slots]; bumped each time a new write claims the slot + + +class _SlotStatusTable: + """Per (KV slot, tensor name): empty, being written, or already in the pool. + + Hidden and a deferred mm field on the same slot are independent. + ``map_slots`` marks a write in progress; if another task still owns + the slot, the manager records those rows as no longer owned by it. + ``commit`` runs after the pool write: still-owned slots become + committed. + """ + + def __init__(self, num_slots: int) -> None: + self.num_slots = num_slots + self.state: dict[TensorName, torch.Tensor] = {} # int8[num_slots] + self.tids: dict[TensorName, torch.Tensor] = {} # Tid per kv slot; 0 = none + self.slot_versions: dict[TensorName, torch.Tensor] = {} # int64[num_slots]; bumped on each new write claim + self.task_slots: dict[Tid, torch.Tensor] = {} # kv slots + self.task_keys: dict[Tid, tuple[TensorName, ...]] = {} + + def init_table(self, key: TensorName) -> None: + """Allocate occupancy tensors for ``key`` if they do not exist.""" + if key in self.state: + return + self.state[key] = torch.zeros(self.num_slots, dtype=torch.int8) + self.tids[key] = torch.zeros(self.num_slots, dtype=torch.int64) + self.slot_versions[key] = torch.zeros(self.num_slots, dtype=torch.int64) + + def get_slot_status(self, key: TensorName) -> _SlotStatus: + """Occupancy tensors for ``key``. ``init_table`` must have run.""" + return _SlotStatus(state=self.state[key], tids=self.tids[key], slot_version=self.slot_versions[key]) + + def map_slots( + self, slots: torch.Tensor, tid: Tid, keys: Iterable[TensorName] + ) -> list[tuple[Tid, TensorName, torch.Tensor]]: + """Record ``tid`` on these (slot, key). Return in-transit rows + another write still owned (caller marks them skipped on it).""" + keys = tuple(keys) + stolen: list[tuple[Tid, TensorName, torch.Tensor]] = [] + for key in keys: + status = self.get_slot_status(key) + cur = status.tids[slots] + stale = (status.state[slots] == _Occupancy.IN_TRANSIT) & (cur != tid) & (cur != 0) + if bool(stale.any()): + for old in {int(o) for o in cur[stale].tolist()}: + stolen.append((old, key, slots[stale & (cur == old)])) + status.state[slots] = _Occupancy.IN_TRANSIT + status.tids[slots] = tid + # Bump the version so a reader that captured an older one detects + # this claim even after it later commits (COMMITTED, not IN_TRANSIT). + status.slot_version[slots] += 1 + prev = self.task_slots.get(tid) + if prev is None: + self.task_slots[tid] = slots + self.task_keys[tid] = keys + else: + # Deferred tasks grow one `_WriteChunk` per step. + self.task_slots[tid] = torch.cat([prev, slots]) + self.task_keys[tid] = tuple(dict.fromkeys(self.task_keys[tid] + keys)) + return stolen + + def commit(self, tids: Iterable[Tid]) -> None: + """Flip still-owned slots to COMMITTED and drop the reverse index.""" + for tid in tids: + slots = self.task_slots.pop(tid, None) + keys = self.task_keys.pop(tid, ()) + if slots is None: + continue + for key in keys: + status = self.get_slot_status(key) + still_ours = status.tids[slots] == tid + idx = slots[still_ours] + status.state[idx] = _Occupancy.COMMITTED + status.tids[idx] = 0 + + +class _RequestTaskTable: + """Per request: still live, which WriteTasks it opened, deferred task. + + Also allocates ``tid`` and increments per-request ``write_n``. + Device→host copy and pool write stay on the controller. + """ + + def __init__(self) -> None: + self._next_tid: Tid = 1 + self.write_n: dict[ReqId, int] = {} # last write_n issued + self.tasks: dict[ReqId, set[Tid]] = {} + self.deferred: dict[ReqId, WriteTask] = {} + self.live_reqs: set[ReqId] = set() + + def alloc_tid(self) -> Tid: + tid = self._next_tid + self._next_tid += 1 + return tid + + def increment_write_n(self, req_id: ReqId) -> int: + n = self.write_n.get(req_id, 0) + 1 + self.write_n[req_id] = n + return n + + def track(self, req_id: ReqId, tid: Tid) -> None: + self.tasks.setdefault(req_id, set()).add(tid) + + def finish(self, req_id: ReqId) -> tuple[set[Tid], WriteTask | None]: + """Drop this request's rows. Returns owned tids + deferred task.""" + self.live_reqs.discard(req_id) + self.write_n.pop(req_id, None) + tids = self.tasks.pop(req_id, set()) + dtask = self.deferred.pop(req_id, None) + return tids, dtask + + def drop_completed(self, tids: Iterable[Tid]) -> None: + done = set(tids) + if not done: + return + for req_tids in self.tasks.values(): + req_tids -= done + + +class OmniPrefixCacheManager: + def __init__( + self, + config: PrefixCacheConfig, + view: FullAttentionGroupView, + *, + eager: bool | None = None, + ): + self._config = config + self._view = view + self._pool = PrefixBlockPool(config) + self._controller = OmniPrefixCacheController(self._pool, config, eager=eager) + self._policy = ModelCachePolicy() + # Serializes engine vs async-builder public entries. Non-reentrant: + # those entries never call each other, and the lock must not cover + # a wait-for-copy, GPU-byte-budget flush, or device→host copy. + self._state_lock = threading.Lock() + + self._slot_status = _SlotStatusTable(config.num_blocks * config.block_size) + # Hidden is known from the default policy; mm rows at init_table. + if (hk := self._policy.hidden_key) is not None: + self._slot_status.init_table(hk) + self._request_tasks = _RequestTaskTable() + # Join worklists — not occupancy, not the request-task table. + self._join_next_step_tids: list[Tid] = [] + self._join_finished_tids: set[Tid] = set() # escalated on finish/abort + + # This step's hits (copied into _StepContext at save). + self._cur_num_scheduled: dict[ReqId, int] = {} + self._hit_spans: dict[ReqId, tuple[int, list[int]]] = {} # (upto, blocks) + self._hit_prefetch: dict[ReqId, dict[TensorName, Future]] = {} + + # Planned-but-unread slot refs. A write reusing their slots copies the + # old rows into ``_SlotRef.preserved`` before overwriting (COW). Guarded + # by ``_state_lock``; each ref is dropped when its fetch completes. + self._pending_reads: list[_SlotRef] = [] + # Sticky fatal set on the first drained write failure; every later + # facade entry re-raises it (see _commit_drained_writes). + self._fatal_write_failure: str | None = None + + # Prefix gather during forward (CPU work releases the GIL). + # One worker: complete in submit order; pop finished work from the head. + self._prefetch_pool = ThreadPoolExecutor(1, thread_name_prefix="omni-prefix-cache-prefetch") + self._prefetch_queue: deque[tuple[Future, _SlotRef]] = deque() + + # One snapshot per step id; consume with materialize or discard_step. + self._next_step_id: StepId = 1 + self._step_ctxs: dict[StepId, _StepContext] = {} + + # ------------------------------------------------------ public entries + + def register_policy(self, policy: ModelCachePolicy) -> None: + self._policy = policy + if (hk := policy.hidden_key) is not None: + self._slot_status.init_table(hk) + + @torch.inference_mode() + def new_step_starts(self, scheduler_output: SchedulerOutput) -> None: + """Handle one scheduler_output. + + Engine thread only; before _update_states removes finished + requests; exactly once per real step. Registers new-request prefix + hits (copying their block tables) and forces finished/aborted + requests' still-open writes onto the high-priority copy queue — + a block hash that entered the batch must land in the cache, + abort included. ``escalate`` (eager: the copy + pool write) runs + after ``_state_lock`` is released. + """ + to_escalate: list[int] = [] + with self._state_lock: + # 1. Publish writes the committer has already written into the pool. + self._commit_drained_writes() + + # 2. Finished/aborted reqs: collect still-open writes. Abort too: + # those block hashes are already in vLLM; dropping the write + # would leave future hits ABSENT. The next save waits + # join_host_ready. (Not leftover_mm — those are this-step reads.) + finished = getattr(scheduler_output, "finished_req_ids", None) or () + for req_id in finished: + tids, dtask = self._request_tasks.finish(req_id) + if dtask is not None: + tids.add(dtask.tid) + pending_tasks = [tid for tid in tids if self._controller.get_task(tid) is not None] + if pending_tasks: + to_escalate.extend(pending_tasks) + self._join_finished_tids.update(pending_tasks) + + # 3. Copy this arrival's prefix-hit block ids. scheduled_new_reqs + # is the only place they appear; after _update_states they sit + # on the live request and grow as decode allocates more blocks. + # materialize (async builder) must not reread that live table. + self._clear_hit_infos() + for new_req in getattr(scheduler_output, "scheduled_new_reqs", ()) or (): + req_id = new_req.req_id + if req_id in self._request_tasks.live_reqs: + # Already live: async_chunk continuation, or a V2-runner + # resume after preemption (V1 resumes via + # scheduled_cached_reqs). Either way num_computed_tokens is + # not mirrored as a hit span; see the design doc. + continue + self._request_tasks.live_reqs.add(req_id) + num_computed = int(getattr(new_req, "num_computed_tokens", 0) or 0) + if num_computed > 0: + # block_ids is per-kv-group; group 0 only. + blocks = getattr(new_req, "block_ids", None) + if blocks is not None and len(blocks) > 0 and not isinstance(blocks[0], int): + blocks = blocks[0] + if not blocks: + # Fail at the cause: a hit we cannot snapshot now would + # crash at materialize time with less context (materialize is + # forbidden from reading the live batch). + raise OmniPrefixCacheUnmatchError( + f"prefix hit for req {req_id} ({num_computed} tokens) carries no block_ids" + ) + bs = self._config.block_size + if num_computed % bs != 0: + raise OmniPrefixCacheUnmatchError( + f"prefix hit not block aligned (req={req_id}, hit_upto={num_computed}, block_size={bs})" + ) + hit_blocks = list(blocks[: num_computed // bs]) + self._hit_spans[req_id] = (num_computed, hit_blocks) + + # 4. Gather those spans on the prefetch thread; overlaps this forward. + while self._prefetch_queue and self._prefetch_queue[0][0].done(): + self._prefetch_queue.popleft() + self._cur_num_scheduled = dict(scheduler_output.num_scheduled_tokens) + if self._hit_spans: + self._prefetch_hit_spans() + if to_escalate: + self._controller.escalate(to_escalate) + + @torch.inference_mode() + def save_outputs( + self, + hidden_states: torch.Tensor | None, + mm_outputs: dict[str, Any] | None, + *, + num_tokens_unpadded: int, + num_tokens_padded: int, + ) -> int: + """Write this step's outputs into the cache; returns the step id. + + Engine thread only, after the forward and before materialize. + Immediately-cached rows: one on-device clone, one whole-step + device→host into the staging pool, then one JOIN_NEXT_STEP + WriteTask per request whose `chunk.host` is a view of that page. + Deferred rows stay on the device clone (JOIN_ON_FINISH); the + committer copies them later. Leftover mm (this-step deferred rows + + mm not written to the pool) is copied to CPU here so materialize + never reads live graph buffers. + Snapshots everything materialize needs. The returned step id MUST + be consumed exactly once — by materialize() or discard_step(). + Every step id claims one staging slot (saves with only leftover mm + included); a later save waits for a free slot and times out if + none return. + + The state lock never covers a blocking wait: the previous step's + JOIN_NEXT_STEP wait, the clone build, the GPU-byte-budget reserve + (which may flush), and the staging-slot claim all run unlocked. + """ + # 1. Join the previous step's host copies (unlocked). + self._wait_for_host_ready() + + # 2. Packed batch layout for this step (req -> [start, end)). + req_order = self._view.batch_req_ids() + num_sched = {r: int(self._cur_num_scheduled.get(r, 0)) for r in req_order} + query_start: dict[str, int] = {} + current_start_idx = 0 + for req_id in req_order: + query_start[req_id] = current_start_idx + current_start_idx += num_sched[req_id] + + slots_cpu: torch.Tensor | None = None + mm_outputs = mm_outputs or {} + freeze_event = None + + # 3. Slot map, then split into immediate / deferred / leftover. + if num_tokens_unpadded > 0: + # Derive the slot mapping on CPU: reading the device one back + # would need a stream sync that waits on the whole forward. + slots_cpu = self._view.step_slots_cpu(req_order, num_sched) + if int(slots_cpu.numel()) != num_tokens_unpadded: + # Fail at the cause: skipping the save would leave rows absent + # behind hashes vLLM already published — a delayed crash at + # some future hit instead of a debuggable one here. + raise OmniPrefixCacheUnmatchError( + f"slot mapping covers {int(slots_cpu.numel())} of {num_tokens_unpadded} scheduled tokens; " + "CPU-side slot derivation out of sync with the batch" + ) + step_outputs = self._split_step_outputs( + hidden_states, + mm_outputs, + num_tokens_unpadded, + num_tokens_padded, + slots_cpu=slots_cpu, + req_order=req_order, + num_sched=num_sched, + query_start=query_start, + ) + + # 4. Freeze the device clones and reserve the GPU-byte budget (unlocked). + freezed_tensors = step_outputs.freeze_targets() + if freezed_tensors: + if torch.cuda.is_available() and any(t.is_cuda for t in freezed_tensors): + freeze_event = torch.cuda.Event() + freeze_event.record() + # One ticket per clone (immediate step clone, shared deferred + # clone); per-request slices are views and charge nothing. + # Reserve may block on a flush: outside the lock. The flush may + # close a deferred task this step appends to; _stage_deferred + # then opens a fresh one, so a long request cannot pin the + # whole budget. + self._controller.reserve(step_outputs.budget_bytes()) + + # 5. Claim a staging slot (unlocked), optional device→host into it, + # register the writes + store the step snapshot (locked), then + # dispatch the queued writes (unlocked: in eager mode dispatch + # is the copy + pool write itself). Saves with only leftover mm + # still claim. Full pool waits; timeout lists unused step ids. + d2h_claim: StepD2HClaim | None = None + step_holder = StagingBufferHolder.for_step(self._next_step_id) + transferred = False + bound_tids: list[int] = [] + try: + d2h_claim = self._stage_step_host(step_outputs.immediate, num_tokens_unpadded, freeze_event, step_holder) + + step_id, queued = self._publish_saved_step( + req_order=req_order, + query_start=query_start, + num_sched=num_sched, + num_tokens_unpadded=num_tokens_unpadded, + step_outputs=step_outputs, + slots_cpu=slots_cpu, + mm_keys=set(mm_outputs.keys()), + freeze_event=freeze_event, + d2h_claim=d2h_claim, + bound_tids=bound_tids, + ) + self._controller.dispatch(queued) + transferred = True + return step_id + finally: + # Slot claim is outside the lock; a later raise must release + # the step and any task that already bound this slot. + if not transferred and d2h_claim is not None: + # A dispatch failure has no caller that can consume this step. + with self._state_lock: + self._step_ctxs.pop(step_holder.owner_id, None) + self._release_staging_on_failed_save(d2h_claim.staging_slot, step_holder, bound_tids) + + @torch.inference_mode() + def materialize(self, step_id: int, req_ids: list[str]) -> StageCacheOutputs: + """Per-request merged outputs for the step saved as `step_id`. + + Any thread. `req_ids` must be (a subset of) the save-time snapshot; + an outside id means the caller is reading the live batch. + A request without a hit is a plain miss and gets exactly + this step's rows — normal path, nothing logged. A hit span that + resolves to absent rows raises OmniPrefixCacheUnmatchError: fatal + by contract (do not pretend it was a miss). + + Two phases: under the lock, publish finished writes and pin every + row source (task refs + masks, absent checks included) — not yet + reading the tensors. Unlocked: wait this step's `step_d2h_event`, + clone the staging views (then drop the step holder), and merge. + The engine thread never waits on this thread's device→host copy. + """ + ctx = None + step_released = False + try: + with self._state_lock: + ctx = self._take_step_ctx(step_id) + self._commit_drained_writes() + + # The builder must pass (a subset of) the req list captured at + # save time — an id outside the snapshot means it is reading the + # live batch, which the contract forbids. Not a fallback that + # serves a miss. + unknown = set(req_ids) - set(ctx.spans) + if unknown: + raise OmniPrefixCacheUnmatchError( + f"materialize(step {step_id}) got req ids outside the save snapshot: {sorted(unknown)[:8]}" + ) + + cached_keys = ctx.cached_keys + + hit_sources: dict[tuple[str, str], _SlotRef | Future] = {} + for req_id in req_ids: + hit = ctx.hits.get(req_id) + if not hit: + continue + hit_upto, hit_blocks = hit + prefetched = ctx.hit_prefetch.get(req_id, {}) + slots = self._get_hit_slots(hit_upto, hit_blocks) + keys = self._policy.get_hit_keys(cached_keys) + for key in keys: + fut = prefetched.get(key) + if fut is not None: + hit_sources[(req_id, key)] = fut + continue + hit_sources[(req_id, key)] = self._slot_ref(slots, key, req_id) + + # ---- unlocked: data movement + merge ---- + if ctx.mm_cpu_snapshot_event is not None: + ctx.mm_cpu_snapshot_event.synchronize() + ctx.mm_cpu_snapshot = _unpin_leftover(ctx.mm_cpu_snapshot) + current: dict[str, torch.Tensor] = {} + if ctx.d2h is not None: + # Whole-step device→host was launched at save. One event wait + # (usually already complete), then a contiguous copy-out per + # key so consumers no longer depend on the reusable slot. + if ctx.d2h.event is not None: + ctx.d2h.event.synchronize() + current = {k: v.clone() for k, v in ctx.d2h.views.items()} + self._release_step_staging(ctx, step_id) + step_released = True + + hidden_out: dict[str, torch.Tensor] | None = None + hidden_key = self._policy.hidden_key + if hidden_key is not None and hidden_key in current: + hidden_out = {} + for req_id in req_ids: + hidden_out[req_id] = self._merge_cached_for_req( + ctx, req_id, hidden_key, current[hidden_key], hit_sources + ) + + mm_out: dict[str, dict[str, Any]] = {} + for key in cached_keys: + cur = current.get(key) + if cur is None: + val = ctx.mm_cpu_snapshot.get(key) + if not isinstance(val, torch.Tensor): + continue + # Leftover snapshot; spans stay within [0, n), no re-slice. + cur = val + mm_out[key] = { + req_id: self._merge_cached_for_req(ctx, req_id, key, cur, hit_sources) for req_id in req_ids + } + + self._merge_uncached_mm(ctx, req_ids, cached_keys, mm_out) + return StageCacheOutputs(hidden_states=hidden_out, mm_outputs=mm_out) + finally: + if ctx is not None and not step_released: + self._release_step_staging(ctx, step_id) + + @_locked + def discard_step(self, step_id: int) -> None: + """Consume the step context when nothing will materialize it. + + Any thread; same exactly-once contract as materialize (unknown or + duplicate id raises). Only the read-side snapshot is dropped — + the cache write proceeds unchanged. + """ + ctx = self._take_step_ctx(step_id) + self._release_step_staging(ctx, step_id) + + def shutdown(self) -> None: + self._prefetch_pool.shutdown(wait=False, cancel_futures=True) + self._prefetch_queue.clear() + self._controller.shutdown() + + # ------------------------------------------------------ new_step + + def _clear_hit_infos(self) -> None: + """Drop the live hit / prefetch tables. Caller holds ``_state_lock``.""" + self._hit_spans.clear() + self._hit_prefetch.clear() + + def _prefetch_hit_spans(self) -> None: + """Caller holds ``_state_lock``. Plan each hit span not yet planned + and gather it on the prefetch thread. + + Called twice per step: at new_step_starts (overlaps the forward) + and again at publish for spans that could not plan then — same-step + hits, whose rows this step's save has only now registered. Starting + the gather at publish, not at materialize, means it runs before the + next step can hand those blocks to a new tenant; a span that still + fails to plan is left to materialize, which raises if unread. + """ + keys = self._policy.get_hit_keys(self._pool.keys()) + for req_id, (hit_upto, hit_blocks) in self._hit_spans.items(): + futs = self._hit_prefetch.setdefault(req_id, {}) + if all(key in futs for key in keys): + continue + n_new = int(self._cur_num_scheduled.get(req_id, 0)) + slots = self._get_hit_slots(hit_upto, hit_blocks) + for key in keys: + if key in futs: + continue + try: + src = self._slot_ref(slots, key, req_id) + except OmniPrefixCacheUnmatchError: + continue + fut = self._prefetch_pool.submit(self._prefetch_hit, src, n_new) + self._prefetch_queue.append((fut, src)) + futs[key] = fut + if not futs: + del self._hit_prefetch[req_id] + + @torch.inference_mode() + def _prefetch_hit(self, src: _SlotRef, n_new: int) -> torch.Tensor: + """Prefetch thread: gather the hit span and pre-build the merged + buffer with the prefix filled. materialize writes only this step's + rows at the tail — the gather AND the prefix copy both happen while + the forward runs, and the cat leaves the critical path.""" + rows = self._fetch_source(src) + out = torch.empty((rows.shape[0] + n_new, rows.shape[-1]), dtype=rows.dtype) + out[: rows.shape[0]] = rows + return out + + # ---------------------------------------------------------- save + + def _stage_step_host( + self, + device_snapshot: dict[str, torch.Tensor], + num_tokens_unpadded: int, + freeze_event: torch.cuda.Event | None, + step_holder: StagingBufferHolder, + ) -> StepD2HClaim: + """Claim a staging slot (wait + timeout) and copy device→host if this step has rows. + + Unlocked. Timeout is annotated with the unconsumed sids and the + task count so a leaked consume or a stuck write is visible. + """ + try: + return self._controller.stage_step_host(device_snapshot, num_tokens_unpadded, freeze_event, step_holder) + except OmniPrefixCacheStagingTimeoutError as e: + with self._state_lock: + ids = sorted(self._step_ctxs) + raise OmniPrefixCacheStagingTimeoutError( + f"{e}; unconsumed step contexts (ids={ids}) or a stuck write " + f"(in_flight_tasks={self._controller.in_flight_tasks()})" + ) from e + + def _wait_for_host_ready(self) -> None: + """Pop last step's join worklists and wait them out, unlocked. + + JOIN_NEXT_STEP waits ``done``, not ``host_ready``: this save may reuse + those slots, and a pool write still pending behind a reused slot is + unreadable for a delayed hit read (scatter skips reassigned rows). + Escalated deferred writes keep ``host_ready`` — their readers hold the + task in ``staged_list``. Lock covers only the pop. + """ + with self._state_lock: + finished_ids = list(self._join_finished_tids) + next_step_ids = list(self._join_next_step_tids) + self._join_finished_tids.clear() + self._join_next_step_tids.clear() + if finished_ids: + self._controller.join_host_ready(finished_ids) + if next_step_ids: + self._controller.join(next_step_ids) + + @_locked + def _publish_saved_step( + self, + *, + req_order: list[str], + query_start: dict[str, int], + num_sched: dict[str, int], + num_tokens_unpadded: int, + step_outputs: _StepOutputs, + slots_cpu: torch.Tensor | None, + mm_keys: set[str], + freeze_event: torch.cuda.Event | None, + d2h_claim: StepD2HClaim, + bound_tids: list[int], + ) -> tuple[StepId, list[WriteTask]]: + """Takes ``_state_lock``. Register this step's writes and store the + consume-once snapshot. Copies live hits into the snapshot, then + clears them. Returns the queued tasks for the caller to dispatch + unlocked; device→host, budget flush and the eager copy stay outside. + """ + self._commit_drained_writes() + for key, storage in step_outputs.new_key_storage.items(): + self._pool.install_key(key, storage) + self._slot_status.init_table(key) + queued: list[WriteTask] = [] + if step_outputs.immediate: + queued = self._submit_step_writes( + req_order, + query_start, + num_sched, + step_outputs.immediate, + step_outputs.immediate_budget, + slots_cpu, + d2h_claim.views, + freeze_event, + d2h_claim.staging_slot, + d2h_claim.event, + bound_tids, + ) + self._stage_deferred(step_outputs.deferred_chunks, freeze_event) + if self._hit_spans: + # Same-step hits: their rows are registered now (IN_TRANSIT on + # this step's tasks); start the gather before the next step. + self._prefetch_hit_spans() + step_id = self._next_step_id + self._next_step_id += 1 + self._step_ctxs[step_id] = _StepContext( + spans={r: (query_start[r], query_start[r] + num_sched[r]) for r in req_order}, + num_tokens_unpadded=num_tokens_unpadded, + hits=dict(self._hit_spans), + hit_prefetch=dict(self._hit_prefetch), + cached_keys=without_hidden(self._pool.keys()) & mm_keys, + mm_cpu_snapshot=step_outputs.leftover, + mm_cpu_snapshot_event=step_outputs.leftover_event, + d2h=d2h_claim, + ) + self._clear_hit_infos() + return step_id, queued + + def _split_step_outputs( + self, + hidden_states: torch.Tensor | None, + mm_outputs: dict[str, Any], + num_tokens_unpadded: int, + num_tokens_padded: int, + *, + slots_cpu: torch.Tensor | None, + req_order: list[str], + num_sched: dict[str, int], + query_start: dict[str, int], + ) -> _StepOutputs: + """Split this step's outputs into immediate / deferred / leftover. + + Unlocked. One pass over ``mm_outputs``. ``n==0`` has only leftover + mm (no device clones). A deferred key whose first dim is this + step's token count is cloned for the JOIN_ON_FINISH write and + CPU-copied into leftover for this-step materialize. Talker + ``codes.audio`` stays unpadded while hidden is padded; both must + open a pool key. + Lists and other shapes stay leftover. + """ + n = num_tokens_unpadded + immediate: dict[str, torch.Tensor] = {} + deferred_tensors: dict[str, torch.Tensor] = {} + new_key_storage: dict[str, torch.Tensor] = {} + + def alloc_key(key: str, val: torch.Tensor) -> None: + # Pinned allocation, unlocked; installed under the lock at publish. + storage = self._pool.alloc_key(key, val.dtype, int(val.shape[-1])) + if storage is not None: + new_key_storage[key] = storage + + if n > 0: + if hidden_states is not None and (hk := self._policy.hidden_key) is not None: + if hidden_states.ndim < 2 or hidden_states.shape[0] < n: + rows = 0 if hidden_states.ndim < 2 else int(hidden_states.shape[0]) + raise OmniPrefixCacheUnmatchError(f"hidden_states has {rows} rows, need {n}") + alloc_key(hk, hidden_states) + immediate[hk] = hidden_states[:n].clone() + for key, val in mm_outputs.items(): + is_step_rows = _is_step_token_tensor(val, n, num_tokens_padded) + if key in self._policy.deferred_keys: + if is_step_rows: + alloc_key(key, val) + deferred_tensors[key] = val[:n].clone() + continue + if self._policy.skip_immediate_mm(key) or not is_step_rows: + continue + alloc_key(key, val) + immediate[key] = val[:n].clone() + leftover, leftover_event = _snapshot_leftover_mm_cpu(mm_outputs, set(immediate), n, num_tokens_padded) + deferred_chunks: list[tuple[str, _WriteChunk]] = [] + if deferred_tensors: + assert slots_cpu is not None + deferred_chunks = self._pack_deferred_chunks(deferred_tensors, slots_cpu, req_order, num_sched, query_start) + immediate_budget = ( + _BudgetTicket(nbytes=sum(t.numel() * t.element_size() for t in immediate.values())) if immediate else None + ) + return _StepOutputs( + immediate=immediate, + deferred_chunks=deferred_chunks, + leftover=leftover, + leftover_event=leftover_event, + immediate_budget=immediate_budget, + new_key_storage=new_key_storage, + ) + + def _pack_deferred_chunks( + self, + deferred_tensors: dict[str, torch.Tensor], + slots_cpu: torch.Tensor, + req_order: list[str], + num_sched: dict[str, int], + query_start: dict[str, int], + ) -> list[tuple[str, _WriteChunk]]: + """Per-req views of already-cloned deferred tensors. No further clone.""" + ticket = _BudgetTicket(nbytes=sum(t.numel() * t.element_size() for t in deferred_tensors.values())) + out: list[tuple[str, _WriteChunk]] = [] + for req_id in req_order: + sched = num_sched[req_id] + if sched <= 0: + continue + start = query_start[req_id] + end = start + sched + out.append( + ( + req_id, + _WriteChunk( + slots_cpu=slots_cpu[start:end], + tensors={k: v[start:end] for k, v in deferred_tensors.items()}, + budget=ticket, + ), + ) + ) + return out + + def _submit_step_writes( + self, + req_order: list[str], + query_start: dict[str, int], + num_sched: dict[str, int], + device_snapshot: dict[str, torch.Tensor], + budget: _BudgetTicket | None, + slots_cpu: torch.Tensor, + host_views: dict[str, torch.Tensor], + freeze_event, + staging_slot: int, + step_d2h_event, + bound_tids: list[int], + ) -> list[WriteTask]: + """Caller holds ``_state_lock``. Register one WriteTask per request; + returns them for the caller to dispatch once the lock is released. + + Per-req views of the shared device snapshot: one on-device clone, req-scoped + finish/abort, reassigned rows, and completion. Appends bound tids to + `bound_tids` as it goes so a mid-loop raise still unwinds holders. + """ + tasks: list[WriteTask] = [] + for req_id in req_order: + start = query_start[req_id] + end = start + num_sched[req_id] + if end == start: + continue + tensors = {k: v[start:end] for k, v in device_snapshot.items()} + tid = self._request_tasks.alloc_tid() + chunk = _WriteChunk(slots_cpu=slots_cpu[start:end], tensors=tensors, budget=budget) + # Host rows are views into the slot; the committer only waits + # the shared step event. Device→host is already in flight. + chunk.host = {k: v[start:end] for k, v in host_views.items()} + task = WriteTask( + tid=tid, + req_id=req_id, + write_n=self._request_tasks.increment_write_n(req_id), + schedule=WriteSchedule.JOIN_NEXT_STEP, + chunks=[chunk], + freeze_event=freeze_event, + staging_slot=staging_slot, + step_d2h_event=step_d2h_event, + ) + self._map_slots(slots_cpu[start:end], tid, tensors.keys()) + # Bind and pin before register: the slot must never be holder-free + # while the task is live (released at its pool write). + self._controller.staging_bind(staging_slot, StagingBufferHolder.for_task(tid)) + bound_tids.append(tid) + if budget is not None: + self._controller.pin_budget(budget, tid) + self._controller.register(task) + self._request_tasks.track(req_id, tid) + self._join_next_step_tids.append(tid) + tasks.append(task) + return tasks + + def _stage_deferred(self, deferred_chunks: list[tuple[str, _WriteChunk]], freeze_event) -> None: + """Caller holds ``_state_lock``. Register pre-built deferred `_WriteChunk`s + (bytes already reserved by save_outputs). + + A new write (``write_n`` + 1) starts when the open task was closed by + the budget flush, or when a chunk re-acquires a (key, slot) the task + already lost — appending there would let its own reassigned tombstone + drop the fresh rows at scatter. Hits read both through ``staged_list``. + """ + for req_id, chunk in deferred_chunks: + task = self._request_tasks.deferred.get(req_id) + if task is not None and ( + task.reassigned_intersects(chunk) + or self._controller.append_chunk(task, chunk, freeze_event) is not None + ): + task = None + if task is None: + task = WriteTask( + tid=self._request_tasks.alloc_tid(), + req_id=req_id, + write_n=self._request_tasks.increment_write_n(req_id), + schedule=WriteSchedule.JOIN_ON_FINISH, + chunks=[chunk], + freeze_event=freeze_event, + ) + self._request_tasks.deferred[req_id] = task + self._request_tasks.track(req_id, task.tid) + self._controller.register(task, queued=False) + if chunk.budget is not None: + # Safe after register: queued=False keeps the task PENDING + # until escalate, which only this thread calls. + self._controller.pin_budget(chunk.budget, task.tid) + # Block reuse across deferred tenants (preemption path) is + # handled inside _map_slots: the old tenant's rows are skipped. + self._map_slots(chunk.slots_cpu, task.tid, chunk.tensors.keys()) + + # ----------------------------------------------------- occupancy + + def _map_slots(self, slots: torch.Tensor, tid: int, keys: Iterable[str]) -> None: + """Caller holds ``_state_lock``. Record `tid` on these (slot, key); + if another write still owns them, mark those rows skipped on it.""" + keys = tuple(keys) + # Copy committed rows aside before this claim overwrites them. + if self._pending_reads: + self._preserve_for_pending_reads(slots, keys) + for old, key, stolen in self._slot_status.map_slots(slots, tid, keys): + old_task = self._controller.get_task(old) + if old_task is not None: + old_task.add_reassigned(key, stolen) + + def _preserve_for_pending_reads(self, new_slots: torch.Tensor, keys: tuple[str, ...]) -> None: + """Caller holds ``_state_lock``. Copy still-valid committed rows into + each pending ref before ``new_slots`` are claimed. + + Only COMMITTED rows at the reserved version need rescuing: the save + barrier joins JOIN_NEXT_STEP writes to ``done`` and publish drains + before claiming, and JOIN_ON_FINISH in-transit rows are read through + ``staged_list`` task refs directly. + """ + new_set = {int(s) for s in new_slots.tolist()} + for ref in self._pending_reads: + if ref.key not in keys or ref.reserved_version is None or not self._pool.has_key(ref.key): + continue + status = self._slot_status.get_slot_status(ref.key) + take: list[tuple[int, int]] = [] # (row index in ref, slot) + for i, s in enumerate(ref.slots.tolist()): + s = int(s) + if s not in new_set or s in ref.preserved: + continue + committed = int(status.state[s]) == _Occupancy.COMMITTED + if committed and int(status.slot_version[s]) == int(ref.reserved_version[i]): + take.append((i, s)) + if not take: + continue + rows = self._pool.rows(ref.key, torch.tensor([s for _, s in take], dtype=torch.long)) + for j, (_, s) in enumerate(take): + ref.preserved[s] = rows[j].clone() + + @torch.inference_mode() + def _commit_drained_writes(self) -> None: + """Fold completed/failed writes into occupancy. Caller holds ``_state_lock``. + + A failed write leaves rows absent behind hashes vLLM already + published — unservable and unrecoverable, so fatal. The fatality is + sticky: the drained record is one-shot and the first raise can land + in a swallowed path (a prefetch Future dropped by discard_step), so + every later entry must keep raising rather than serve zeros. + """ + if self._fatal_write_failure is not None: + raise OmniPrefixCacheUnmatchError(self._fatal_write_failure) + failed = self._controller.drain_failed() + if failed: + self._fatal_write_failure = ( + f"prefix cache write failed for task(s) {failed}; cached rows lost behind published hashes" + ) + raise OmniPrefixCacheUnmatchError(self._fatal_write_failure) + drained = self._controller.drain_completed() + if drained: + self._slot_status.commit(drained) + self._request_tasks.drop_completed(drained) + + # --------------------------------------------------- consume-once + + def _take_step_ctx(self, step_id: int) -> _StepContext: + """Pop the context for this step id (exactly once). Caller holds ``_state_lock``.""" + ctx = self._step_ctxs.pop(step_id, None) + if ctx is None: + raise OmniPrefixCacheUnmatchError( + f"step context {step_id} missing (have {sorted(self._step_ctxs)}); already consumed or never saved" + ) + return ctx + + def _release_step_staging_slot(self, slot: int, step_id: int) -> None: + self._controller.staging_release(slot, StagingBufferHolder.for_step(step_id)) + + def _release_step_staging(self, ctx: _StepContext, step_id: int) -> None: + """Drop this step's staging hold. Does not require ``_state_lock``. + + materialize/discard: task holds leave at their pool write. + """ + if ctx.d2h is not None: + self._release_step_staging_slot(ctx.d2h.staging_slot, step_id) + + def _release_staging_on_failed_save( + self, slot: int, step_holder: StagingBufferHolder, bound_tids: list[int] + ) -> None: + """save raised after claiming the slot: drop the step hold and any + task holds. Does not require ``_state_lock``. + """ + self._release_step_staging_slot(slot, step_holder.owner_id) + for tid in bound_tids: + self._controller.staging_release(slot, StagingBufferHolder.for_task(tid)) + + # -------------------------------------------------- slot ref / fetch + + def _get_hit_slots(self, hit_upto: int, hit_blocks: list[int]) -> torch.Tensor: + """Prefix-hit block ids → KV slot ids. Alignment is checked at + ``new_step_starts``. Does not require ``_state_lock``. + """ + bs = self._config.block_size + block_ids = torch.tensor(hit_blocks, dtype=torch.int64) + return (block_ids.unsqueeze(1) * bs + torch.arange(bs)).reshape(-1)[:hit_upto] + + def _slot_ref(self, slots: torch.Tensor, key: str, req_id: str) -> _SlotRef: + """Caller holds ``_state_lock``. Pin a ``_SlotRef`` for `slots` (no data movement). + + Rows still being written win over the CPU pool: they may not have + landed yet, and a pool read would return zero/stale values. + JOIN_NEXT_STEP tasks go in ``join_tids`` (wait-then-pool at + fetch). JOIN_ON_FINISH tasks stay as refs for fetch_host. + + Hidden rejects any empty hole (prefetch skips; materialize + raises). Other keys only need a source — holes fall to the pool. + + The ref is registered in ``_pending_reads`` so a write that later + reuses one of these slots copies the old row aside first (COW). Reads + are planned synchronously under the lock at hit time (prefetch runs at + ``new_step_starts`` / publish), before the block can be reused, so this + registration always precedes the reuse. ``_fetch_source`` unregisters. + """ + status = self._slot_status.get_slot_status(key) + states = status.state[slots] + tids = status.tids[slots] + reserved_version = status.slot_version[slots].clone() + staged_mask = states == _Occupancy.IN_TRANSIT + + staged: list[tuple[WriteTask, torch.Tensor]] = [] + join_tids: list[int] = [] + for tid in {int(t) for t in tids[staged_mask].tolist()}: + task = self._controller.get_task(tid) if tid != 0 else None + if task is None: + _raise_unreadable_hit(req_id, key, f"in-transit entry {tid} cannot serve them") + if task.schedule is WriteSchedule.JOIN_NEXT_STEP: + join_tids.append(task.tid) + else: + staged.append((task, staged_mask & (tids == tid))) + + already_staged = self._pool.has_key(key) + has_source = already_staged or bool(staged) or bool(join_tids) + if is_hidden_key(key): + n_abs = int((states == _Occupancy.ABSENT).sum()) + if n_abs or not has_source: + _raise_unreadable_hit(req_id, key, f"{n_abs} absent slots") + ref = _SlotRef( + slots=slots, + key=key, + req_id=req_id, + already_staged=already_staged, + staged_list=staged, + join_tids=join_tids, + reserved_version=reserved_version, + ) + self._pending_reads.append(ref) + return ref + + def _apply_preserved_rows(self, src: _SlotRef, out: torch.Tensor) -> torch.Tensor: + """Write ``src.preserved`` rows over the matching slots in ``out``.""" + if not src.preserved: + return out + for i, s in enumerate(src.slots.tolist()): + row = src.preserved.get(int(s)) + if row is not None: + out[i] = row + return out + + def _unregister_pending_read(self, src: _SlotRef) -> None: + """Drop ``src`` from ``_pending_reads``. Identity, not dataclass eq + (tensors in the ref make ``==`` unusable). Safe if already removed.""" + with self._state_lock: + self._pending_reads = [ref for ref in self._pending_reads if ref is not src] + + def _fetch_source(self, src: _SlotRef) -> torch.Tensor: + """Fetch a planned row source (execute phase, no lock). + + One key is one schedule: ``join_tids`` (JOIN_NEXT_STEP) and + ``staged_list`` (JOIN_ON_FINISH) do not coexist. Immediate: wait + ``done``, drain, read the pool. Deferred: pool rows already + written, overlay ``fetch_host`` on the still-in-progress mask. + Slots a later write reused are served from ``src.preserved`` + (copied under the lock before that write claimed them). + """ + try: + return self._fetch_source_inner(src) + finally: + self._unregister_pending_read(src) + + def _fetch_source_inner(self, src: _SlotRef) -> torch.Tensor: + # For JOIN_NEXT_STEP, wait `done`, drain, read the pool + if src.join_tids: + self._controller.join(src.join_tids) + with self._state_lock: + self._commit_drained_writes() + joined = self._apply_preserved_rows(src, self._pool.rows(src.key, src.slots)) + self._ensure_not_reassigned( + src.slots, + src.key, + req_id=src.req_id, + planned_version=src.reserved_version, + preserved_slots=src.preserved, + ) + return joined + + # For JOIN_ON_FINISH, pool rows already written, overlay `fetch_host` + n = int(src.slots.numel()) + out: torch.Tensor | None = None + if src.already_staged: + out = self._pool.rows(src.key, src.slots) + in_transit = None + for task, mask in src.staged_list: + try: + rows = self._controller.fetch_host(task, src.slots[mask], src.key) + except KeyError: + _raise_unreadable_hit( + src.req_id, + src.key, + f"entry {task.tid} (req {task.req_id}, write_n {task.write_n}) cannot serve them", + ) + if out is None: + out = torch.zeros((n, rows.shape[-1]), dtype=rows.dtype) + out[mask] = rows + in_transit = mask if in_transit is None else in_transit | mask + if out is not None: + out = self._apply_preserved_rows(src, out) + self._ensure_not_reassigned( + src.slots, + src.key, + in_transit_mask=in_transit, + req_id=src.req_id, + planned_version=src.reserved_version, + preserved_slots=src.preserved, + ) + if out is None: + _raise_unreadable_hit(src.req_id, src.key, "no source") + return out + + def _ensure_not_reassigned( + self, + slots: torch.Tensor, + key: str, + *, + in_transit_mask: torch.Tensor | None = None, + req_id: str = "?", + planned_version: torch.Tensor | None = None, + preserved_slots: dict[int, torch.Tensor] | None = None, + ) -> None: + """Takes ``_state_lock``. Post-fetch check: rows read unlocked may + have been given to a newer write between plan and read (block reuse). + + vLLM frees a request's blocks when it finishes or is aborted, one + step before ``finished_req_ids`` reaches us, and may reuse them at + once. A planned ``_SlotRef`` is registered in ``_pending_reads``; + the reusing write copy-on-writes those COMMITTED rows into + ``preserved`` before it claims the slots. Those slots are safe. + A version mismatch with no preserved copy means the pool rows are + a newer tenant's — raise for live and finished alike, never serve + them. + + ``planned_version`` is the per-slot version captured at plan time. + ``in_transit_mask`` excludes JOIN_ON_FINISH slots this task itself + still serves through ``fetch_host``. + """ + with self._state_lock: + status = self._slot_status.get_slot_status(key) + if planned_version is not None: + violated = status.slot_version[slots] != planned_version + else: + violated = status.state[slots] == _Occupancy.IN_TRANSIT + if in_transit_mask is not None: + violated &= ~in_transit_mask + if preserved_slots: + for i, s in enumerate(slots.tolist()): + if int(s) in preserved_slots: + violated[i] = False + if not bool(violated.any()): + return + live = req_id in self._request_tasks.live_reqs + _raise_unreadable_hit( + req_id, + key, + f"{int(violated.sum())} hit slots reassigned before this delayed read " + f"({'live' if live else 'finished'} req, block reuse)", + ) + + # ---------------------------------------------------------- merge + + def _merge_cached_for_req( + self, + ctx: _StepContext, + req_id: str, + key: str, + current_cpu: torch.Tensor, + hit_sources: dict[tuple[str, str], _SlotRef | Future], + ) -> torch.Tensor: + """Hit prefix + this step's rows for one (req, key). + + No hit → this step's slice only. Prefetch Future → write the + slice into the reserved tail. Else cat(fetch, new). + """ + start, end = ctx.spans[req_id] + new_rows = current_cpu[start:end] + src = hit_sources.get((req_id, key)) + if src is None: + return new_rows + if isinstance(src, Future): + # Prefetched during the forward, prefix already in place; only + # this step's rows land here. result() re-raises fetch/validation + # errors — unread hits still raise after the thread hop. + merged = src.result() + merged[merged.shape[0] - new_rows.shape[0] :] = new_rows + return merged + cached = self._fetch_source(src) + return torch.cat([cached, new_rows], dim=0) + + def _merge_uncached_mm( + self, + ctx: _StepContext, + req_ids: list[str], + cached_keys: set[str], + mm_out: dict[str, dict[str, Any]], + ) -> None: + """Write leftover mm that is not a pool key into mm_out. + + No hit concat: leftover mm was already copied to CPU at save + (``ctx.mm_cpu_snapshot``). cached_keys already went through + _merge_cached_for_req. ``req_ids`` is a subset of ``ctx.spans``. + """ + leftover = {k: v for k, v in ctx.mm_cpu_snapshot.items() if k not in cached_keys and not is_hidden_key(k)} + if not leftover: + return + from vllm_omni.utils.mm_outputs import to_payload_element + + order = list(ctx.spans) + total_length = sum(e - s for s, e in ctx.spans.values()) + for key, val in leftover.items(): + per_req: dict[str, Any] = {} + for req_id in req_ids: + idx = order.index(req_id) + start, end = ctx.spans[req_id] + per_req[req_id] = to_payload_element( + val, + idx, + start=start, + end=end, + pass_lists_through=True, + seq_len=total_length, + ) + mm_out[key] = per_req diff --git a/vllm_omni/core/prefix_cache/runner_mixin.py b/vllm_omni/core/prefix_cache/runner_mixin.py new file mode 100644 index 000000000..25fa4fe2b --- /dev/null +++ b/vllm_omni/core/prefix_cache/runner_mixin.py @@ -0,0 +1,148 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Runner-facing access layer for the omni prefix cache. + +One copy of the integration rules shared by the GPU and NPU runners: +one-shot manager construction, the model cache-policy snapshot, and the +per-step lifecycle entry. The manager/controller state machines and the +platform output-building flows stay in their owners. + +Not exported from the package ``__init__``; runners import this module +directly. No module-level vLLM import (``get_pp_group`` is resolved at +call time) so ``tests/core`` keeps loading the package without vLLM. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +from vllm_omni.core.prefix_cache.group_view import get_prefix_cache_group_view +from vllm_omni.core.prefix_cache.interface import ( + ModelCachePolicy, + OmniPrefixCacheUnmatchError, + PrefixCacheConfig, +) +from vllm_omni.core.prefix_cache.manager import OmniPrefixCacheManager +from vllm_omni.data_entry_keys import flatten_payload + +if TYPE_CHECKING: + import torch + from vllm.v1.core.sched.output import SchedulerOutput + + +class PrefixCacheRunnerMixin: + """Prefix-cache integration for a model runner (no state machine here). + + Host-class contract — the runner must provide: + ``input_batch`` live batch with the group-0 block table + ``kv_cache_config`` for ``kv_cache_groups`` at construction + ``omni_prefix_cache`` set here; ``None`` until first-step build + ``_omni_prefix_cache_cfg`` staged by the runner's kv-cache init + (``stage_prefix_cache_config``), consumed here + + No cooperative ``__init__``: the class attributes below are safe + defaults (``ModelCachePolicy`` is frozen); runners keep initializing + their instance fields as before. + """ + + # Host-class fields (annotation only; the runner initializes them). + input_batch: Any + kv_cache_config: Any + is_pooling_model: bool + + omni_prefix_cache: OmniPrefixCacheManager | None = None + _omni_prefix_cache_cfg: PrefixCacheConfig | None = None + # Snapshotted once at load_model; the hot path must not re-probe model + # attributes every step. The same object is registered on the manager. + _omni_cache_policy: ModelCachePolicy = ModelCachePolicy() + + def _snapshot_prefix_cache_model_policy(self, model) -> None: + """Freeze the model's cache policy at load_model.""" + self._omni_cache_policy = ModelCachePolicy.from_model(model) + + 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 hidden-state write into the + omni prefix cache 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. + """ + return self._omni_cache_policy.needs_full_hidden_states + + def _ensure_omni_prefix_cache(self) -> None: + """One-shot construction (caller gates on the staged config). + + Only the last PP rank writes and materializes the cache. Other ranks + would register hits they can never serve (save_outputs is last-rank + only), so skip construction there. + """ + from vllm.distributed.parallel_state import get_pp_group + + cfg = self._omni_prefix_cache_cfg + self._omni_prefix_cache_cfg = None + if cfg is None or not get_pp_group().is_last_rank: + return + view = get_prefix_cache_group_view( + self.input_batch, + cfg.block_size, + kv_cache_groups=getattr(self.kv_cache_config, "kv_cache_groups", None), + ) + if view is None: + raise OmniPrefixCacheUnmatchError( + "omni prefix caching requires a block table on the input batch; " + "disable enable_prefix_caching for this model" + ) + manager = OmniPrefixCacheManager(cfg, view) + manager.register_policy(self._omni_cache_policy) + self.omni_prefix_cache = manager + + def _prefix_cache_step_begin(self, scheduler_output: SchedulerOutput) -> None: + """Per-step lifecycle entry; must run before ``_update_states`` + removes finished requests (finished/abort escalation happens inside + ``new_step_starts``). Builds the manager on the first real step.""" + if self.omni_prefix_cache is None and self._omni_prefix_cache_cfg is not None: + self._ensure_omni_prefix_cache() + if self.omni_prefix_cache is not None: + self.omni_prefix_cache.new_step_starts(scheduler_output) + + def _prefix_cache_save_step( + self, + hidden_states: torch.Tensor, + multimodal_outputs: dict | None, + *, + num_tokens_unpadded: int, + num_tokens_padded: int, + ) -> int | None: + """Write this step's outputs into the cache. + + None when this stage/rank does not write (pooling, cache off, or not + the last PP rank). The returned step id must be consumed exactly once + — materialize or discard_step — by the output path. + """ + from vllm.distributed.parallel_state import get_pp_group + + if self.is_pooling_model or self.omni_prefix_cache is None or not get_pp_group().is_last_rank: + return None + return self.omni_prefix_cache.save_outputs( + hidden_states, + flatten_payload(multimodal_outputs) if multimodal_outputs else {}, + num_tokens_unpadded=num_tokens_unpadded, + num_tokens_padded=num_tokens_padded, + ) + + def _prefix_cache_materialize( + self, step_id: int | None, req_ids: list[str] + ) -> tuple[dict[str, torch.Tensor] | None, dict | None]: + """Per-request merged outputs for a saved step. + + ``req_ids`` must be the save-time snapshot, never the live + ``input_batch`` (under async output this runs a step late). + ``step_id`` None means save did not run — nothing to consume. + ``mm_outputs`` is all-or-nothing; empty only when the step had no mm. + """ + if step_id is None or self.omni_prefix_cache is None: + return None, None + outs = self.omni_prefix_cache.materialize(step_id, list(req_ids)) + return outs.hidden_states, (outs.mm_outputs or None) diff --git a/vllm_omni/deploy/qwen3_tts.yaml b/vllm_omni/deploy/qwen3_tts.yaml index 4019f10aa..8f5ab584d 100644 --- a/vllm_omni/deploy/qwen3_tts.yaml +++ b/vllm_omni/deploy/qwen3_tts.yaml @@ -63,11 +63,12 @@ stages: max_num_seqs: 64 gpu_memory_utilization: 0.3 trust_remote_code: true - # 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. + # Omni prefix cache (reuse of talker codec rows on shared prefixes). Off + # by default; set true when text repeats across requests (same voice or + # shared ref_audio). The talker opts out of hidden-state caching and + # defers codes.audio writes, so per-step cost is small. Stage 1 (codec + # decoder) is not a supported prefix-cache stage: keep it false there. + # See docs/design/feature/prefix_caching.md for the stage matrix. enable_prefix_caching: false async_scheduling: true max_num_batched_tokens: 32768 diff --git a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py index f78b777a5..18353971c 100644 --- a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py +++ b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py @@ -107,6 +107,9 @@ class ExecuteModelState(NamedTuple): cudagraph_stats: CUDAGraphStat | None batch_desc: BatchDescriptor multimodal_outputs: Any # Omni-Specific + # Omni: prefix-cache step id; its context must be consumed exactly + # once (materialize or discard_step) in sample_tokens. + prefix_cache_step_id: int | None = None class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin, DuplexSamplingRunnerMixin): """Autoregressive NPU model runner that returns hidden states per request.""" @@ -211,88 +214,6 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin, Duplex 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 _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 _stage_deferred_prefix_cache_mm_outputs( - self, - *, - scheduler_output: SchedulerOutput, - multimodal_outputs: Any, - query_start_loc_cpu: Any, - ) -> None: - """See gpu_ar_model_runner._stage_deferred_prefix_cache_mm_outputs.""" - if self.omni_prefix_cache is None: - return - - deferred_mm_cache_keys = self._deferred_prefix_cache_mm_keys() - if not deferred_mm_cache_keys: - return - - 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, - ) - - def _maybe_update_prefix_cache( - self, - hidden_states: torch.Tensor, - multimodal_outputs: dict, - num_tokens_unpadded: int, - num_tokens_padded: int, - ): - if self.omni_prefix_cache is not None and get_pp_group().is_last_rank: - hs_for_cache = hidden_states if self._model_needs_full_prefix_hidden_states() else None - slot_mapping_gpu = self.input_batch.block_table[0].slot_mapping.gpu - slot_mapping_cpu = slot_mapping_gpu[:num_tokens_padded].cpu() - self.omni_prefix_cache.update_omni_tensor_prefix_cache( - 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=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( - self, - hidden_states: torch.Tensor, - multimodal_outputs: dict, - num_scheduled_tokens: dict[str, int], - ) -> tuple[dict[str, torch.Tensor] | None, dict | None]: - combined_hidden_states, combined_multimodal_outputs = None, None - if self.omni_prefix_cache is not None: - 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, - multimodal_outputs=flatten_payload(multimodal_outputs) if multimodal_outputs else multimodal_outputs, - num_scheduled_tokens=num_scheduled_tokens, - ) - return combined_hidden_states, combined_multimodal_outputs - @staticmethod def _resolve_req_hidden_states( hidden_states_cpu: torch.Tensor, @@ -452,11 +373,8 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin, Duplex 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, - ) + # Exactly once per real scheduler_output, before _update_states. + self._prefix_cache_step_begin(scheduler_output) # -------------------------------------- Omni-new ------------------------------------------------- if self.speculative_config is not None and self.speculative_config.use_ngram_gpu(): @@ -782,9 +700,9 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin, Duplex hidden_states, aux_hidden_states = hidden_states # -------------------------------------- Omni-new ------------------------------------------------- - self._maybe_update_prefix_cache( - hidden_states=hidden_states, - multimodal_outputs=multimodal_outputs, + prefix_cache_step_id = self._prefix_cache_save_step( + hidden_states, + multimodal_outputs, num_tokens_unpadded=num_tokens_unpadded, num_tokens_padded=num_tokens_padded, ) @@ -864,6 +782,7 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin, Duplex cudagraph_stats, batch_desc, multimodal_outputs, # Omni-specific + prefix_cache_step_id, ) self.kv_connector_output = kv_connector_output @@ -956,6 +875,7 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin, Duplex cudagraph_stats, batch_desc, multimodal_outputs, # Omni-Specific + prefix_cache_step_id, ) = self.execute_model_state # Clear ephemeral state. self.execute_model_state = None @@ -1106,9 +1026,15 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin, Duplex hidden_states_cpu = None req_hidden_states_cpu: dict[str, torch.Tensor] | None = None audio_sparse_output = engine_output_type == "audio" and sparse_mm_req_ids is not None + _omni_cache_on = self.omni_prefix_cache is not None needs_scheduled_hidden_payload = needs_pooler_payload and ( - self.omni_prefix_cache is None or not self._model_needs_full_prefix_hidden_states() + not _omni_cache_on or not self._model_needs_full_prefix_hidden_states() ) + if not needs_pooler_payload and prefix_cache_step_id is not None: + # No consumer for this step's merge: consume the step context by + # id (exactly-once contract). The cache write still lands. + self.omni_prefix_cache.discard_step(prefix_cache_step_id) + prefix_cache_step_id = None if needs_scheduled_hidden_payload: num_valid_tokens = min( int(scheduler_output.total_num_scheduled_tokens), @@ -1131,27 +1057,17 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin, Duplex if callable(query_start_loc_cpu): query_start_loc_cpu = query_start_loc_cpu() - self._stage_deferred_prefix_cache_mm_outputs( - scheduler_output=scheduler_output, - multimodal_outputs=multimodal_outputs, - query_start_loc_cpu=query_start_loc_cpu, - ) - pooler_output: list[dict[str, object]] | None = None if needs_pooler_payload: combined_hidden_states = None combined_multimodal_outputs = None mm_cpu = None - if self.omni_prefix_cache is not None: + if _omni_cache_on: ( combined_hidden_states, combined_multimodal_outputs, - ) = self._maybe_get_combined_prefix_cache_tensors( - hidden_states, - multimodal_outputs, - scheduler_output.num_scheduled_tokens, - ) - if self.omni_prefix_cache is None or combined_multimodal_outputs is None: + ) = self._prefix_cache_materialize(prefix_cache_step_id, list(req_ids_output_copy)) + if not _omni_cache_on or combined_multimodal_outputs is None: mm_cpu = build_mm_cpu( flatten_payload(multimodal_outputs) if multimodal_outputs else multimodal_outputs ) diff --git a/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py b/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py index 59dfc88a3..d796dec8e 100644 --- a/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py +++ b/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py @@ -459,6 +459,7 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin cudagraph_stats, _batch_desc, multimodal_outputs_raw, # Omni-Specific + _prefix_cache_step_id, # generation stages never save to the prefix cache ) = self.execute_model_state # Clear ephemeral state. self.execute_model_state = None diff --git a/vllm_omni/platforms/npu/worker/npu_model_runner.py b/vllm_omni/platforms/npu/worker/npu_model_runner.py index bd7c7dce4..d8e352d48 100644 --- a/vllm_omni/platforms/npu/worker/npu_model_runner.py +++ b/vllm_omni/platforms/npu/worker/npu_model_runner.py @@ -20,7 +20,7 @@ from vllm_ascend.ops.rotary_embedding import update_cos_sin from vllm_ascend.utils import enable_sp, lmhead_tp_enable from vllm_ascend.worker.model_runner_v1 import SEQ_LEN_WITH_MAX_PA_WORKSPACE -from vllm_omni.core.prefix_cache import OmniTensorPrefixCache +from vllm_omni.core.prefix_cache import stage_prefix_cache_config from vllm_omni.model_executor.models.output_templates import OmniOutput from vllm_omni.platforms.npu._310p import is_310p from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner @@ -36,35 +36,32 @@ else: class OmniNPUModelRunner(OmniGPUModelRunner, NPUModelRunner): def initialize_kv_cache(self, kv_cache_config) -> None: - """Create the omni tensor prefix cache. - - The omni prefix cache is used to store the hidden states of the prefix tokens - """ + """Stage the omni prefix-cache config (hidden / mm tensors reused on hits).""" NPUModelRunner.initialize_kv_cache(self, kv_cache_config) - if self.omni_prefix_cache is None and self.cache_config.enable_prefix_caching: - # Read num_blocks back off self.kv_cache_config: vllm-ascend - # deepcopies the config it was handed, so the value it stored is the - # authoritative one, not our caller's argument. - num_blocks = self.kv_cache_config.num_blocks - self.omni_prefix_cache = OmniTensorPrefixCache( - num_blocks=num_blocks, - block_size=self.cache_config.block_size, - hidden_size=self.model_config.get_hidden_size(), - hs_dtype=self.dtype, - ) - logger.info( - "Initialized omni prefix cache on NPU (num_blocks=%d, block_size=%d, hidden_size=%d). " - "Hidden-state cache is pinned host memory of roughly %.1f GiB; each per-token " - "multimodal output key allocates another tensor of the same block shape.", - num_blocks, - self.cache_config.block_size, - self.model_config.get_hidden_size(), - num_blocks - * self.cache_config.block_size - * self.model_config.get_hidden_size() - * self.dtype.itemsize - / (1024**3), + if getattr(self, "_omni_prefix_cache_cfg", None) is None: + # Same gate as the GPU runner (pooling stage, kv_consumer / + # kv_both, hybrid kv groups). Read the config back off + # self.kv_cache_config: vllm-ascend deepcopies the one it was + # handed, so the stored value is the authoritative one. + # Controller runs in eager mode on NPU (no CUDA streams: + # dispatch() completes the copy+scatter synchronously). Built once + # on the first step via the inherited _ensure_omni_prefix_cache. + cfg = stage_prefix_cache_config( + kv_cache_config=self.kv_cache_config, + cache_config=self.cache_config, + kv_transfer_config=getattr(self.vllm_config, "kv_transfer_config", None), + scheduler_config=self.scheduler_config, + model_config=self.model_config, + is_pooling_model=self.is_pooling_model, + speculative_config=self.speculative_config, ) + if cfg is not None: + self._omni_prefix_cache_cfg = cfg + logger.info( + "Initialized omni prefix cache on NPU (eager mode, num_blocks=%d, block_size=%d).", + cfg.num_blocks, + cfg.block_size, + ) def load_model(self, *args, **kwargs) -> None: if is_310p(): @@ -84,6 +81,8 @@ class OmniNPUModelRunner(OmniGPUModelRunner, NPUModelRunner): if callable(candidate): override_fn = candidate self._sampled_token_ids_cpu_override = override_fn + self._snapshot_prefix_cache_model_policy(model) + self._pooler_payload_include_hidden_flag = bool(getattr(model, "omni_pooler_payload_include_hidden", True)) self._omni_query_start_loc_model_kwarg = bool(getattr(model, "supports_omni_query_start_loc", False)) self._maybe_enable_output_token_ids_for_model_sampler() self._init_talker_mtp() diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index 06fdfa18b..b918b57f1 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -342,6 +342,9 @@ class ExecuteModelState(NamedTuple): multimodal_outputs: Any # slot_mappings for attention/drafter (aligned with upstream v1 API) slot_mappings: dict[str, torch.Tensor] | list[dict[str, torch.Tensor]] | None = None + # OMNI: prefix-cache step id; the step context saved under it must be + # consumed exactly once (materialize or discard_step) in sample_tokens. + prefix_cache_step_id: int | None = None class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, DuplexSamplingRunnerMixin): @@ -565,7 +568,14 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex CUDAGraphWrapper.clear_all_graphs() BreakableCUDAGraphWrapper.clear_all_graphs() - # 6. Delegate to upstream shutdown (model = None, KV caches, workspace). + # 6. Stop the prefix-cache committer before the CUDA context and the + # tensors it copies from are torn down. + cache = self.omni_prefix_cache + if cache is not None: + cache.shutdown() + self.omni_prefix_cache = None + + # 7. Delegate to upstream shutdown (model = None, KV caches, workspace). super().shutdown() def _capture_talker_mtp_graphs(self) -> None: @@ -625,18 +635,6 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex 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 _get_runner_assisted_full_attention_metadata_request( self, *, @@ -724,132 +722,34 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex ) return 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, - hidden_states_cpu: torch.Tensor | None, - multimodal_outputs: dict, - num_tokens_unpadded: int, - num_tokens_padded: int, - ): - """If prefix caching is enabled and it's the last pipeline parallelism rank, - retrieve the hidden states & multimodal outputs from the prefix cache based - on our batch slot mappings. - """ - # Cache hidden states if we've enabled hidden state prefix caching - # unless this isn't the last pipeline parallelism 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. - hs_for_cache = hidden_states if self._model_needs_full_prefix_hidden_states() else None - # FIX: The .cpu attribute of slot_mapping is stale (not updated by the Triton - # _compute_slot_mapping_kernel which only writes to .gpu). We must use .gpu and - # sync back to CPU to get the correctly computed slot mapping. - slot_mapping_gpu = self.input_batch.block_table[0].slot_mapping.gpu - slot_mapping_cpu = slot_mapping_gpu[:num_tokens_padded].cpu() - self.omni_prefix_cache.update_omni_tensor_prefix_cache( - 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=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]: - """If prefix caching is enabled, extract the merged hidden states and multimodal outputs for - all requests in the batch (including those that aren't a hit on Prefix cache). - """ - # 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() - ): - 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, - hidden_states_cpu=hidden_states_cpu, - 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, - multimodal_outputs=flatten_payload(multimodal_outputs) if multimodal_outputs else multimodal_outputs, - num_scheduled_tokens=num_scheduled_tokens, - ) - return combined_hidden_states, combined_multimodal_outputs - - def _stage_deferred_prefix_cache_mm_outputs( - self, - *, - scheduler_output: SchedulerOutput, - multimodal_outputs: Any, - query_start_loc_cpu: Any, - ) -> None: - if self.omni_prefix_cache is None: - return - - deferred_mm_cache_keys = self._deferred_prefix_cache_mm_keys() - if not deferred_mm_cache_keys: - return - - 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, - ) - def _prepare_prefix_cache_pooler_payload_sources( self, *, - hidden_states: torch.Tensor, staged_hidden_states_cpu: torch.Tensor | None, - multimodal_outputs: Any, - scheduler_output: SchedulerOutput, needs_scheduled_hidden_payload: bool, + req_ids: list[str], + step_id: int | None, ) -> tuple[torch.Tensor | None, dict[str, torch.Tensor] | None, dict | None]: + """Prefix-cache payload sources for the pooler payload build. + + req_ids must be the step's snapshot, not ``input_batch.req_ids``: under + async omni output this runs a step late, when the live batch has moved on. + """ + if not get_pp_group().is_last_rank: + raise RuntimeError("Omni prefix-cache merge is only valid on the last pipeline parallel rank.") hidden_states_cpu = None if needs_scheduled_hidden_payload: 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 - - combined_hidden_states, combined_multimodal_outputs = self._maybe_get_combined_prefix_cache_tensors( - hidden_states, - staged_hidden_states_cpu, - multimodal_outputs, - scheduler_output.num_scheduled_tokens, - ) - return hidden_states_cpu, combined_hidden_states, combined_multimodal_outputs + if step_id is None: + # save_outputs did not run this step (no last-rank AR output), so + # there is no step context to consume; serve the fresh slice. + if hidden_states_cpu is None and staged_hidden_states_cpu is not None: + hidden_states_cpu = staged_hidden_states_cpu + return hidden_states_cpu, None, None + combined_hidden, combined_mm = self._prefix_cache_materialize(step_id, list(req_ids)) + return hidden_states_cpu, combined_hidden, combined_mm def _build_omni_pooler_payload( self, @@ -956,15 +856,9 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex if hasattr(self.model, "_clear_warmup_state"): self.model._clear_warmup_state() - # Async-write pipeline: apply any pending GPU->CPU prefix-cache writes - # whose copy event has already fired. Non-blocking — entries whose D2H - # is still in flight stay queued and will be picked up on the next - # step's drain. This guarantees any downstream read of - # ``omni_prefix_cache.hidden_states_cache`` / - # ``omni_prefix_cache.mm_outputs_cache`` in this step sees the - # state produced no later than the previous forward step. - if self.omni_prefix_cache is not None: - self.omni_prefix_cache.drain_ready_async_writes() + # Exactly once per real scheduler_output, before _update_states + # (dummy/warmup runs never reach execute_model). + self._prefix_cache_step_begin(scheduler_output) # [Omni] Handle KV transfer BEFORE updating states (which removes finished requests) finished_reqs = getattr(scheduler_output, "finished_requests_needing_kv_transfer", {}) @@ -988,12 +882,6 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex 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 = self.routed_experts_capturer if capturer is not None and hasattr(capturer, "finalize_pending_copy"): @@ -1294,30 +1182,22 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex else: hidden_states, multimodal_outputs = self.extract_multimodal_outputs(model_output) hidden_states_cpu = None - - # Async-write pipeline (replaces the per-step blocking - # ``.to("cpu")`` + ``aten::index_put_`` on pageable host memory). - # Schedules non-blocking GPU->CPU copies on a dedicated stream; - # the actual CPU scatter into ``hidden_states_cache`` / - # ``mm_outputs_cache`` happens in ``drain_ready_async_writes`` - # at the top of subsequent execute_model() calls. - if not self.is_pooling_model and self.omni_prefix_cache is not None and get_pp_group().is_last_rank: - hs_for_cache = hidden_states if self._model_needs_full_prefix_hidden_states() else None - # Some models (e.g. qwen3-tts-talker) opt out of full-hidden-state - # prefix caching but the downstream pooler payload path still - # needs a CPU hidden-states view. Materialize it synchronously - # in that case; the legacy behavior is preserved. - if hs_for_cache is None and self._model_omni_pooler_payload_include_hidden(): - hidden_states_cpu = hidden_states[:num_tokens_unpadded].detach().to("cpu").contiguous() - slot_mapping_gpu = self.input_batch.block_table[0].slot_mapping.gpu - self.omni_prefix_cache.schedule_async_write( - hidden_states_gpu=hs_for_cache, - multimodal_outputs_gpu=(flatten_payload(multimodal_outputs) if multimodal_outputs else None), - slot_mapping_gpu=slot_mapping_gpu, - num_tokens_unpadded=num_tokens_unpadded, - num_tokens_padded=num_tokens_padded, - skip_mm_cache_keys=self._deferred_prefix_cache_mm_keys(), - ) + # Prefix-cache write: freeze + submit; policy gating (full + # hidden, skip/deferred keys) happens inside the manager. + prefix_cache_step_id = self._prefix_cache_save_step( + hidden_states, + multimodal_outputs, + num_tokens_unpadded=num_tokens_unpadded, + num_tokens_padded=num_tokens_padded, + ) + if ( + prefix_cache_step_id is not None + and not self._model_needs_full_prefix_hidden_states() + and self._model_omni_pooler_payload_include_hidden() + ): + # Opt-out models keep the scheduled-slice CPU view for the + # pooler payload (runner-side bypass, legacy parity). + hidden_states_cpu = hidden_states[:num_tokens_unpadded].detach().to("cpu").contiguous() if not self.broadcast_pp_output: # Common case. @@ -1391,6 +1271,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex cudagraph_stats, multimodal_outputs, slot_mappings, # OMNI: pass slot_mappings for drafter + prefix_cache_step_id, ) self.kv_connector_output = kv_connector_output @@ -1673,12 +1554,12 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex return raw def _model_omni_pooler_payload_include_hidden(self) -> bool: - return self._runner_model_omni_flag("omni_pooler_payload_include_hidden", default=True) + return self._pooler_payload_include_hidden_flag def _should_defer_full_payload_d2h(self) -> bool: """Keep opted-in full payloads on device until request completion.""" return ( - getattr(self, "omni_prefix_cache", None) is None + self.omni_prefix_cache is None and self._runner_model_omni_flag("omni_payload_at_request_end") and self._should_accumulate_full_payload_output() ) @@ -1707,8 +1588,6 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex def _should_use_async_omni_output(self) -> bool: if not self.use_async_scheduling: return False - if self.omni_prefix_cache is not None: - return False if self.speculative_config is not None: return False @@ -1854,6 +1733,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex num_scheduled_tokens_np: np.ndarray, query_start_loc_cpu: Any, postprocess_already_applied: bool = False, + prefix_cache_step_id: int | None = None, ) -> OmniModelRunnerOutput: combined_hidden_states = None combined_multimodal_outputs = None @@ -1877,12 +1757,11 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex and needs_pooler_payload and (self.omni_prefix_cache is None or not self._model_needs_full_prefix_hidden_states()) ) - self._stage_deferred_prefix_cache_mm_outputs( - scheduler_output=scheduler_output, - multimodal_outputs=multimodal_outputs, - query_start_loc_cpu=query_start_loc_cpu, - ) - + if not needs_pooler_payload and prefix_cache_step_id is not None: + # No consumer for this step's merge: consume the step context by + # id (exactly-once contract). The cache write still lands. + self.omni_prefix_cache.discard_step(prefix_cache_step_id) + prefix_cache_step_id = None if self.omni_prefix_cache is None and needs_scheduled_hidden_payload and not audio_sparse_output: num_valid_tokens = min( int(scheduler_output.total_num_scheduled_tokens), @@ -1916,11 +1795,10 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex combined_hidden_states, combined_multimodal_outputs, ) = self._prepare_prefix_cache_pooler_payload_sources( - hidden_states=hidden_states, staged_hidden_states_cpu=staged_hidden_states_cpu, - multimodal_outputs=multimodal_outputs, - scheduler_output=scheduler_output, needs_scheduled_hidden_payload=needs_scheduled_hidden_payload, + step_id=prefix_cache_step_id, + req_ids=req_ids_output_copy, ) if combined_multimodal_outputs is None: flat_mm = flatten_payload(multimodal_outputs) if multimodal_outputs else multimodal_outputs @@ -2060,6 +1938,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex cudagraph_stats, multimodal_outputs, slot_mappings, # OMNI: unpack slot_mappings for drafter + prefix_cache_step_id, ) = self.execute_model_state self.execute_model_state = None @@ -2259,6 +2138,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex num_scheduled_tokens_np=num_scheduled_tokens_np, query_start_loc_cpu=query_start_loc_cpu, postprocess_already_applied=omni_postprocess_already_applied, + prefix_cache_step_id=prefix_cache_step_id, ) if not use_async_omni_output: diff --git a/vllm_omni/worker/gpu_generation_model_runner.py b/vllm_omni/worker/gpu_generation_model_runner.py index 4d80d3f8d..6d780925c 100644 --- a/vllm_omni/worker/gpu_generation_model_runner.py +++ b/vllm_omni/worker/gpu_generation_model_runner.py @@ -407,6 +407,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin cudagraph_stats, multimodal_outputs_raw, slot_mappings, # OMNI: unpack slot_mappings for upstream v1 API compatibility + _prefix_cache_step_id, # generation stages never save to the prefix cache ) = self.execute_model_state self.execute_model_state = None diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index 37d2ac122..48e9764f9 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/vllm_omni/worker/gpu_model_runner.py @@ -33,7 +33,8 @@ from vllm.v1.worker.gpu_input_batch import CachedRequestState from vllm.v1.worker.gpu_model_runner import GPUModelRunner, PerLayerAttnMetadata from vllm.v1.worker.ubatch_utils import maybe_create_ubatch_slices -from vllm_omni.core.prefix_cache import OmniTensorPrefixCache +from vllm_omni.core.prefix_cache import stage_prefix_cache_config +from vllm_omni.core.prefix_cache.runner_mixin import PrefixCacheRunnerMixin from vllm_omni.data_entry_keys import OmniPayload from vllm_omni.engine.serialization import deserialize_additional_information from vllm_omni.model_executor.layers.rotary_embedding.mrope import OmniMRotaryEmbedding as MRotaryEmbedding @@ -79,7 +80,7 @@ def _filter_mrope_kwargs_for_model(model: object, kwargs: dict[str, Any]) -> dic return {key: value for key, value in kwargs.items() if key in accepted} -class OmniGPUModelRunner(GPUModelRunner): +class OmniGPUModelRunner(PrefixCacheRunnerMixin, GPUModelRunner): intermediate_tensors: IntermediateTensors | None def __init__(self, *args, **kwargs): @@ -87,11 +88,16 @@ class OmniGPUModelRunner(GPUModelRunner): self.model_intermediate_buffer: dict[str, dict[str, Any]] = {} self._omni_num_scheduled_tokens_np: np.ndarray | None = None self._omni_last_model_output: object | None = None - # The Omni tensor prefix cache will be allocated - # when we initialize the metadata builders if enabled + # Omni prefix cache (hidden / per-token mm tensors reused on prefix + # hits). Built once, on the first real step after load_model: the + # config is staged at kv-cache init and cleared after construction. self.omni_prefix_cache = None + self._omni_prefix_cache_cfg = None self._sampled_token_ids_cpu_override = None self._omni_query_start_loc_model_kwarg = False + # Output-payload constant snapshotted once in load_model; the cache + # policy counterpart lives on PrefixCacheRunnerMixin. + self._pooler_payload_include_hidden_flag = True def _to_list(self, sampled_token_ids: torch.Tensor) -> list[list[int]]: override_fn = self._sampled_token_ids_cpu_override @@ -172,15 +178,20 @@ class OmniGPUModelRunner(GPUModelRunner): device=sm.device, ) - # Initialize the wrapper for both multimodal output tensors - # and for hidden states to be passed between stages - if self.cache_config.enable_prefix_caching: - self.omni_prefix_cache = OmniTensorPrefixCache( - num_blocks=kv_cache_config.num_blocks, - block_size=self.cache_config.block_size, - hidden_size=self.model_config.get_hidden_size(), - hs_dtype=self.dtype, - ) + # Stage the config; the manager is built on the first step once + # input_batch exists. The gate (pooling stage, kv_consumer, hybrid + # kv groups) is shared with the NPU runner. + cfg = stage_prefix_cache_config( + kv_cache_config=kv_cache_config, + cache_config=self.cache_config, + kv_transfer_config=getattr(self.vllm_config, "kv_transfer_config", None), + scheduler_config=self.scheduler_config, + model_config=self.model_config, + is_pooling_model=self.is_pooling_model, + speculative_config=self.speculative_config, + ) + if cfg is not None: + self._omni_prefix_cache_cfg = cfg @instrument(span_name="Loading (GPU)") def load_model(self, *args, **kwargs) -> None: @@ -192,6 +203,8 @@ class OmniGPUModelRunner(GPUModelRunner): if callable(candidate): override_fn = candidate self._sampled_token_ids_cpu_override = override_fn + self._snapshot_prefix_cache_model_policy(model) + self._pooler_payload_include_hidden_flag = bool(getattr(model, "omni_pooler_payload_include_hidden", True)) self._omni_query_start_loc_model_kwarg = bool(getattr(model, "supports_omni_query_start_loc", False)) self._maybe_enable_output_token_ids_for_model_sampler() self._init_talker_mtp() @@ -515,10 +528,6 @@ class OmniGPUModelRunner(GPUModelRunner): The SamplingMetadata is updated and copied to the GPU if there is a new/resumed/paused/finished request in the batch. """ - # Used for prefix cache - if self.omni_prefix_cache is not None: - self.omni_prefix_cache.reset_prefix_cached_new_req_ids() - # Remove finished requests from the cached states. # cleanup_finished_request lives on OmniConnectorModelRunnerMixin and # is only safe to call once init_omni_connectors() has finished @@ -535,8 +544,6 @@ 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"): @@ -600,13 +607,6 @@ class OmniGPUModelRunner(GPUModelRunner): reqs_to_add.append(req_state) continue - # Since this is the first time the request has been scheduled, - # num_computed_tokens > 0 means that we have a hit in prefix - # caching; mark it so that we can manage the hidden states - # later on as needed. - if self.omni_prefix_cache is not None and new_req_data.num_computed_tokens > 0: - self.omni_prefix_cache.add_prefix_cached_new_req_id(req_id) - sampling_params = new_req_data.sampling_params pooling_params = new_req_data.pooling_params