[Bugfix] Fix TOCTOU race in KV block allocator causing prefix-cache block theft

[AbhiOnGithub]

Summary

Fixes #37076

When the V1 scheduler processes multiple WAITING requests in a single scheduling step, a use-after-free / TOCTOU window in the KV block allocator allows one request to silently steal another request's cached prefix block, causing KV data corruption (token bleed between requests).

Root Cause

The scheduler loop calls two methods for each waiting request:

for req in waiting_requests:
    computed_blocks, n = manager.get_computed_blocks(req)   # Step 1: lookup
    ...
    manager.allocate_slots(req, ..., computed_blocks)        # Step 2: pin + alloc

Between Step 1 and Step 2, other requests in the loop are also processed. The bug:

1. get_computed_blocks(req_A) finds cached block_X (ref_cnt=0, eviction candidate in the free queue)
2. Before allocate_slots(req_A) pins it via touch(), req_B's allocate_new_blocks() steals block_X from the free queue and erases its hash via _maybe_evict_cached_block()
3. allocate_slots(req_A) now touch()es a block that belongs to req_B → req_A reads req_B's KV data

Timeline (single-threaded scheduler loop):
  get_computed_blocks(A)   →  block_X found (ref_cnt=0)
                                              ← TOCTOU WINDOW OPENS
  get_computed_blocks(B)   →  no cache hit
  allocate_new_blocks(B)   →  steals block_X from free queue!
                                              ← block_X now has B's data
  allocate_slots(A)        →  touches stale block_X → TOKEN BLEED

Fix

Pre-touch (pin) cached blocks immediately inside get_computed_blocks(), closing the TOCTOU window before any other request can run.

• Added touch_computed_blocks() and release_computed_blocks() to KVCacheCoordinator
• get_computed_blocks() calls touch_computed_blocks() right after find_longest_cache_hit() so ref_cnt goes 0→1 and blocks are removed from the free queue
• allocate_slots() on the failure path (not enough free blocks) calls release_computed_blocks() to undo the pin — no ref-count leak
• allocate_new_computed_blocks() no longer calls touch() (blocks already pinned); handles the sliding-window skipped-block case by freeing pre-touched skipped blocks before slicing

Why the free-block budget check is unchanged

Before the fix: num_blocks_to_allocate = num_new_blocks + N (N evictable computed blocks) is compared against F free blocks (which include those N blocks). Equivalent to num_new_blocks > F - N.

After the fix: pre-touch removes N blocks from the free queue, so get_num_free_blocks() = F - N, and _get_num_evictable_blocks() = 0. Check becomes num_new_blocks > F - N. Identical condition — scheduling decisions are unchanged.

Test Plan

• Added regression test test_prefix_cache_block_not_stolen_between_get_and_alloc that directly reproduces the TOCTOU scenario and verifies the prefix block's ref_cnt == 1 after get_computed_blocks()
• All 56 unit tests pass (test_prefix_caching.py + test_single_type_kv_cache_manager.py)
• All pre-commit hooks pass (ruff, mypy, typos, SPDX, etc.)

[gemini-code-assist]

Code Review

This pull request introduces a fix for a critical TOCTOU race condition in the KV block allocator by pre-touching (pinning) cached blocks immediately after they are found. The changes are logical and well-structured, and the new regression test effectively reproduces and verifies the fix. However, I've identified a potential critical resource leak scenario where pre-touched blocks may not be released if a request is aborted or preempted before allocation, which needs to be addressed.

I am having trouble creating individual review comments. Click here to see my feedback.

vllm/v1/core/kv_cache_manager.py (348-354)

There appears to be a potential resource leak with the new pre-touch mechanism. Blocks are pinned in get_computed_blocks(), but they are only released here inside allocate_slots() on the specific failure path of having insufficient free blocks.

If a request is aborted or preempted by the scheduler after get_computed_blocks() has been called but before allocate_slots() is attempted, the pre-touched blocks will not be released. This will lead to a leak of KV cache blocks over time.

To fix this, the scheduler logic must be updated to ensure release_computed_blocks() is called for any request that has had blocks pre-touched but does not proceed to allocation for any reason (e.g., preemption, client disconnect). The entity that calls get_computed_blocks should be responsible for calling release_computed_blocks on all non-successful paths.

[mergify]

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

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

[mergify]

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

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

[mergify]

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

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