[2/N][Core] support partial prefix cache hit for hybrid model

[ZJY0516]

Purpose

This PR builds on the partial prefix-cache primitives from #45939 and completes fine-grained prefix-cache hits for aligned full-attention + Mamba hybrid models. The overall design is described in RFC #45702.

The PR adds:

• a public prefix_match_unit that selects prefix-match granularity independently from physical KV-cache block sizes;
• fine-grained lookup in the full-attention and Mamba cache managers;
• token-accurate convergence in the hybrid coordinator;
• scheduler boundaries needed to materialize and resume partial Mamba states;
• copy-on-write for requests that continue from a shared partial block;
• worker-side block copies with scheduler-side lifetime fencing.

It e2e works now

vllm serve Qwen/Qwen3.5-35B-A3B-FP8 --enable-prefix-caching --mamba-cache-mode align --prefix-match-unit 16

Use 2 sequential request to simulate multi-turn:

• 10000 in, 1 out
• 10000 + 200 in, 1 out

w/ partial cache hit
"first_latency_sec": 0.3047793720033951, "second_latency_sec": 0.09551339800236747

w/o partial cache hit
"first_latency_sec": 0.2865299719851464, "second_latency_sec": 0.13373498094733804

Design

Match unit and cache-manager hash views

Request.block_hashes is computed at prefix_match_unit granularity. Each KV-cache group continues to allocate its existing physical block size.

For example:

prefix_match_unit          = 2
full-attention block_size  = 4
Mamba block_size           = 8

request hashes: H2, H4, H6, H8, H10, ...

Managers using physical-block lookup receive a grouped view such as [H4, H8, ...]. Managers supporting fine-grained lookup keep the original hash list and may probe H2 or H6 inside a physical block.

The configured match unit must divide every cache-group block size. If it is not specified, vLLM derives it from the GCD of the resolved group block sizes.

Explicit token-length cache hits

Cache-manager lookup now returns both the physical blocks and the exact number of matched tokens:

(blocks_by_group, hit_length)

The separate length is required because a hit at H6 includes the physical block [4, 8), but only tokens [4, 6) in that block are valid.

FullAttentionManager first matches complete physical blocks from the start of the request, then probes hash boundaries inside the next block from longest to shortest. MambaManager keeps its existing right-to-left state lookup, but can now search at match-unit boundaries and return the physical block containing the matched recurrent state.

Managers without fine-grained support continue to return block-aligned hit lengths through the same API.

Hybrid-coordinator convergence

The hybrid coordinator tracks a token hit length independently from each group's block list. It repeatedly reduces the candidate length until every cache group can cover the same logical prefix.

For example:

full attention: 6-token hit, 2 physical blocks
Mamba:          4-token hit, 1 physical block
------------------------------------------------
coordinator:    4-token hit

Full-attention results are downward-closed, so an existing result can be trimmed without another lookup. Physical block lists are trimmed with ceiling division so the block containing a partial tail is retained.

EAGLE also uses the explicit token length. In fine-grained mode it drops one match unit, rather than one physical block, before the result is returned.

The explicit-length contract is also propagated to the Mooncake store coordinator. Its block-pool adapter exposes hash_block_size so it can reuse the same manager lookup interface.

Creating partial tail entries

Full attention registers the final hash-aligned prompt boundary once the containing physical block is available.

Aligned Mamba requires the recurrent state from exactly that boundary. The scheduler therefore stops once at the final match-unit boundary before the end of a partial prompt block. MambaManager registers that state only when the step lands on this boundary; it does not retain every intermediate state inside the physical block.

When execution resumes from a partial hit, the first chunk stops at the next physical block boundary. Subsequent chunks return to the existing block-aligned scheduling path.

Copy-on-write

A request cannot append to a physical block while the prefix cache exposes an entry for an earlier boundary in that block. This PR uses copy-on-write, with two allocation paths.

For a new request hitting a partial entry, the request is redirected to a private block:

