[Core][Bugfix] Roll back eager prefix-cache registration on uncommitted free

[ZJY0516]

Purpose

KVCacheManager.allocate_slots calls coordinator.cache_blocks, which inserts block hashes into BlockPool.cached_block_hash_to_block before the worker writes the K/V bytes. The registration covers total_computed_tokens + num_new_tokens — the range the worker is expected to compute this step, not what it has already written.

If the request is preempted/aborted before the worker actually writes those bytes (priority-preempt mid-schedule(), KV-connector load failure, etc.), the standard free path only decrements ref_cnt and pushes the block onto the free queue — the hash entry and block.block_hash survive. A later request whose request.block_hashes matches will "cache-hit" a block that was never actually written to, reading uninitialized memory.

This PR adds an explicit lifecycle for eager registrations using a FIFO bucket queue:

• BlockPool._uncommitted: list[dict[req_id, list[KVCacheBlock]]] holds one bucket per in-flight scheduler step.
• BlockPool.begin_step() pushes a new bucket; called at the top of Scheduler.schedule().
• BlockPool.commit_step() pops the oldest bucket; called at the top of Scheduler.update_from_output() once the worker has confirmed its writes for that step.
• BlockPool.rollback_uncommitted(req_id) iterates every bucket, pops the request's pending entries, and evicts them from the cache map via the existing _maybe_evict_cached_block helper.

Preempt and abort paths call rollback_uncommitted before free:

# _preempt_request and KV-connector load-failure path
self.kv_cache_manager.rollback_uncommitted(request.request_id)
self.kv_cache_manager.free(request)

The normal finish path calls only free(request) (unchanged from main). By that point the worker has confirmed the writes and commit_step has already popped this step's bucket, so the entries stay in the cache map for future hits.

Why FIFO buckets (not a single global dict)

The first version used a single dict[req_id, list[block]] cleared at step boundaries. Codex review correctly flagged that this is unsafe in two ways:

1. Sync, same-step finish: a request that hits EOS in update_from_output would be freed after the worker wrote its bytes but before the next schedule() could mark them committed — so the eager rollback would erase valid cache entries.
2. Async batch queue (EngineCore.step_with_batch_queue): multiple schedule() calls may be in flight before the oldest one's worker completes, so a single global commit at the next schedule() would erase entries whose worker writes haven't been confirmed yet.

The per-step FIFO bucket queue closes both windows: commit happens at the top of update_from_output for the specific step whose worker just confirmed; rollback iterates every still-pending bucket so a request preempted/aborted at any point removes its zombies from every in-flight step.

Test Plan

Test Result

---

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.

[gemini-code-assist]

Code Review

This pull request introduces a mechanism to prevent 'zombie' cache entries caused by eager registration before worker confirmation. It implements an uncommitted tracking system in the BlockPool, allowing the KVCacheManager to roll back cache entries if a request is preempted or aborted before the worker executes. The Scheduler now triggers a commit_step at the start of each cycle to finalize registrations from the previous step. Comprehensive tests were added to verify rollback behavior in various scenarios, including hybrid models and scheduler-level preemption. I have no feedback to provide.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: b5c7dff954

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

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

[ZJY0516]

@codex review

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 904629a634

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

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

[chatgpt-codex-connector]

Track only speculative registrations in uncommitted buckets

cache_full_blocks now adds every newly cached block to _uncommitted, but some cache writes happen after the worker has already confirmed data (e.g., AsyncScheduler._update_request_with_output() calls kv_cache_manager.cache_blocks() in vllm/v1/core/sched/async_scheduler.py). Those already-committed blocks are being recorded in whatever bucket is currently last, so a later preemption can call rollback_uncommitted() and evict valid cache entries (resetting block_hash) even though their KV bytes are real, causing avoidable cache misses and recomputation under async scheduling.

Useful? React with 👍 / 👎.

[heheda12345]

@codex can you help to check one case:
block_size=4, in the same step, the following things happen

schedule request A: [A B C D E F] -> cache [A B C D]
schedule request B: [A B C D G H] -> prefix cache hit length=4, reuse request A's [A B C D] and only compute [G H]
A is preempted so we don't compute kv cache of [A B C D] in this step, causing request B's wrong result

[chatgpt-codex-connector]

To use Codex here, create a Codex account and connect to github.

[ZJY0516]

@codex can you help to check one case:
block_size=4, in the same step, the following things happen

schedule request A: [A B C D E F] -> cache [A B C D]
schedule request B: [A B C D G H] -> prefix cache hit length=4, reuse request A's [A B C D] and only compute [G H]
A is preempted so we don't compute kv cache of [A B C D] in this step, causing request B's wrong result

[chatgpt-codex-connector]

💡 Codex Review

vllm/vllm/v1/core/block_pool.py

Lines 284 to 285 in 904629a

	blk.block_hash = block_hash_with_group_id
	self.cached_block_hash_to_block.insert(block_hash_with_group_id, blk)

Hide uncommitted blocks from prefix-cache lookups

Publishing the block hash to cached_block_hash_to_block before the step is committed lets other requests cache-hit data that has not been written yet. In the A:[A B C D E F] / B:[A B C D G H] case, B can touch A’s eager block, then A is preempted and rolled back, but B still remains scheduled with that prefix treated as computed, which can yield incorrect outputs from unwritten KV. The new _uncommitted buckets only prevent future hits; they do not protect already scheduled consumers, so uncommitted entries need to be hidden from lookup until commit (or dependent requests must be invalidated).

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

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

[mergify]

Hi @ZJY0516, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy failing?

mypy is run differently in CI. If the failure is related to this check, please use the following command to run it locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @ZJY0516.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork