[Bugfix][KV Offload] Stop offloading the final sampled token's KV slot by almogtavor · Pull Request #54288 · vllm-project/vllm

almogtavor · 2026-08-28T23:21:42Z

Purpose

_build_store_jobs lets a finished request offload up to req.num_tokens. The last of those tokens has no KV of its own because it was produced by the forward pass at the position before it and the request ends before any pass covers its own slot. The last position with KV that anything wrote is therefore num_tokens - 1 and a finish landing on a block boundary hands the offload tier a full block whose final slot sits past that point.

Offloaded blocks are keyed by a hash of their token ids and every id involved is legitimate. A later request that replays the same tokens in normal multi-turn chat will therefore match the hash and load the bad block. In the reproduction below the GPU prefix cache commits one block where the offload path stored two.

Under speculative decoding the junk is dangerous since Spec decode proposes several draft tokens, runs them all through the model, then throws away the rejected ones. Slot 3 that was never written for that token (although we generated the [0,1,2,3] tokens), can hold the KV of the first rejected draft. It’s a KV vector for a token that vllm decided not to keep, but downstream it looks legitimate.
The cache key is a hash of [a, b, c, EOS] so later request that happens to start with those same four tokens gets a hit and loads the poisoned block, so the model attends to a token that wasnt accepted.

Fix

num_tokens_after_batch = max(req.num_prompt_tokens, req.num_tokens - 1) in the finished branch.

I compared that against num_computed_tokens - num_output_placeholders, the confirmed-token formula used by both AsyncScheduler._update_request_with_output and simple_kv_offload. One H100 running Qwen3-30B-A3B-W4A16 with 11 finished requests observed per configuration:

configuration	confirmed-token formula matched `num_tokens - 1`	`max()` changed the answer
sync, no spec decode	11/11	never
async, no spec decode	11/11	never
sync, ngram spec decode	5/11	never
async, ngram_gpu spec decode	11/11	never

The confirmed-token formula runs ahead by up to num_speculative_tokens under synchronous speculative decoding because num_computed_tokens counts the whole scheduled batch before rejected drafts are rolled back. Logged cases had num_computed_tokens at 28 and 60 and 98 where num_tokens - 1 was 27 and 57 and 95. The max() suggested in the issue applies only to a request that finishes without generating anything and nothing reaches this branch in that state because aborted requests take the branch above it and pooling requests finish in the same step they prefill so they take the running branch instead. A pooling probe of 12 requests saw none arrive here. The max() therefore never changed an answer but it costs nothing and covers the case if a future caller does reach it.

#48596 added this branch so that the block completed at finish would stop being skipped. Skipping it is the correct behaviour and the branch still runs for every other block so the worker side fence fix from that PR is untouched.

Reproduction

CPU only. A 7-token prompt finishing on a block boundary has the GPU prefix cache committing 1 block while the offload path stores 2. A second request whose prompt is the first request's full token stream including EOS then gets content-hash hits on both.

Measured on one H100

Qwen3-30B-A3B-W4A16 with offload_prompt_only: false because the default of true clamps the watermark down to the prompt and hides this branch completely.

A request finishing at exactly 64 tokens was snapshotted at the step before it finished and again at the step it finished. Slot 63 of its last GPU block was bit identical across that final forward pass and held all zeros while slot 62 went from zero to a norm of 526. The store job emitted at that same step copied that GPU block to the offload tier.

A later request whose prompt starts with those 64 tokens then loaded the block back into fresh GPU blocks. Position 63 came back all zeros while its fifteen neighbours ranged from 506 to 538 so the model decoded while attending to an empty slot.

Changing the one line to num_tokens - 1 stops that store job. The same later request recomputes the block instead and position 63 comes back at 520.

Test plan

pytest tests/v1/kv_connector/unit/offloading_connector -q

241 passed, 2 skipped. Three tests encoded the old watermark and now assert the corrected contract, with test_last_block_offloaded_at_request_finish renamed to test_final_sampled_token_does_not_complete_an_offloaded_block. Running those tests against unfixed source gives 4 failed, 237 passed, 2 skipped.

ruff check and ruff format --check at v0.14.0 pass on both changed files. No model eval because the measurement above reads the KV cache directly and is exact where sampled output on this mixture of experts model drifts between runs of identical code.

Not a duplicate: no other open PR touches this branch.