request table: shared source -> private destination
copy:          shared source -> private destination

The shared cached block remains unchanged and the request continues from the private copy.

An already-running aligned-Mamba request may register a partial entry for its current state and then continue in the next step. Its worker block table is append-only, so that request must remain on the original block. Instead, the cache entry is moved to a new block:

cache hashes: original request block -> cache-owned destination
copy:         original request block -> cache-owned destination
request:      continues writing the original block

BlockPool.move_block_hashes() re-points all hashes without emitting a remove/store event. The destination is excluded from same-step lookup until the copy has populated it.

Copy execution and lifetime

Managers queue (source, destination) block pairs. KVCacheManager converts them to block-ID copies on SchedulerOutput, and workers execute them after zeroing newly allocated blocks and before model forward:

zero new blocks -> copy partial blocks -> model forward

The worker copies complete block-major storage pages. Storage aliases across layer views are deduplicated, so each backing allocation is copied once.

Both copy endpoints are retained until the non-empty scheduler step carrying the copy has completed. Their release uses the existing deferred-free queue, preventing either block from being recycled while asynchronous execution is still reading or writing it.

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 55de9c44e1

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 7327f58bea

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: cb2b1592cd

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[Sunshine-llh]

Purpose

This PR builds on the partial prefix-cache primitives from #45939 and completes fine-grained prefix-cache hits for aligned full-attention + Mamba hybrid models. The overall design is described in RFC #45702.

The PR adds:

• a public prefix_match_unit that selects prefix-match granularity independently from physical KV-cache block sizes;
• fine-grained lookup in the full-attention and Mamba cache managers;
• token-accurate convergence in the hybrid coordinator;
• scheduler boundaries needed to materialize and resume partial Mamba states;
• copy-on-write for requests that continue from a shared partial block;
• worker-side block copies with scheduler-side lifetime fencing.

It e2e works now

vllm serve Qwen/Qwen3.5-35B-A3B-FP8 --enable-prefix-caching --mamba-cache-mode align --prefix-match-unit 16

Use 2 sequential request to simulate multi-turn:

• 10000 in, 1 out
• 10000 + 200 in, 1 out

w/ partial cache hit
"first_latency_sec": 0.3047793720033951, "second_latency_sec": 0.09551339800236747

w/o partial cache hit
"first_latency_sec": 0.2865299719851464, "second_latency_sec": 0.13373498094733804

Design

Match unit and cache-manager hash views

Request.block_hashes is computed at prefix_match_unit granularity. Each KV-cache group continues to allocate its existing physical block size.

For example:

prefix_match_unit          = 2
full-attention block_size  = 4
Mamba block_size           = 8

request hashes: H2, H4, H6, H8, H10, ...

Managers using physical-block lookup receive a grouped view such as [H4, H8, ...]. Managers supporting fine-grained lookup keep the original hash list and may probe H2 or H6 inside a physical block.

The configured match unit must divide every cache-group block size. If it is not specified, vLLM derives it from the GCD of the resolved group block sizes.

Explicit token-length cache hits

Cache-manager lookup now returns both the physical blocks and the exact number of matched tokens:

(blocks_by_group, hit_length)

The separate length is required because a hit at H6 includes the physical block [4, 8), but only tokens [4, 6) in that block are valid.

FullAttentionManager first matches complete physical blocks from the start of the request, then probes hash boundaries inside the next block from longest to shortest. MambaManager keeps its existing right-to-left state lookup, but can now search at match-unit boundaries and return the physical block containing the matched recurrent state.

Managers without fine-grained support continue to return block-aligned hit lengths through the same API.

Hybrid-coordinator convergence

The hybrid coordinator tracks a token hit length independently from each group's block list. It repeatedly reduces the candidate length until every cache group can cover the same logical prefix.

For example:

full attention: 6-token hit, 2 physical blocks
Mamba:          4-token hit, 1 physical block
------------------------------------------------
coordinator:    4-token hit

