[Bugfix] Fix hybrid attention cache miss due to eagle drop by wzhao18 · Pull Request #51295 · vllm-project/vllm

wzhao18 · 2026-08-06T18:15:36Z

Purpose

This PR fixes an unexpected prefix cache miss bug in hybrid sparse attention models.

Bug

Setup:

Kimi K3, TP8 B300
Block size: 1536
Prefix-match unit (PMU): 128
Cache retention interval: 0 (only cache at prompt boundary)

Turn 1: prompt_length=24576, cache_hit_length=0 [expected]
Turn 2: prompt_length=50000, cache_hit_length=23040 [expected]

Turn 1: prompt_length=24577, cache_hit_length=0 [expected]
Turn 2: prompt_length=50000, cache_hit_length=0 [unexpected - full miss]

The cause:

Mamba caches state at two prompt-boundary positions:

the last block-aligned position for prompt_length - 1; and
the last PMU-aligned position, registered as a partial block entry.

With prompt_length=24576 those are 23040 and 24576.
With prompt_length=24577 both collapse to 24576.

Full attention is cached up to 24576. The current prefix cache implementation requires Eagle to match one block (128 in this case) past the candidate and then dropping it. The 23040 mamba checkpoint in case 1 satisfies a cache hit where 23040 + 128 is cached in full attention. But in case 2 there is no mamba checkpoint that can satisfy it, and thus leads to a full miss.

Fix

The PR fixes by moving the mamba/SWA cache state position forward when using eagle so that it can match an extra unit/block of full attention cache to drop.

Test Plan

Test Result

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.

BEFORE SUBMITTING, PLEASE READ https://docs.vllm.ai/en/latest/contributing (anything written below this line will be removed by GitHub Actions)

wzhao18 opened the upstream version of the private commit we hand-trimmed: 'Simplify hybrid attention eagle retention and lookup and fix unexpected cache miss'. Its reproducer is our configuration verbatim -- Kimi K3, TP8 B300, block size 1536, prefix-match unit 128, retention interval 0. It carries the two files we could not port. The refactor in the title is what deleted MooncakeStoreCoordinator.eagle_attn_group_indices, the symbol whose absence in our nightly forced the trim; get_prefix_replay_checkpoint and eagle_replay_tokens replace it. Under the trimmed d87 the DSpark arms stopped hitting the kv_cache_manager assert and served real traffic (61431494: c32 858 requests at 5,128 tok/s/GPU), but still lost EngineCore to a sporadic CUDA illegal memory access -- the scheduler and cache manager were fixed while the Mooncake connector kept doing its own un-rewound lookup. That is the gap this closes. Applies with zero fuzz to our pinned nightly cb8104839c and with offsets only to the latest nightly 821717118f. Tests stripped; #51113 deliberately not stacked (it applies on top only with fuzz, and #51295 supersedes that region).

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.

mergify · 2026-08-07T02:41:52Z

Hi @wzhao18, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

ZJY0516

IIUC, this will hurt the latency

Let's say block size = 4, PMU = 2, prompt length = 9

  main  
  0 ---- 8 | 8 - 9

  PR
  0 -- 6 | 6 -- 8 | 8 - 9

We'll need one more forward

mergify · 2026-08-12T13:07:53Z

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

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.

mergify · 2026-08-16T04:25:26Z

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

akshaver · 2026-08-17T19:49:16Z

@wzhao18 I left an MR against this to cover to another failure mode: wzhao18#2

wzhao18 · 2026-08-17T20:29:04Z

Hi @akshaver Thanks for the PR! Based on discussion with @ZJY0516, there may be some active work to refactor the eagle block drop logic to simplify things up. The current logic is very complicated and requires a lot of work to handle things properly. Therefore, I plan to wait for the refactor to land before pursuing any bug fixes.

@ZJY0516 Can you confirm if that is the plan and is there an ETA for the refactor?

