[natsala13]

Purpose

Fix repeated prompts missing a saved Mamba prefix-cache checkpoint with EAGLE/MTP.

There are two mismatches:

1. Lookup: the coordinator caps the proof scan at the reusable-token limit (request.num_tokens - 1). Attention lookup then drops another hash unit for speculative replay. With a 16,000-token prompt and 16-token hashes, this can stop at 15,968 and miss the saved Mamba checkpoint at 15,984. Allowing the scan through token 16,000 lets the existing drop land on 15,984.
2. Registration: when only attention groups are marked as draft groups, Mamba partial-tail registration uses its per-group use_eagle flag while the scheduler uses the global drop_eagle_checkpoint_block flag. The saved state and its registered boundary can disagree.

This change gives attention lookup its extra proof margin, bounded by the request's available hashes, in both reconciled and per-group connector lookup. Mamba receives no extra margin because its finder does not drop a unit. Mamba tail registration uses the scheduler's global flag. The returned reusable hit stays within the caller's limit, and the speculative proof drop remains enabled.

Attribution and related work

Builds on work by Brian Yu, with expanded regression coverage.

Related upstream PR #53802 addresses the boundary mismatch by moving checkpoint registration to positions reachable under the existing lookup cap. This patch instead preserves the saved replay checkpoint and corrects the proof-scan cap, and aligns Mamba registration with the global draft setting. This proposes an alternative at the lookup boundary, rather than the checkpoint-placement approach in #53802; the overlap is explicit so maintainers can choose or reconcile the approaches.

AI assistance was used for implementation, review, and test execution. This draft is for human review; it does not claim the submitter has already reviewed every changed line or personally rerun the tests.

Test Plan

The regression extends the existing scheduler/cache-manager test file. A newer retention test is also updated: a 128-token resend should recover the saved 96-token checkpoint, rather than codifying the old 64-token fallback. Its checks that both retained Mamba states exist are preserved. It covers both all-group fallback and attention-only draft marking, fine/coarse hashes, aligned/unaligned prompts, repeated/extended prompts, changed proof suffixes, and reusable-token limits in local and connector bookkeeping paths.

To reproduce red/green in an isolated checkout of this PR, using the repository's prepared vLLM test environment:

suite=(tests/v1/core/prefix_cache/test_mamba_eagle_resume_checkpoint.py
       tests/v1/core/test_prefix_caching.py::test_hybrid_mamba_retention_mtp_resend_of_aligned_prompt)
# Keep the new tests, but restore only the two runtime files to the tested PR base.
git restore --source=f0fbc98c488b40a07d3b269aeba2bda3c4c3b03b --worktree -- \
  vllm/v1/core/kv_cache_coordinator.py \
  vllm/v1/core/single_type_kv_cache_manager.py
.venv/bin/python -m pytest "${suite[@]}" -q --tb=short
# Restore the fix and run exactly the same tests.
git restore --source=HEAD --worktree -- \
  vllm/v1/core/kv_cache_coordinator.py \
  vllm/v1/core/single_type_kv_cache_manager.py
.venv/bin/python -m pytest "${suite[@]}" -q --tb=short
.venv/bin/python -m pytest \
  tests/v1/core/prefix_cache \
  tests/v1/core/test_prefix_caching.py \
  tests/v1/core/test_single_type_kv_cache_manager.py \
  tests/v1/core/test_mamba_align_chunk_split.py -q --tb=short

Test Result

The tested base is upstream commit f0fbc98c488b40a07d3b269aeba2bda3c4c3b03b; head is c6a1b11096. Commit metadata was updated after testing; the code and tests are byte-for-byte identical to the tested version. This PR targets vllm-project/vllm:main. The results below are for the exact tested base and code, not a fresh run after rebasing onto the latest upstream main. Final-base red/green validation completed successfully in 3m46s with a five-minute walltime limit:

Check on tested base	Original runtime + final tests	Fixed runtime + same tests
Focused suite plus updated aligned-resend test	8 failed, 24 passed (10.41s)	32 passed (8.54s)
Related cache/scheduler suites	—	314 passed (68.24s), including all 32 above

All four tested files match the Git blobs in the pushed head. The eight failures are cache-length assertions. The other cases protect behavior that should already pass; they are not claimed as eight additional independent bugs.

The first wider-suite run found an existing test from #54713 that asserted the old 64-token fallback. Its expected hit and block counts now require the recovered 96-token checkpoint. The final run above includes that update.

Before porting onto fork main, the original three-file patch on the reproduction base 2a02f6efe319c885e3ccbcecde402e0028f9ec1e produced:

Check	Original runtime + new tests	Fixed runtime + same tests
Focused suite	7 failed, 24 passed	31 passed
Related cache/scheduler suites	—	255 passed, including the focused suite

The failures are cache-length assertions, not import or collection errors. Seven cases demonstrate the regression; the remaining positive, negative, and boundary cases guard behavior that must continue to work. All applicable pre-commit hooks, including Ruff and type checks, pass on the final PR files.

Real-model reproduction

On the original reproduction base, ran real Qwen/Qwen3.5-0.8B-Base weights with BF16, MTP=3, prefix matching at 16 tokens, Mamba align mode, fine-grained caching, eager synchronous execution, one request at a time, and greedy generation of 32 tokens. Both runs compare cold and warm outputs across repeated, extended, and final-token-changed prompts at lengths 16,000 and 16,001.

Measurement	Original runtime	Fixed runtime
Repeated 16,000-token prompt: cached tokens	15,120	15,984
Prompt tokens recomputed	880	16
Cold/warm output token IDs	6/6 comparisons equal, 32 tokens each	6/6 comparisons equal, 32 tokens each
Job elapsed time	3:19 (five-minute limit)	3:18

The original runner explicitly required a lost repeated-prompt tail and printed BUG REPRODUCED. This isolates the lookup-cap defect; attention-only draft registration is demonstrated by the unit tests. The result is 864 fewer recomputed prompt tokens in this example, not a measured throughput claim.

Raw test logs are retained separately and are not included in this PR.

These are focused cache correctness and small real-model checks. Model-task accuracy, serving throughput, and real distributed KV transfer still need separate qualification.

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

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

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