Full-attention results are downward-closed, so an existing result can be trimmed without another lookup. Physical block lists are trimmed with ceiling division so the block containing a partial tail is retained.

EAGLE also uses the explicit token length. In fine-grained mode it drops one match unit, rather than one physical block, before the result is returned.

The explicit-length contract is also propagated to the Mooncake store coordinator. Its block-pool adapter exposes hash_block_size so it can reuse the same manager lookup interface.

Creating partial tail entries

Full attention registers the final hash-aligned prompt boundary once the containing physical block is available.

Aligned Mamba requires the recurrent state from exactly that boundary. The scheduler therefore stops once at the final match-unit boundary before the end of a partial prompt block. MambaManager registers that state only when the step lands on this boundary; it does not retain every intermediate state inside the physical block.

When execution resumes from a partial hit, the first chunk stops at the next physical block boundary. Subsequent chunks return to the existing block-aligned scheduling path.

Copy-on-write

A request cannot append to a physical block while the prefix cache exposes an entry for an earlier boundary in that block. This PR uses copy-on-write, with two allocation paths.

For a new request hitting a partial entry, the request is redirected to a private block:

request table: shared source -> private destination
copy:          shared source -> private destination

The shared cached block remains unchanged and the request continues from the private copy.

An already-running aligned-Mamba request may register a partial entry for its current state and then continue in the next step. Its worker block table is append-only, so that request must remain on the original block. Instead, the cache entry is moved to a new block:

cache hashes: original request block -> cache-owned destination
copy:         original request block -> cache-owned destination
request:      continues writing the original block

BlockPool.move_block_hashes() re-points all hashes without emitting a remove/store event. The destination is excluded from same-step lookup until the copy has populated it.

Copy execution and lifetime

Managers queue (source, destination) block pairs. KVCacheManager converts them to block-ID copies on SchedulerOutput, and workers execute them after zeroing newly allocated blocks and before model forward:

zero new blocks -> copy partial blocks -> model forward

The worker copies complete block-major storage pages. Storage aliases across layer views are deduplicated, so each backing allocation is copied once.

Both copy endpoints are retained until the non-empty scheduler step carrying the copy has completed. Their release uses the existing deferred-free queue, preventing either block from being recycled while asynchronous execution is still reading or writing it.

Essential Elements of an Effective PR Description Checklist

Thanks again for the partial-prefix-cache work. We integrated the approach and confirmed that it enables partial cache reuse for aligned full-attention + Mamba hybrid models.

However, we still observe a relatively low prefix-cache hit rate on our real workload.

Our understanding is that each request retains the final prefix_match_unit-aligned boundary of its prompt, rather than retaining every intermediate Mamba checkpoint inside the
physical Mamba block. For example, with:

• prefix_match_unit = 16
• a request prompt length of 1946

the cache can retain a checkpoint around token 1936, but not necessarily checkpoints such as H560.

This is problematic for our workload because many requests share a common system/context prefix near the beginning, then diverge with different user content:

Request A: [shared prefix: 560 tokens] + [unique content A]  -> total 1946 tokens
Request B: [shared prefix: 560 tokens] + [unique content B]

Even though the first 560 tokens are identical, Request B may not be able to reuse the prefix cached by Request A because the cached partial Mamba state is associated with the
final aligned boundary (H1936), not the internal shared boundary (H560).

We understand that materializing a Mamba state at every prefix_match_unit boundary would have significant memory, copy, and metadata overhead. Still, do you see room for further
optimization here? For example:

- retaining multiple selected intermediate checkpoints within a physical Mamba block;
- adaptively materializing checkpoints at frequently reused prefix boundaries;
- allowing applications to mark known reusable boundaries, such as system prompts or document-prefix boundaries;
- using a coarser checkpoint interval than prefix_match_unit, while still providing better coverage than only the final prompt boundary.

Any guidance on the intended next optimization direction would be very helpful. Thank you again for the work on this PR; it already solves an important limitation for partial
prefix reuse in hybrid Mamba models.