When a request finishes, _build_store_jobs sets the store watermark to req.num_tokens. That counts the final sampled token, but no forward pass ever writes its KV slot: the token was produced by the previous position's forward, and the request ends before another one runs. If the finish lands on a block boundary the last block is stored with an unwritten slot in it. Under spec decode that slot holds the KV of the first rejected draft token, so the stored block is plausible garbage rather than obviously wrong. The block is content-addressed over token ids, all of which are valid, so a later request whose prompt replays the same tokens (the multi-turn shape) hits it and loads the unwritten slot. The engine's own prefix cache already refuses to commit that far. KVCacheManager caps caching at min(computed + new, num_tokens), and the async scheduler commits num_computed_tokens - num_output_placeholders. Both stop one token short of where the offload path was storing. Measured at the finishing step with block_size 4 and a 3 token prompt plus EOS: sync gives num_tokens=4, num_computed=3, placeholders=0; async gives num_tokens=4, num_computed=4, placeholders=1. The frontier is 3 in both, while the code used 4. num_tokens - 1 is exact in both modes, which plain num_computed_tokens is not. Reproduced without a GPU: with a 7 token prompt finishing on a boundary, the GPU prefix cache commits 1 block while the offload path stored 2. A second request whose prompt is the first request's full token stream, EOS included, then gets content-hash hits and loads both blocks. Three existing tests asserted the old watermark; they are updated to the corrected contract. test_on_request_finished_fires_after_final_block_store now declines the first prepare_store so its ordering assertion still has a store to order against. Fixes vllm-project#54193 Signed-off-by: almogtavor <almogtavor@gmail.com>

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.

tomylin890 · 2026-08-29T02:26:01Z

Thanks for jumping on this so fast. The fix is essentially identical to what we've been running in our fork since we hit it — same branch, same -1 — and it's been serving production traffic behind decode-side offloading for about a day now with soak checks green, so you can count this as a field confirmation.

One thing that might be worth a thought while it's in review: we wrote the watermark as max(req.num_prompt_tokens, req.num_tokens - 1) instead of the bare -1. The two only diverge for a request that finishes without ever sampling a token, where the bare form also drops the last prompt position — a slot prefill did write. We never managed to construct a reachable path to that state (aborts take the num_computed_tokens branch), so this is probably paranoia, but the clamp costs nothing and would save someone a confused hour if such a path ever shows up.

Also, declining the first prepare_store so the retry lands on the finishing step is a much cleaner way to pin the ordering than what we had in our tests. Stealing that.

almogtavor · 2026-08-29T09:46:33Z

@tomylin890 good idea I'll add that

A request with zero output tokens never reaches this branch today, but the clamp costs nothing and keeps the last prompt position, which prefill did write, offloadable if one ever does. Signed-off-by: almogtavor <almogtavor@gmail.com>

almogtavor · 2026-08-29T11:46:17Z

Added in 79290b6. I also ran it on an H100 against num_computed_tokens - num_output_placeholders with 11 finished requests observed per configuration:

configuration	confirmed-token formula matched `num_tokens - 1`	`max()` changed the answer
sync, no spec decode	11/11	never
async, no spec decode	11/11	never
sync, ngram spec decode	5/11	never
async, ngram_gpu spec decode	11/11	never

max() only applies to a request that finishes without generating anything and nothing reached this branch in that state so it never changed an answer. I kept it since it costs nothing (as guard). The confirmed-token formula is what I avoided because it runs ahead by up to the draft length under synchronous spec decode.

…ler.py Co-authored-by: Or Ozeri <or@ozery.com> Signed-off-by: Almog Tavor <70065337+almogtavor@users.noreply.github.com>

coderabbitai · 2026-09-06T09:51:38Z

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Repository UI

Review profile: CHILL

Plan: Team

Run ID: d46a8e0a-78f9-4775-af4e-d57c1ebbf8c7

📥 Commits

Reviewing files that changed from the base of the PR and between a205224 and f96dfe2.

📒 Files selected for processing (1)

tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py

Included review availability: Your plan provides up to 10 included reviews per hour; 7 remain after this review.

📝 Summary

Summary by CodeRabbit

Bug Fixes
- Prevented the final sampled or rejected draft token from being incorrectly stored in the GPU prefix cache.
- Improved completed-request handling so only KV slots that were actually written are cached.
- Preserved existing caching behavior for aborted and active requests, ensuring expected cache updates throughout request processing.

Walkthrough

The scheduler now excludes the final sampled token from finished-request KV storage while preserving the full prompt boundary. Tests cover block completion, store retries, request-finish ordering, and sliding-window behavior.

Changes

Finished-request KV storage

Layer / File(s)	Summary
Finished-request store boundary `vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py`	Finished requests now store through `max(num_prompt_tokens, num_tokens - 1)`. Aborted and active requests retain their existing boundaries.
Scheduler regression coverage `tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py`	Tests verify that the final sampled token does not complete an offloaded block, that store retries occur before request completion callbacks, and that the final EOS block is not stored.

