[wzhao18]

Summary

This PR fixes various undesirable prefix cache miss issues for hybrid full-attention + Mamba models, including inconsistent prefix-cache boundaries and incomplete Mooncake tail saves.

Many changes in this PR are referenced from #56615 and #57180. Given prefix caching for hybrid mamba is quite involved, the PR aims to take a more thorough approach to address the problems.

Problems addressed

Setup: FlashKDA prefill with DSpark/EAGLE, PMU/hash unit 128, GPU attention and Mamba block size 896, and scheduler LCM / Mooncake attention block size 7,168 tokens.

1. Incomplete external tail (Mooncake)

Sequence: 7,449-token prompt → GPU cache reset → resend the same prompt.

The replay on main recorded:

Cached GPU Mamba state:          7,296
Cached GPU attention proof:      7,424
Cached Mooncake Mamba state:     7,296
Cached Mooncake attention tail:  7,296
GPU local hit before reset:      7,296
Mooncake hit after reset:            0

Normal Mooncake saves reach the LCM boundary at 7,168. Its separate partial-tail save registers full-attention KV under the Mamba replay boundary, 7,296. EAGLE requires attention proof through 7,424, so the external lookup returns a complete miss despite retaining the Mamba state.

2. Attention lookup capped below available EAGLE proof

Sequence: 7,040-token prompt → resend the same prompt, without GPU reset.

The replay on main recorded:

Cached Mamba state:          6,784
Cached attention proof:     7,040
Current local hit:              0
Expected usable hit:         6,784

The Mamba checkpoint is derived from P - 1 and EAGLE's drop; attention caches through P. Both are available, but lookup caps the attention proof scan at 7,039 tokens (max_cache_hit_length = P - 1). The largest hash boundary within that limit is 6,912, so the cached attention entry at 7,040 is never inspected and the reusable prefix is missed.

3. Mamba checkpoint applies the resend limit twice with EAGLE

Sequence: 7,296-token prompt → resend the same prompt, without GPU reset.

Boundary calculation for the existing implementation (illustrative, not a measured replay):

Cached Mamba state:        7,040
Cached attention proof:   7,296
Current hit upper bound:  7,040
Expected usable hit:      7,168

The Mamba checkpoint calculation rounds down from P - 1, then subtracts another 128 tokens for EAGLE, while attention caches through P. This limits the hit to 7,040 tokens, although the cached attention proof at 7,296 supports a 7,168-token hit after EAGLE's drop. Since both the P - 1 limit and EAGLE's drop leave tokens to recompute on resend, this effectively applies the resend cap twice.

Proposed design

1. Use one resend-safe Mamba replay checkpoint. For prompt length P and checkpoint unit U:

   Without EAGLE: floor((P - 1) / U) × U
   With EAGLE:    max(floor(P / U) × U - U, 0)
   

   Without EAGLE, cap at P - 1 so an exact resend can reuse the checkpoint while leaving at least one token to recompute. With EAGLE, the block drop already leaves tokens to recompute, so no additional P - 1 cap is needed.

   Prompt extension could use a later checkpoint without the P - 1 cap. However, current KDA prefill backends support only one intermediate checkpoint without extra prefill splits. We therefore use the resend-safe boundary as the common Mamba replay checkpoint for both resend and extension.

2. Save the matching attention KV in Mooncake.

   Save the full-attention KV needed to reuse the Mamba checkpoint, including the extra proof required by EAGLE. Keep the source blocks referenced until the asynchronous save completes.

3. Allow lookup to inspect proof beyond the hit limit.

   Bound the attention proof scan by the available hashes, not the final hit limit. For a 7,296-token prompt, confirming a 7,168-token hit requires inspecting the attention hash at 7,296—beyond the prompt_len - 1 limit of 7,295. This is safe because we only inspect the proof; after EAGLE's drop, the returned hit is 7,168 and remains within the limit.

E2E agentic benchmark results

Kimi K3, 1P1D DCP8 + DCP8, DSpark k4, PMU 128.

Concurrency	Main TPS/GPU	PR TPS/GPU	Main prefix hit rate	PR prefix hit rate
64	13,762	13,963	95.763%	96.009%
80	15,163	15,253	94.817%	95.080%
96	15,973	16,669	94.413%	95.031%
112	13,976	15,668	92.432%	94.165%
128	11,371	15,754	90.466%	94.041%

Implementation review order

Implementation files only, in suggested review order:

