Repository navigation
[Bugfix][KV Offload] Do not offload the partially computed block of a request that finished with an error - #63937
Conversation
… request that finished with an error _build_store_jobs clamps the storable token count to num_computed_tokens only for FINISHED_ABORTED. Any other finish uses max(num_prompt_tokens, num_tokens - 1), which is the whole prompt for a request that produced no output, so a request that finished with an error mid-prefill stored its half-written block under the full block hash. Clamp finished requests to num_computed_tokens as well; for a request that finished normally the clamp is a no-op. Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: Sohom Chakraborty <16609933+sohom-cs@users.noreply.github.com>
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
|
@codex review |
|
Codex Review: Something went wrong. Try again later by commenting “@codex review”. ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
|
@codex review |
|
Codex Review: Something went wrong. Try again later by commenting “@codex review”. ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
|
/ci run |
|
✅ @sohom-cs, CI is now available for this PR.
|
|
✅ Triggered Buildkite CI #98257 for commit ℹ️ This PR is 198 commits behind upstream |
Overview
When a request finishes with an error before its prefill is complete, the block it was computing is stored under the hash of the full block. A later request with that prefix hits the offloaded block and loads KV for tokens that were never computed. One-expression fix in
_build_store_jobs; sibling of #59096 and #59329.Claims
FINISHED_ERRORmid-prefill stores only the blocks it fully computed, asFINISHED_ABORTEDalready does.Validation
test_finish_mid_prefill_stores_only_computed_blocks[FINISHED_ERROR|FINISHED_ABORTED]runs a 16-token prompt through chunked prefill at 6 tokens a step (block size 4), finishes the request after the first chunk, and checks that the finish step stores nothing beyond the one full block step 1 already stored.On main (6fe3d44) the
FINISHED_ERRORcase fails:With the fix both cases pass; the file's other tests still pass (CPU only, macOS, no GPU run). Model outputs are unaffected, so no evals.
Details
_build_store_jobs(offloading/scheduler.py:1874-1883) clamps the storable token count tonum_computed_tokensonly forFINISHED_ABORTED; every other finish usesmax(num_prompt_tokens, num_tokens - 1), which is the whole prompt for a request that produced no output.storable_chunksthen only bounds by allocated blocks, so the half-written block goes out under its full hash. A request can finish this way mid-prefill throughfinish_requests(..., FINISHED_ERROR)(sched/scheduler.py:2525-2540): an encoder-cache mismatch, or a failed KV load under thefailpolicy.Fix:
min(req.num_computed_tokens, max(req.num_prompt_tokens, req.num_tokens - 1)). The spec-decode clamp is kept.Not a duplicate: #63010 (store cursor reset on a restart from token 0) and #60810 (the block table of a resumed streaming session), both merged on 10-25, and the open #62795 (flush promotions before a request's tiers are finalized) are adjacent; none touches this clamp.
AI assistance: I used an AI coding assistant (Claude) to trace this path and write the test; I reviewed every changed line and ran the tests myself.
Pull Request Checklist
I used vLLM's
/pr-checklistskill. (Mandatory for agents, optional for humans).AI assistance was used during the creation of this PR.
Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.
Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.
Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.
Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.