Estimated code review effort: 2 (Simple) | ~10 minutes

Merge Risk: ⚪ Minimal · up to f96df

Finished-request KV offloading no longer stores the unwritten final sampled-token slot, preventing invalid cached blocks from being reused. Boundary and retry behavior are covered by updated scheduler tests, with no remaining merge-blocking risk.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 40.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 5 functions across 2 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly and concisely identifies the primary fix: preventing offloading of the final sampled token's KV slot.
Description check	✅ Passed	The description directly explains the bug, the scheduler fix, the poisoning risk, and the test validation.
Linked Issues check	✅ Passed	The scheduler change implements the linked issue's requested watermark of max(req.num_prompt_tokens, req.num_tokens - 1), and the updated tests cover block-boundary finishes and invalid final slots. […
Out of Scope Changes check	✅ Passed	The changes are limited to the offloading scheduler behavior and related unit tests. They directly support the linked issue objectives, with no unrelated code changes identified.

Fix all pre-merge checks with AI

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

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

coderabbitai · 2026-09-06T09:58:02Z

Note

GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer.

…ount Signed-off-by: almogtavor <almogtavor@gmail.com>

orozery · 2026-09-06T11:06:06Z

/ci run

github-actions · 2026-09-06T11:06:13Z

✅ @almogtavor, CI is now available for this PR.

/ci run starts upstream CI; /amd-ci run starts AMD CI only.
/ci retry retries failed jobs in the CI build for the current PR head. If the current head has no CI build, it starts a new CI build for the current head containing only jobs that failed in the latest earlier CI build for this PR.
/amd-ci retry retries failed jobs in AMD CI for the current PR head. Use /amd-ci run when the current head has no AMD CI build.
/ci cancel cancels scheduled or running CI builds for this PR branch; /amd-ci cancel does the same for AMD CI only.

github-actions · 2026-09-06T11:06:23Z

✅ Triggered Buildkite CI #87438 for commit f96dfe20cbf0.

vllm-project#54288) Signed-off-by: almogtavor <almogtavor@gmail.com> Signed-off-by: Almog Tavor <70065337+almogtavor@users.noreply.github.com> Co-authored-by: Or Ozeri <or@ozery.com> Signed-off-by: Jyotirmoy Roy <jyotirmoyroy649@gmail.com>

OffloadingConnector under MTP/EAGLE: serve stored hits instead of vetoing the whole request, and read the load boundary from the computed offset. Backports three merged upstream fixes onto the 0.28.0 tree (all land after the 0.28.0 wheel, so this image has neither): vllm-project#52771 (merged 2026-09-07) -- OffloadingConnector: stop zeroing offload hits under MTP/EAGLE spec decode. Three changes: the all-groups fallback no longer marks every KV group as a drafter group (it logs instead); the volatile-tail store exclusion is lifted once a request is finished (the final chunk is then stored); the lookup query is widened for every eagle group, not only sliding-window ones. vllm-project#52807 (merged 2026-09-03) -- do not let a recurrent group's unhashed block truncate the load boundary: scan from the computed boundary, not 0, so a sparse (Mamba / SWA) group's legitimate unhashed blocks below it are not mistaken for the load start. The confirmed symptom (upstream vllm-project#52735, reproduced there on 3090 Ti, H100 NVL on stock v0.27.1, and GB200): with a hybrid-GDN model under MTP the CPU offload tier stores but never serves. The full-attention group hits every chunk while the Mamba groups return 0 and veto the whole request, so 41 GB is written and 0 bytes ever read back. This fork's own offload-dflash-eagle-groups.patch narrows the fallback for the dflash path only; under MTP every group is still flagged, which is exactly the reporter's boot line in syv-ai#95. Applies on top of offload-dflash-eagle-groups.patch (the hunks were generated against the tree with that patch already applied). Third fix (upstream vllm-project#54288, merged 2026-09-06, also absent from 0.29.0): a finished request offloaded up to num_tokens, but the final sampled token has no KV of its own (no forward pass covers its slot; under spec decode it holds a rejected draft), so a block ending there was stored under a hash of legitimate token ids and a later request replaying those tokens loaded the unwritten slot. Clamp the finished watermark to num_tokens - 1 (never below the prompt). Reachable only with offload_prompt_only=false; the default clamps to the prompt first. Upstream wrote vllm-project#52771 against a tree that already had this fix, and its EAGLE trailing-block tests fail without it. Ported to 0.29.0 from the fork's PR#100 branch (syv-ai/qwen38-27b-rtx3090, patches/offload-mtp-serve.patch at 3284992); all seven hunks apply unchanged. Neither upstream fix is in the v0.29.0 tag. Signed-off-by: Michael Stufflebeam <cpuchip@gmail.com> Co-authored-by: Claude <noreply@anthropic.com>

almogtavor requested review from ApostaC, NickLucche, ivanium, orozery and xuechendi as code owners August 28, 2026 23:21

claude Bot reviewed Aug 28, 2026

View reviewed changes

mergify Bot added bug Something isn't working kv-connector labels Aug 28, 2026

almogtavor mentioned this pull request Aug 28, 2026

[Bug]: KV offloading connector: finished-request store watermark includes a KV slot never written by any forward pass (silent cross-request poisoning with offload_prompt_only=False) #54193

Closed

1 task

Clamp the finished watermark at the prompt length

79290b6

A request with zero output tokens never reaches this branch today, but the clamp costs nothing and keeps the last prompt position, which prefill did write, offloadable if one ever does. Signed-off-by: almogtavor <almogtavor@gmail.com>

orozery reviewed Sep 1, 2026

View reviewed changes

Comment thread vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py Outdated

Comment thread tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py Outdated

This was referenced Sep 2, 2026

[Bugfix][V1] Register accepted spec-decode tokens before terminal cleanup #54979

Open

[Bugfix] OffloadingConnector: stop zeroing offload hits under MTP/EAGLE spec decode #52771

Merged

Update vllm/distributed/kv_transfer/kv_connector/v1/offloading/schedu…

d97a7dc

…ler.py Co-authored-by: Or Ozeri <or@ozery.com> Signed-off-by: Almog Tavor <70065337+almogtavor@users.noreply.github.com>

Merge branch 'main' into offload-finished-watermark

a205224

gate the finished-store test on request state instead of an attempt c…

f96dfe2

…ount Signed-off-by: almogtavor <almogtavor@gmail.com>

orozery added the ready ONLY add when PR is ready to merge/full CI is needed label Sep 6, 2026

orozery approved these changes Sep 6, 2026

View reviewed changes

orozery merged commit dc02934 into vllm-project:main Sep 6, 2026
37 checks passed

cpuchip mentioned this pull request Sep 12, 2026

Offload tier: backport the MTP/EAGLE serve fixes (#52771, #52807, #54288), boot on native from every launcher, size-checked oracle syv-ai/qwen38-27b-rtx3090#100

Open

Navigation Menu

Uh oh!

Uh oh!

[Bugfix][KV Offload] Stop offloading the final sampled token's KV slot - #54288

[Bugfix][KV Offload] Stop offloading the final sampled token's KV slot#54288
orozery merged 5 commits into
vllm-project:mainfrom
almogtavor:offload-finished-watermark

almogtavor commented Aug 28, 2026 •

edited

Loading

Uh oh!

claude Bot left a comment

Uh oh!

tomylin890 commented Aug 29, 2026

Uh oh!

almogtavor commented Aug 29, 2026

Uh oh!

almogtavor commented Aug 29, 2026 •

edited

Loading

Uh oh!

Uh oh!

Uh oh!

coderabbitai Bot commented Sep 6, 2026 •

edited

Loading

Summary by CodeRabbit

Walkthrough

Changes

❌ Failed checks (1 warning)

Uh oh!

coderabbitai Bot commented Sep 6, 2026

Uh oh!

orozery commented Sep 6, 2026

Uh oh!

github-actions Bot commented Sep 6, 2026

Uh oh!

github-actions Bot commented Sep 6, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Uh oh!

Conversation

almogtavor commented Aug 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Purpose

Fix

Reproduction

Measured on one H100

Test plan

Uh oh!

claude Bot left a comment

Choose a reason for hiding this comment

Claude Code Review

Uh oh!

tomylin890 commented Aug 29, 2026

Uh oh!

almogtavor commented Aug 29, 2026

Uh oh!

almogtavor commented Aug 29, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Uh oh!

coderabbitai Bot commented Sep 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary by CodeRabbit

Walkthrough

Changes

❌ Failed checks (1 warning)

Uh oh!

coderabbitai Bot commented Sep 6, 2026

Uh oh!

orozery commented Sep 6, 2026

Uh oh!

github-actions Bot commented Sep 6, 2026

Uh oh!

github-actions Bot commented Sep 6, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

almogtavor commented Aug 28, 2026 •

edited

Loading

almogtavor commented Aug 29, 2026 •

edited

Loading

coderabbitai Bot commented Sep 6, 2026 •

edited

Loading