1. vllm/v1/kv_cache_interface.py
   Unifies the resend-safe Mamba checkpoint calculation, avoiding the redundant P - 1 cap when EAGLE already drops a cache unit.

2. vllm/v1/core/kv_cache_utils.py
   Centralizes partial-hash-hit eligibility so checkpoint generation, local caching, and Mooncake use consistent alignment rules.

3. vllm/config/cache.py
   Adds resolved hash and Mamba checkpoint block sizes as runtime fields, excluded from the configuration hash.

4. vllm/v1/worker/gpu_model_runner.py and vllm/v1/worker/gpu/model_runner.py
   Populates those sizes before checkpoint metadata is built, including during CUDA-graph memory profiling, using the actual cache layout and DCP configuration.

5. vllm/model_executor/layers/mamba/checkpoint.py
   Makes the checkpoint builder use the resolved checkpoint unit, matching the scheduler and cache lookup granularity.

6. vllm/v1/core/sched/scheduler.py
   Aligns prefill splitting with the shared resend-safe checkpoint calculation while preserving support for checkpoints materialized within a prefill chunk.

7. vllm/v1/core/block_pool.py
   Allows the existing partial-entry registration path to register physical-block-aligned checkpoints too.

8. vllm/v1/core/single_type_kv_cache_manager.py
   Aligns attention and Mamba tail registration with the resend-safe boundary. Prevents transient intermediate states from being published as reusable checkpoints, while preserving retained checkpoints and shared junctions.

9. vllm/v1/core/kv_cache_coordinator.py
   Allows EAGLE attention lookup to inspect the required proof beyond the final hit limit, bounded by available hashes.

The following paths are under vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/:

10. data.py
    Adds completion and tail-publication metadata, separating the original prompt length from the expanded prefill horizon after preemption. Introduces a shared helper for identifying companion attention blocks.

11. coordinator.py
    Applies the same EAGLE proof-lookup rules to external caching and identifies the exact attention blocks needed to accompany the Mamba tail.

12. scheduler.py
    Emits tail-only saves when the prompt completes and keeps their source blocks referenced through asynchronous storage, including request cleanup.

13. worker.py
    Saves the Mamba tail together with its matching attention proof, including EAGLE's margin and gaps beyond normal saves. Handles proof publication even when no new Mamba offload accompanies the completion step.

---

Pull Request Checklist

• I used vLLM's /pr-checklist skill. (Mandatory for agents, optional for humans).

• AI assistance was used during the creation of this PR.

• Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.

• Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.

• Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.

• Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[ZJY0516]

I added some e2e tests in #59229, could you also add some?

[ZJY0516]

Do we need to unify this and prefix_match_unit?

[wzhao18]

Do you mean to get rid of --prefix-match-unit and make hash_block_size the prefix caching granularity for the tail? I think that would be reasonable. Ideally we could automatically enable the fine-grained caching without needing to specify prefix-match-unit. But I think if we want to do it, we should do it in a separate PR.

[wzhao18]

/ci run

[github-actions]

✅ Triggered Buildkite CI #93174 for commit 288b4c358cb3.

[ZJY0516]

IIUC, this PR has several parts

• code refactor like eagle_proof_margin
• more accurate position in get_mamba_prefill_checkpoint_position
• extend max_cache_hit_length in find_longest_cache_hit
• mooncake bugfix

Could you seperate this into 2 PR, one for core, one for mooncake.

Now core changes LGTM

[ZJY0516]

please add some comments here, why it's num_tokens when drop_eagle_block==True

[Dao007forever]

Suggestion: compute the partial tail once per job on the scheduler.

• partial_tail_block_ranges currently runs on both the scheduler (tail_attention_block_ids, for pinning) and the worker (for puts), using separately built coordinators. Computing it in _reference_save_blocks and register_finished_partial_tail, then sending it as ReqMeta.partial_tail, means pinning and puts always come from the same computation.
• The worker’s readiness condition (num_prompt_tokens <= completed_token_len) is true on every save after the prompt finishes. Each later save rebuilds the tail puts, and batch_is_exist drops them, wasting 1 check.

[Dao007forever]

nit: this looks repeated, can we extract to 1 function?

[wzhao18]

@Dao007forever Thanks for the review and feedback. Based on @ZJY0516 's suggestion, I will split this PR into two PR, one for the core and one for mooncake (which will come after the core part merged). I also see you have some clean up changes to the mooncake part, maybe you could help create the mooncake-part fix given you have more familiarity?

[Dao007forever]

Sounds good @wzhao18 ! I'll stack the Mooncake after #60533!

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @wzhao18.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork