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

[starkwj]

Purpose

Fix MooncakeStore cache misses caused by an EAGLE-induced boundary mismatch between Mamba/KDA states and attention KV cache.

For a reusable Mamba/KDA token boundary B, EAGLE attention lookup probes the original cache boundary B + H (H is the prefix-match-unit) and then drops one hash unit H due to the eagle drop, while other cache type (i.e., Mamba) directly looks up state at B.

For a request at its prompt's tail, the existing boundary offload path (_sub_block_tail_puts) stores Mamba and Attention both at B.
Thus, when the same or the subsequent request lookup MooncakeStore, it get Attention at B, and drop H due to eagle drop, then lookup Mamba at B-Hand missed. (When the HBM prefix cache is evicted.)

This PR makes the store layout match lookup behavior:

• Store Mamba/KDA state at B.
• Store eligible EAGLE attention KV at B + H.
• Only store B + H when GPU computation has completed through that position.

Changes

• Share the per-group EAGLE peek-margin calculation between lookup and store.
• Add completed_token_len to store metadata so the worker only publishes an attention companion when B + H has been computed.
• Handle each boundary independently and deduplicate overlapping MooncakeStore keys.

The H is the EAGLE peek margin for each attention group.
For fine-grained lookup, H is the hash block size; for block-granularity groups, it may be the group's physical block size.

The boundary handoff logic covers three cases:

Boundary case	Mamba/KDA	Eagle-group attention written by the boundary handoff
Sub-block boundary	Store at B	Store through B + H when completed; otherwise only through B
Block-aligned prompt replay boundary	Store at B	Store the B + H companion when completed; otherwise publish no companion
Block-aligned internal checkpoint	Store at B	Publish no additional attention companion as the original code

_sub_block_tail_puts is renamed to _boundary_tail_puts because it now
handles both sub-block boundaries and block-aligned prompt replay boundaries
that require an EAGLE attention companion.

Test Plan

• Kimi-K3, TP 8, DCP 8, DSpark (Inferact/Kimi-K3-DSpark), MooncakeStore (standalone mode), and other settings following vllm recipe (e.g., prefix-match-unit 128).
  Test with some requests, then /reset_prefix_cache clean the HBM cache, send the same requests again.
• unit tests.

Test Result

Request, prompt length	version	send 1st time	after HBM cache reset and send 2nd time
A, 7449	main	cached_tokens: 0 created_cache_tokens: 7296	cached_tokens: 0 created_cache_tokens: 7296
A, 7449	this PR	same as above	cached_tokens: 7296 created_cache_tokens: 0
B, 23196	main	cached_tokens: 0 created_cache_tokens: 23040	cached_tokens: 0 created_cache_tokens: 23040
B, 23196	this PR	—	cached_tokens: 23040 created_cache_tokens: 0
C, 41300	main	cached_tokens: 0 created_cache_tokens: 41088	cached_tokens: 30464 created_cache_tokens: 10624
C, 41300	this PR	—	cached_tokens: 41088 created_cache_tokens: 0

Thus, for the main current code, request A and B totally missed from MooncakeStore.
Request C got partial hit at an internal point (FlashKDA internal checkpoint), but still missed a large portion.

unit tests passed:

• tests/v1/kv_connector/unit/test_mooncake_store_worker.py
• tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py

---

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.

Duplicate-work check

No open PR implementing the same MooncakeStore write-side B / B + H boundary pairing was found.

There are some PRs targeting at the similar issue around eagle drop, including #51295, #55036, #53945, etc.
But not related to external KV offloading as Mooncake.

• [Bugfix][Mooncake] Save exact Mamba boundary states #51358 (merged): save mamba state at B, but not save tail attention KV at B+H.

AI assistance

AI was used for problem analysis, code implementation and review.

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

[Copilot]

🔵 Needs a closer look

One or more issues must be addressed before approval.

Pull request overview

Fixes MooncakeStore prefix-cache misses by aligning Mamba/KDA boundary writes with EAGLE attention lookup behavior.

Changes:

• Adds shared EAGLE peek-margin calculation.
• Tracks completed token length before storing attention proofs.
• Handles per-boundary tail writes and deduplicates keys.
• Adds focused scheduler and worker tests.

File summaries

File	Description
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py	Updated as part of this pull request.
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py	Updated as part of this pull request.
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py	Updated as part of this pull request.
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py	Updated as part of this pull request.
tests/v1/kv_connector/unit/test_mooncake_store_worker.py	Updated as part of this pull request.
tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py	Updated as part of this pull request.

Review details

Suppressed comments (2)

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py:762

• This hardcodes the fine-grained case (eagle_margin == hash_block_size), but _eagle_peek_margin returns spec.block_size for non-fine-grained EAGLE groups such as SlidingWindowManager. Lookup then probes B + spec.block_size and drops that block, while this store keeps that group at B, so a hybrid Mamba + EAGLE-SWA boundary cannot be reused. Apply the per-group margin here.

            eagle_margin = self.coord.eagle_peek_margin_by_group.get(g_idx)
            if (
                eagle_margin == hash_block_size
                and boundary + eagle_margin <= req_meta.completed_token_len
            ):
                group_boundary += eagle_margin

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py:838

• This dispatch gate only recognizes an EAGLE peek of one hash unit and checks completion only through B + H. For a block-granularity EAGLE group, lookup requires B + spec.block_size; with no hash-unit-margin group, an aligned boundary takes _boundary_snapshot_puts and never publishes the companion key even when that block is complete. Derive the gate from each group's actual margin, matching the lookup-side calculation.

        hash_unit_peek = any(
            margin == self.coord.hash_block_size
            for margin in self.coord.eagle_peek_margin_by_group.values()
        )

• Files reviewed: 6/6 changed files
• Comments generated: 0
• Review effort level: Lite

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.