Under the EAGLE/MTP last-unit drop, a consumer resumes one unit below the prompt's last reusable hash boundary. Materialize and retain the Mamba state at that exact boundary, and share the drop-unit calculation between cache publication, scheduling, and lookup. For hybrid DCP layouts, propagate the validated hash alignment to each cache manager. Otherwise sparse Mamba retention uses the much larger scheduler LCM and discards an otherwise valid replay checkpoint. This is the DCP-specific gap not covered by vllm-project#51295. Co-authored-by: wzhao18 <wzhao18.sz@gmail.com> Co-authored-by: andyluo7 <andy.luo@amd.com> Assisted-by: OpenAI Codex Signed-off-by: Yichao Zhu <Yichao.Zhu@amd.com>

Under the EAGLE/MTP last-unit drop, a consumer resumes one unit below the prompt's last reusable hash boundary. Materialize and retain the Mamba state at that exact boundary, and share the drop-unit calculation between cache publication, scheduling, and lookup. Separately, propagate the validated hash alignment to each cache manager. `find_longest_cache_hit` reports hits at `_cache_hit_alignment_tokens`, but `cache_blocks` aligned retention to `scheduler_block_size`, so the producer retained a position the consumer never asks for. This is not EAGLE-specific: it floors every reachable boundary, including the plain `num_prompt_tokens - 1` replay boundary, and it only surfaces once the two differ -- i.e. under a DCP-scaled scheduler LCM with sparse retention on. This is the DCP-specific gap not covered by vllm-project#51295. Coverage is split by what each case actually proves: - the DCP row of the EAGLE test and the new drafter-free test are the behavioral regressions (both drop to a full miss without the fix); - the non-DCP EAGLE row is the non-regression side, where lowering the alignment must change nothing; - a dedicated test pins the alignment hand-off itself, including the disabled branch that keeps `SlidingWindowManager` from being handed an alignment finer than its block size. Co-authored-by: wzhao18 <wzhao18.sz@gmail.com> Co-authored-by: andyluo7 <andy.luo@amd.com> Assisted-by: OpenAI Codex Signed-off-by: Yichao Zhu <Yichao.Zhu@amd.com>

…ibling resumes Second of the two EAGLE + `--mamba-cache-mode align` prefix-cache defects pinned by vllm-project#52371. Full attention hits where it holds a key; EAGLE prunes one hash unit off that candidate and drops it. Mamba materializes state only on its own block grid, so nothing exists at the pruned position and the hit floors back a whole Mamba block -- on a 4,416-token block a 64-token drop costs 4,416 tokens of reuse, per request. vllm-project#51295 (commit 1 here) covers the case where the shared prefix runs to the end of the producer's prompt. When it ends earlier -- a system prompt followed by a per-request suffix, the deployed shape -- the producer's tail sits over tokens no sibling shares, so a check-point there is unreachable. The sibling stops at the last shared boundary and EAGLE drops one hash unit below it; that position is `request.shared_prefix_boundary`, already recorded by `get_computed_blocks`. So: the scheduler ends a chunk at the junction as observed rather than floored to the block grid, which rounds the resume point away, and `MambaManager` accepts the junction in addition to the prompt tail. The two must agree -- the junction stop is the earliest mandatory stop, so it REPLACES the block-boundary stop, and a junction the manager then refuses leaves less cached than not splitting at all. `KVCacheManager` therefore computes one predicate that both sides read, rather than each deriving its own: - EAGLE must be annotated on the model. The manager takes the drop per KV cache group while the scheduler flag is model-wide, so deriving them separately lets the split fire where the manager will not check-point. - Fine-grained partial hash hits must be on, or nothing can look the entry up. - No multi-module MTP. `cache_blocks` then hands the manager `num_computed - num_reprefillable`, not the chunk end, so the position it would publish is not the one it holds. Registration is bounded to the prompt chunk being computed: during decode the target is the running state block, mutated in place, which equals what its key promises only after that step's forward. A junction past the prompt falls back to the stock block-floored stop rather than being dropped. The manager writes nothing past the prompt, but stock vLLM still check-points at the block boundary, so zeroing the stop lost a snapshot a resumed request's siblings could have reused. Off by default behind `--enable-mamba-fine-grained-prefix-cache` (`CacheConfig.enable_mamba_fine_grained_prefix_cache`), listed in `compute_hash`'s ignored factors alongside the other prefix-cache knobs so it does not perturb the compiled-graph cache. Signed-off-by: Adam Shaver <ashaver@nvidia.com> Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

@wzhao18

Reverts the `mooncake/store/coordinator.py` hunk that commit 1 (vllm-project#51295) carries, taking the reviewer's offer: "I don't have an issue with leaving disagg support for a later PR altogether." Done as a commit on top rather than by editing commit 1, so commit 1 stays byte-identical to vllm-project#51295 (@wzhao18) and this PR's own opinion is visible as its own change. After this the PR touches nothing under `vllm/distributed/` at all. What this defers: the Mooncake offload STORE mirrors the engine's sparse retention (it is the only other caller of `SingleTypeKVCacheManager.reachable_block_mask`), so while EAGLE moves the engine's replay boundary the store still retains at `num_prompt_tokens - 1`. The consequence is missed reuse for store users under EAGLE + align, not incorrect output: the store is a cache, and the engine does not depend on its retention for correctness. What is NOT deferred because it is not the same thing: the `num_prompt_tokens - 1` in `nixl/base_scheduler.py` and `mooncake_connector.py::_get_remote_prefill_token_count` is the D-side half of a P/D handoff whose P-side half pops exactly one token (`_truncate_mamba_request_for_prefill`). That offset is hard-coded on both nodes and never sent on the wire; shifting it desynchronises them -- an SSM slot-pairing assert under `mamba_cache_mode="all"`, silent state corruption otherwise. Note the two Mooncake files differ in kind: `mooncake/store/` is the prefix-cache store (a peer of the engine), `mooncake_connector.py` is the P/D connector (a peer of NIXL). The follow-up wants the store change plus, if the connectors are in scope, a `kv_transfer_params` field so P and D cannot diverge, a rework of the slot pairing in `nixl/base_worker.py`, and validation on `tests/v1/kv_connector/nixl_integration/run_mamba_prefix_cache_test.sh`, which needs two GPUs. None of that is testable by this PR's single-engine suite. Signed-off-by: Adam Shaver <ashaver@nvidia.com> Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

wzhao18 · 2026-09-14T17:48:21Z

Superseded by #53945

mergify Bot added the kv-connector label Aug 6, 2026

wzhao18 force-pushed the wzhao/prefix-replay-checkpoints branch 2 times, most recently from a9330c3 to 943bee2 Compare August 6, 2026 22:21

wzhao18 changed the title ~~Simplify hybrid attention eagle retention and lookup and fix unexpected cache miss~~ [Bugfix] Cache Mamba states where the EAGLE drop can reach Aug 6, 2026

mergify Bot added the bug Something isn't working label Aug 6, 2026

wzhao18 closed this Aug 6, 2026

wzhao18 reopened this Aug 6, 2026

wzhao18 marked this pull request as ready for review August 6, 2026 23:38

wzhao18 requested review from ApostaC, WoosukKwon, alexm-redhat, heheda12345, ivanium, njhill, orozery, robertgshaw2-redhat and ywang96 as code owners August 6, 2026 23:38

claude Bot reviewed Aug 6, 2026

View reviewed changes

wzhao18 force-pushed the wzhao/prefix-replay-checkpoints branch 3 times, most recently from bda78df to 3e828eb Compare August 7, 2026 02:20

wzhao18 changed the title ~~[Bugfix] Cache Mamba states where the EAGLE drop can reach~~ [Bugfix] Cache Sparse KV cache where the EAGLE drop can reach Aug 7, 2026

ZJY0516 added the verified Run pre-commit for new contributors without triggering other tests label Aug 7, 2026

ZJY0516 reviewed Aug 7, 2026

View reviewed changes

wzhao18 marked this pull request as draft August 7, 2026 03:29

wzhao18 changed the title ~~[Bugfix] Cache Sparse KV cache where the EAGLE drop can reach~~ [Bugfix] Fix hybrid attention cache miss due to eagle drop Aug 7, 2026

wzhao18 mentioned this pull request Aug 7, 2026

[Bug]: Hybrid sparse attention + Eagle results in full cache miss for certain prompt lengths #51441

Open

1 task

wzhao18 force-pushed the wzhao/prefix-replay-checkpoints branch from 86bd0d2 to 0d93dac Compare August 13, 2026 21:17

wzhao18 marked this pull request as ready for review August 13, 2026 21:17

wzhao18 requested review from NickLucche and xuechendi as code owners August 13, 2026 21:17

claude Bot reviewed Aug 13, 2026

View reviewed changes

mergify Bot removed the needs-rebase label Aug 13, 2026

wzhao18 marked this pull request as draft August 14, 2026 16:48

mergify Bot added the needs-rebase label Aug 16, 2026

akshaver mentioned this pull request Aug 17, 2026

Shaver/pr51295 plus block grid wzhao18/vllm#2

Closed

mergify Bot added scheduler kv-cache-manager labels Aug 19, 2026

This was referenced Aug 26, 2026

[Bugfix][Spec Decode] Cache the Mamba state at the block-grid position of EAGLE resume #53923

Closed

[Bugfix][Spec Decode] Cache the Mamba state at the block-grid position of EAGLE resume #53945

Merged

puririshi98 mentioned this pull request Aug 27, 2026

[Perf][Spec Decode] dflash/dspark: prefix-cache last-block-drop exemption + decode-path overhead reductions #54092

Closed

GirasoleY mentioned this pull request Aug 29, 2026

[ROCm][DSpark][DCP] Serve prefix cache hits under DCP for Kimi-K3 #53598

Merged

YukioZzz mentioned this pull request Aug 31, 2026

[Bugfix][DCP] Handle hybrid cache geometry in offload recovery #53917

Closed

This was referenced Sep 2, 2026

[Bugfix][Spec Decode] EAGLE replay boundary in the offload store; one definition of the Mamba P/D handoff #55036

Closed

[Bugfix][Spec Decode] EAGLE replay boundary in the offload store; one definition of the Mamba P/D handoff akshaver/vllm#1

Open

wangyicong52 mentioned this pull request Sep 11, 2026

[KV Cache] Preserve EAGLE SWA replay blocks for external APC #56449

Draft

starkwj mentioned this pull request Sep 12, 2026

[Bugfix][Spec Decode][Mooncake] Fix MooncakeStore cache misses from EAGLE/DSpark-drop-induced Mamba/a… #56615

Open

4 tasks

wzhao18 closed this Sep 14, 2026

Navigation Menu

Uh oh!

Uh oh!

[Bugfix] Fix hybrid attention cache miss due to eagle drop - #51295

[Bugfix] Fix hybrid attention cache miss due to eagle drop#51295
wzhao18 wants to merge 1 commit into
vllm-project:mainfrom
wzhao18:wzhao/prefix-replay-checkpoints

wzhao18 commented Aug 6, 2026 •

edited

Loading

Uh oh!

claude Bot left a comment

Uh oh!

mergify Bot commented Aug 7, 2026

Uh oh!

ZJY0516 left a comment

Uh oh!

mergify Bot commented Aug 12, 2026

Uh oh!

claude Bot left a comment

Uh oh!

mergify Bot commented Aug 16, 2026

Uh oh!

akshaver commented Aug 17, 2026 •

edited

Loading

Uh oh!

wzhao18 commented Aug 17, 2026 •

edited

Loading

Uh oh!

wzhao18 commented Sep 14, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Uh oh!

Conversation

wzhao18 commented Aug 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Purpose

Bug

Fix

Test Plan

Test Result

Uh oh!

claude Bot left a comment

Choose a reason for hiding this comment

Claude Code Review

Uh oh!

mergify Bot commented Aug 7, 2026

Uh oh!

ZJY0516 left a comment

Choose a reason for hiding this comment

Uh oh!

mergify Bot commented Aug 12, 2026

Uh oh!

claude Bot left a comment

Choose a reason for hiding this comment

Claude Code Review

Uh oh!

mergify Bot commented Aug 16, 2026

Uh oh!

akshaver commented Aug 17, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

wzhao18 commented Aug 17, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

wzhao18 commented Sep 14, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

wzhao18 commented Aug 6, 2026 •

edited

Loading

akshaver commented Aug 17, 2026 •

edited

Loading

wzhao18 commented Aug 17, 2026 •

edited

Loading