Skip to content

[Core] Make hybrid KV load recovery attention-aware - #51731

Closed
Dao007forever wants to merge 2 commits into
vllm-project:mainfrom
Dao007forever:dao/recompute
Closed

Dao007forever wants to merge 2 commits into
vllm-project:mainfrom
Dao007forever:dao/recompute

Conversation

@Dao007forever

@Dao007forever Dao007forever commented Aug 10, 2026

Copy link
Copy Markdown
Contributor

Purpose

Related to #50687 and the precise-recovery RFC in #47812.

KV-load failure recovery currently assumes a single KV-cache group. With the hybrid memory allocator, get_block_ids() returns one block table per group, so the failure path raises ValueError: too many values to unpack instead of recovering the request.

Simply mapping a failed block index to a token boundary is not sufficient for hybrid models. Whether a boundary is resumable depends on the state required by every attention group at that exact token count:

  • Rewinding a sliding-window group shifts its window left and may expose null blocks that were already released.
  • A Mamba group may resume from a previous retained state even though older entries are null.
  • Groups can have different physical block sizes but must agree on a common cache-hit boundary.
  • EAGLE's cache-hit rewind has already been applied and must not be applied a second time during load-failure recovery.

This change adds an attention-aware search for the longest safe recovery prefix. It walks backward over the coordinator's common cache-hit alignment and, for every group, uses get_num_skipped_tokens(candidate) to identify the blocks required at that candidate. A candidate is accepted only when all of those blocks exist and none were reported invalid.

_update_requests_with_invalid_blocks now:

  • Detects reported failures across every cache group using each group's block size.
  • Assigns shared failed blocks to one request for recomputation as before.
  • Rewinds an owning request to the longest group-consistent prefix.
  • Collects invalid and downstream eviction candidates across all groups while excluding null block IDs.

The coordinator exposes its cache-hit alignment so recovery uses the same boundary contract as prefix-cache lookup, including fine-grained hybrid hits.

Overlap with existing PRs

The mandatory duplicate-work check found #45497, #48216, #50388, and #50742. The distinction was documented in #50687 before opening this draft.

Known async-scheduling limitation

This change does not claim to solve synchronous connector-load failures under scheduler runahead. With AsyncScheduler and a synchronous loader such as LMCacheConnectorV1, request.num_computed_tokens may include several in-flight speculative steps, while the current recovery path subtracts only num_scheduled_tokens for the output reporting the failure. That can overestimate the settled prefix. Later in-flight outputs derived from the failed load must also be marked stale and drained.

Solving that requires coordinated in-flight output invalidation in addition to changing the prefix calculation, so it is called out here rather than folded silently into the hybrid recovery change.

No documentation or model evaluation is required: this changes scheduler bookkeeping only on the KV-load failure path and does not change model kernels, weights, or nominal model output behavior.

AI assistance disclosure

AI assistance was used to analyze the recovery invariants and prepare this draft and its description. The PR remains a draft until the human submitter has reviewed every changed line and can defend the change and its tests end to end.

Test Plan

Run the focused invalid-block recovery suite:

.venv/bin/python -m pytest \
  tests/v1/kv_connector/unit/test_invalid_blocks_correctness.py -v

Run all pre-commit hooks applicable to the changed files:

.venv/bin/pre-commit run --files \
  vllm/v1/core/kv_cache_coordinator.py \
  vllm/v1/core/sched/scheduler.py \
  tests/v1/kv_connector/unit/test_invalid_blocks_correctness.py

Test Result

Focused recovery tests:

7 passed, 17 warnings in 2.24s

The new coverage verifies:

  • A failure in a smaller-block hybrid group uses the actual failed block and common alignment.
  • Rewinding a pruned sliding window continues to zero when earlier required blocks are null.
  • Mamba recovery finds the previous retained state.
  • EAGLE recovery does not apply its prefix-cache drop twice.

All applicable pre-commit hooks passed.


Essential Elements of an Effective PR Description Checklist
  • The purpose and related issues are described.
  • The test plan contains exact commands.
  • Test results are included.
  • No documentation update is required.

Assisted-by: OpenAI Codex
Signed-off-by: Dao Le <Dao007forever@gmail.com>
@mergify mergify Bot added the kv-connector label Aug 10, 2026
Signed-off-by: Dao Le <Dao007forever@gmail.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant