[Core] Make hybrid KV load recovery attention-aware - #51731
Closed
Dao007forever wants to merge 2 commits into
Closed
Dao007forever wants to merge 2 commits into
Dao007forever wants to merge 2 commits into
Conversation
Assisted-by: OpenAI Codex Signed-off-by: Dao Le <Dao007forever@gmail.com>
Signed-off-by: Dao Le <Dao007forever@gmail.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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 raisesValueError: too many values to unpackinstead 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:
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_blocksnow: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
AsyncSchedulerand a synchronous loader such asLMCacheConnectorV1,request.num_computed_tokensmay include several in-flight speculative steps, while the current recovery path subtracts onlynum_scheduled_tokensfor 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:
Run all pre-commit hooks applicable to the changed files:
Test Result
Focused recovery tests:
The new coverage verifies:
All applicable pre-commit hooks passed.
Essential Elements of an Effective PR Description Checklist