[BugFix] Preserve sparse hybrid replay boundaries under EAGLE/MTP

[logprobz]

Purpose

This draft builds directly on #556. It preserves both
original commits unchanged, so @tobymao remains their author. Thank you to Toby
Mao for finding and implementing the original Mamba sparse-retention fix.

The follow-up implements findings from
logprobz's review comment
and two independent Claude Opus 5 reviews:

• Retain both positions an EAGLE/MTP attention lookup can make reachable: the
  aligned boundary and its predecessor. A manager capability now identifies
  lookups that actually drop a proof block, including full attention, MLA, and
  SWA, while excluding Mamba and unsupported local-attention managers. This
  refines the bool(eagle_group_ids) predicate introduced by [BugFix] Mamba sparse retention keeps a state below each boundary under EAGLE/MTP #556.
• Pass the coordinator's _cache_hit_alignment_tokens into sparse-retention
  masks. For Mamba, retain both the fine lookup position and a
  scheduler-aligned state-materialization fallback, so a mathematically valid
  fine boundary does not displace the only state the scheduler produced.
• Apply the same boundary and alignment rules to Mooncake store masks. The
  connector now shares the engine's sparse-retention input resolver instead of
  offloading states that local lookup cannot consume. It also mirrors the
  engine's compatibility gate, so mixed SWA and Mamba layouts fall back to
  block-aligned hits instead of asserting during store.

The PR targets dev/jovian-judgement, so its diff contains Toby's two original
commits from #556 followed by four logprobz follow-up commits.

The BlockStored sparse-event concern from the linked review remains out of
scope because it changes connector event semantics. This PR does not file a
separate issue; that work remains tracked in the review comment pending an
independent issue or PR.

Related work

• [BugFix] Retain both replay boundaries so an EAGLE resend of a block-aligned prompt still hits vllm-project/vllm#54713 is the upstream form of [BugFix] Mamba sparse retention keeps a state below each boundary under EAGLE/MTP #556.
• [Bugfix] Register hybrid prefix-cache boundaries at hittable positions vllm-project/vllm#53802 is related but not a duplicate. It approaches the
  symptom from registration: it floors the scheduler tail split and
  _cache_partial_tail_block from num_tokens - 1, and adds an SWA tail at the
  fine hash boundary inside reachable_block_mask. This draft leaves
  registration unchanged. It passes the coordinator's actual lookup alignment
  into cache_blocks and retains EAGLE-reachable positions consistently in the
  engine and Mooncake store masks.

Test Plan

The pytest runs use the pinned r22 container's isolated /opt/venv after adding
the tblib test dependency, which the image does not include:

/opt/venv/bin/python -m pytest -q -p no:cacheprovider \
  tests/v1/core/test_prefix_caching.py \
  tests/v1/core/prefix_cache \
  tests/v1/core/test_mamba_align_chunk_split.py \
  tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py \
  tests/v1/kv_connector/unit/test_mooncake_store_worker.py \
  tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py

The lint runs use an isolated uv environment and the repository's pinned hook
configuration:

pre-commit run --files \
  tests/v1/core/test_prefix_caching.py \
  tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py \
  vllm/v1/core/kv_cache_coordinator.py \
  vllm/v1/core/single_type_kv_cache_manager.py \
  vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py

pre-commit run mypy-3.12 --hook-stage manual --files \
  tests/v1/core/test_prefix_caching.py \
  tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py \
  vllm/v1/core/kv_cache_coordinator.py \
  vllm/v1/core/single_type_kv_cache_manager.py \
  vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py

Container image:
voipmonitor/vllm@sha256:284784e685aa0377f1cf63a312a364fc884b02beb98949d6886624edbddb3806.

Test Result

Baseline on the unchanged #556 head, with the first two follow-up tests
overlaid:

• SWA replay returned 0 cached tokens instead of the reachable 64-token prefix.
• Sparse retention received 64-token scheduler alignment instead of the
  32-token fine-hit alignment.
• Result: 2 failed, 95 deselected.

The Opus review added focused coverage against the prior PR head. Four cases
failed before the review fixes:

