[Core] Refactor omni prefix cache into Manager/Controller (#6654)

Signed-off-by: ZhengWG <zwg0606@gmail.com>
Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

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

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: <hidden states vector of len 2 corresponding to 'The'>
 1: <hidden states vector of len 2 corresponding to 'quick'>
 2: <hidden states vector of len 2 corresponding to 'brown'>
@@ -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 <empty> ]
 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)

tests/core/test_prefix_cache_async_write.py

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

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

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

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
 
 

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

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
 

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

vllm_omni/core/prefix_cache.py

@@ -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]

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",
+]

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)

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

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)

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,
+        )

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)

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

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
                 )

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

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

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:

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
 

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