• SWA-driven hybrid replay returned 0 instead of 224 tokens for its Mamba group.
• Pure SWA under EAGLE returned 0 instead of 224 tokens on replay.
• Mooncake omitted the Mamba and SWA states at the EAGLE-reachable predecessor.
• Mooncake retained Mamba block 5 with coarse alignment instead of block 6 with
  fine-hit alignment.

A closure review then found that Mooncake enabled fine hits for a mixed SWA and
Mamba layout that the engine correctly rejects. The new regression test failed
before the compatibility fix on assert not coord.enable_partial_hash_hits.
Without that guard, store_mask reached the SWA alignment assertion.

A GPT-6 Astra review found that fine alignment could retain an unmaterialized
Mamba state while removing its usable scheduler-aligned fallback. In its
480-token reproduction, identical replay improved from 0 to 384 tokens after
retaining both positions. Engine and Mooncake regression tests cover this case.

Current head:

• 382 tests passed across the six suites in the pytest command.
• All focused EAGLE, SWA, Mamba, and Mooncake regressions passed, including
  three SWA annotation routes and the mixed-layout compatibility fallback.
• The complete changed-file hook set passed, including Ruff, typos, mypy 3.10,
  SPDX, import, configuration, and forbidden-API checks.
• The manual CI-style mypy 3.12 hook passed.
• git diff --check passed.

No GPU model evaluation was run for the follow-up commits. This PR remains a
draft because the repository requires a serving evaluation and human review
before it is ready. logprobz must review every changed line, run the relevant
tests, and add the GLM-5.3-Flash replay results before marking it ready.

AI assistance

OpenAI Codex assisted with implementation, tests, and this description. Two
Claude Opus 5 agents independently reviewed the implementation and the PR's
guideline compliance. The original review posted by logprobz on #556 credits
Claude Fable 5.1. GPT-6 Astra independently reviewed the current PR and found
the fine-hit materialization regression described above.

The first Codex follow-up commit was already published without attribution and
sign-off trailers. It remains unamended to preserve published history without a
force-push. All three later follow-up commits include
Assisted-by: OpenAI Codex and Signed-off-by: logprobz trailers.

---

Essential Elements of an Effective PR Description Checklist

• The purpose and distinction from related PRs are described above.
• Test commands and current results are included.
• GLM-5.3-Flash model evaluation is pending while this remains a draft.
• No documentation update is required for this internal cache correction.

BEFORE SUBMITTING, PLEASE READ https://docs.vllm.ai/en/latest/contributing

Superseded

Superseded by #669, now merged on dev/jovian-judgement. The coordinated port preserves endpoint checkpoint guards and includes the cache, scheduler, connector and event repairs from this earlier branch. The combined release passed full MTP3 and DFlash2 serving qualification and was promoted, with zero preemptions and all LP26 5% performance gates satisfied. Closing this older proposal to avoid duplicate integration.

[coderabbitai]

Warning

Review limit reached

Next included review available in 12 minutes.

Check out review usage here.

View limit details

Limit details: You’ve used the included review currently available.

You've used all free OSS reviews for now. Wait for the free limit to reset to keep reviewing this public repository.

Learn how review limits work.

Review configuration:

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Team

Run ID: ca604b37-ed9d-49b0-abc2-aecc438be364

📥 Commits

Reviewing files that changed from the base of the PR and between a84f907 and 1279e4f.

📒 Files selected for processing (5)

• tests/v1/core/test_prefix_caching.py
• tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py
• vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py
• vllm/v1/core/kv_cache_coordinator.py
• vllm/v1/core/single_type_kv_cache_manager.py

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[voipmonitor]

@coderabbitai review

[coderabbitai]

⚠️ Action not completed

Review rate limited.

Note: CodeRabbit is an incremental review system and does not re-review already reviewed commits. This command is applicable only when automatic reviews are paused.

[voipmonitor]

Status: superseded by PR #669 and intentionally closed.

PR #669 consolidates the sparse hybrid replay-boundary behavior from this pull request, preserves the contributing authorship in commit trailers, and was merged into dev/jovian-judgement as 86acae9. Do not merge or cherry-pick PR #643 separately.