iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Bugfix][KV Connector] Fix SimpleCPUOffloadScheduler TOCTOU between Phase A and Phase B (#42289)

Browse Source 
Signed-off-by: Qiuyang Yue <yueqiuyang1389@gmail.com>
Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com>
Co-authored-by: gemini-code-assist <noreply@google.com>
This commit is contained in:
Qiuyang Yue
2026-05-18 21:22:33 -07:00
committed by GitHub
co-authored by Claude Opus 4.7 gemini-code-assist
parent 3ca8db2ef8
commit fab07e4d0f
2 changed files with 189 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/v1/simple_kv_offload/test_scheduler.py
+105
View File
@@ -1241,3 +1241,108 @@ def test_partial_gpu_prefix_plus_cpu_load() -> None:
assert bid in ext_block_ids, (
f"Load GPU block {bid} should be an ext_comp block, not a comp or new block"
)
# ---------------------------------------------------------------------------
# Test 11: TOCTOU between Phase A and Phase B (regression for #39702)
# ---------------------------------------------------------------------------
def test_toctou_cpu_hit_evicted_between_phases_no_crash() -> None:
"""Regression for vllm-project/vllm#39702.
When ``get_num_new_matched_tokens`` (Phase A) reports a CPU cache hit
of ``N`` tokens but ``update_state_after_alloc`` (Phase B) runs after
other requests have caused LRU eviction of those exact blocks, the
second ``find_longest_cache_hit`` call returns 0 while
``num_external_tokens`` is still ``N``, triggering
``AssertionError: Expected N hit tokens, got 0``.
Setup: ``num_cpu_blocks=5`` (4 usable; null_block takes 1).
1. Store req_a's 2 blocks to CPU. CPU: [a0, a1, _, _].
2. Phase A on req_b (same prompt as req_a) reports a 2-block hit.
Without the fix this does NOT pin a0/a1 — they remain at LRU front.
3. Store req_c (2) + req_d (2) — 4 more blocks into 4 slots. With
a0/a1 unpinned and at LRU front, they get evicted. CPU: [c0, c1,
d0, d1].
4. Phase B on req_b: re-searches the CPU coordinator, finds 0
cached hashes, asserts.
After the fix, Phase A pins the hit blocks so step 3 evicts req_c's
blocks instead, and Phase B reads the cached ``(cpu_hit_blocks,
hit_length)`` tuple without re-searching.
"""
# 5 total = 4 usable (null_block takes 1)
fix = make_scheduler(num_cpu_blocks=5, num_gpu_blocks=16, lazy=False)
sched = fix.scheduler
# --- Step 1: Store req_a's 2 blocks to CPU cache ---
req_a = make_request(num_blocks=2)
kv_a = _alloc_and_register(fix, req_a, 2)
sched.update_state_after_alloc(req_a, kv_a, num_external_tokens=0)
sched_out_a = make_scheduler_output(
{req_a.request_id: 2 * BLOCK_SIZE},
new_reqs={req_a.request_id: kv_a.get_block_ids()},
)
meta_a = sched.build_connector_meta(sched_out_a)
assert meta_a.store_event >= 0
simulate_store_completion(sched, meta_a.store_event)
# --- Step 2: Phase A — req_b (same prompt as req_a) reports a CPU hit ---
req_b = Request(
request_id="req-b-toctou",
prompt_token_ids=req_a.prompt_token_ids,
sampling_params=req_a.sampling_params,
pooling_params=None,
mm_features=None,
block_hasher=req_a._block_hasher,
)
hit_tokens, is_async = sched.get_num_new_matched_tokens(
req_b, num_computed_tokens=0
)
assert hit_tokens == 2 * BLOCK_SIZE, (
f"Phase A should report 2 blocks of CPU hit, got {hit_tokens}"
)
assert is_async is True
# --- Step 3: TOCTOU window — fill CPU cache so LRU evicts req_a's blocks
# (in production this corresponds to other concurrent requests landing
# between Phase A and Phase B for req_b). 4 usable slots, req_a occupies 2;
# req_c (2) + req_d (2) require evicting 2 LRU blocks. ---
req_c = make_request(num_blocks=2)
req_d = make_request(num_blocks=2)
kv_c = _alloc_and_register(fix, req_c, 2)
kv_d = _alloc_and_register(fix, req_d, 2)
sched.update_state_after_alloc(req_c, kv_c, num_external_tokens=0)
sched.update_state_after_alloc(req_d, kv_d, num_external_tokens=0)
sched_out_pressure = make_scheduler_output(
{
req_c.request_id: 2 * BLOCK_SIZE,
req_d.request_id: 2 * BLOCK_SIZE,
},
new_reqs={
req_c.request_id: kv_c.get_block_ids(),
req_d.request_id: kv_d.get_block_ids(),
},
)
meta_pressure = sched.build_connector_meta(sched_out_pressure)
assert meta_pressure.store_event >= 0
simulate_store_completion(sched, meta_pressure.store_event)
# --- Step 4: Phase B — must not crash ---
# Before fix: AssertionError: Expected 32 hit tokens, got 0
# After fix: Phase A pinned the hits → Step 3 evicted req_c instead,
# Phase B consumes the cached (cpu_hit_blocks, hit_length) tuple.
gpu_blocks_b = fix.gpu_block_pool.get_new_blocks(2)
kv_blocks_b = KVCacheBlocks(blocks=(gpu_blocks_b,))
sched.update_state_after_alloc(req_b, kv_blocks_b, num_external_tokens=hit_tokens)
# --- Step 5: with the fix, the load is queued correctly ---
sched_out_b = make_scheduler_output(
{req_b.request_id: 1},
new_reqs={req_b.request_id: kv_blocks_b.get_block_ids()},
)
meta_b = sched.build_connector_meta(sched_out_b)
assert meta_b.load_event >= 0, (
"Phase B should queue a load event using the pinned CPU hit blocks"
)
assert len(meta_b.load_gpu_blocks) == 2
assert len(meta_b.load_cpu_blocks) == 2
vllm/v1/simple_kv_offload/manager.py
+84 -12
View File
@@ -131,6 +131,12 @@ class SimpleCPUOffloadScheduler:
# the worker reports completions by event index, not request id.
self._load_event_to_reqs: dict[int, list[str]] = {}
# Pending (cpu_hit_blocks, hit_length) tuples from find_longest_cache_hit,
# kept pinned via touch() while awaiting update_state_after_alloc().
self._pending_cpu_hits: dict[
str, tuple[tuple[list[KVCacheBlock], ...], int]
] = {}
# Store metadata
self._lazy_mode = lazy_offload
# Lazy mode: use a cursor to track the last scanned block in the GPU free queue.
@@ -212,7 +218,17 @@ class SimpleCPUOffloadScheduler:
def get_num_new_matched_tokens(
self, request: "Request", num_computed_tokens: int
) -> tuple[int | None, bool]:
"""Return (num_new_tokens, is_async) from consecutive CPU cache hits."""
"""Return (num_new_tokens, is_async) from consecutive CPU cache hits.
Pins found CPU blocks via touch() so they survive LRU eviction until
update_state_after_alloc() consumes them. Any pin from an earlier
call on the same request (e.g. retry after a failed allocate_slots)
is dropped first.
"""
stale = self._pending_cpu_hits.pop(request.request_id, None)
if stale is not None:
self._free_pending_cpu_hit(stale)
skipped = num_computed_tokens // self.block_size
remaining_hashes = request.block_hashes[skipped:]
@@ -223,11 +239,19 @@ class SimpleCPUOffloadScheduler:
max_hit_len = request.num_tokens - 1 - num_computed_tokens
if max_hit_len <= 0:
return 0, False
_, hit_length = self.cpu_coordinator.find_longest_cache_hit(
cpu_hit_blocks, hit_length = self.cpu_coordinator.find_longest_cache_hit(
remaining_hashes, max_hit_len
)
if hit_length > 0:
pin_blocks = [
blk for grp in cpu_hit_blocks for blk in grp if not blk.is_null
]
self.cpu_block_pool.touch(pin_blocks)
self._pending_cpu_hits[request.request_id] = (
cpu_hit_blocks,
hit_length,
)
return hit_length, True
return 0, False
@@ -253,29 +277,60 @@ class SimpleCPUOffloadScheduler:
num_stored_blocks=[0] * num_groups,
)
# Pop the CPU hit cached by get_num_new_matched_tokens(). The
# found blocks were pinned there to survive LRU eviction in the window
# between get_num_new_matched_tokens() and this matching call.
pending = self._pending_cpu_hits.pop(req_id, None)
if num_external_tokens == 0:
if pending is not None:
logger.warning(
"SimpleCPUOffloadScheduler: update_state_after_alloc "
"called for req_id=%s with no external tokens but "
"get_num_new_matched_tokens() unexpectedly recorded "
"a pending CPU hit; releasing the stale pin.",
req_id,
)
self._free_pending_cpu_hit(pending)
return
if pending is None:
logger.warning(
"SimpleCPUOffloadScheduler: update_state_after_alloc called "
"for req_id=%s with num_external_tokens=%d but no pending "
"CPU hit from get_num_new_matched_tokens(); skipping load.",
req_id,
num_external_tokens,
)
return
cpu_hit_blocks_full, _ = pending
num_blocks_to_load = num_external_tokens // self.block_size
assert num_blocks_to_load > 0
skipped = sum(blk.block_hash is not None for blk in blocks.blocks[self.fa_gidx])
num_computed_tokens = skipped * self.block_size
hashes_to_load = request.block_hashes[skipped : skipped + num_blocks_to_load]
# Find CPU cached blocks across all groups.
max_hit_len = len(hashes_to_load) * self.block_size
cpu_hit_blocks, hit_length = self.cpu_coordinator.find_longest_cache_hit(
hashes_to_load, max_hit_len
)
assert hit_length == num_external_tokens, (
f"Expected {num_external_tokens} hit tokens, got {hit_length}"
)
# Build transfer pairs across all groups.
total_computed_tokens = num_computed_tokens + num_external_tokens
kv_cache_groups = self.cpu_kv_cache_config.kv_cache_groups
# The scheduler may have accepted fewer blocks than
# get_num_new_matched_tokens() reported.
# (e.g. due to token budget in test_partial_gpu_prefix_plus_cpu_load).
# Take only the leading N blocks per group matching num_external_tokens;
# the rest will be released along with the temp pin below.
cpu_hit_blocks: list[list[KVCacheBlock]] = []
for g in range(num_groups):
g_block_size = kv_cache_groups[g].kv_cache_spec.block_size
assert num_external_tokens % g_block_size == 0, (
f"num_external_tokens={num_external_tokens} not aligned to "
f"group {g} block_size={g_block_size}"
)
n_take_g = num_external_tokens // g_block_size
cpu_hit_blocks.append(cpu_hit_blocks_full[g][:n_take_g])
gpu_block_ids: list[int] = []
cpu_block_ids: list[int] = []
cpu_blocks_to_touch: list[KVCacheBlock] = []
@@ -304,6 +359,8 @@ class SimpleCPUOffloadScheduler:
# Touch CPU blocks to prevent eviction during async load.
self.cpu_block_pool.touch(cpu_blocks_to_touch)
# Release the temporary pin held since get_num_new_matched_tokens().
self._free_pending_cpu_hit(pending)
# Touch GPU blocks to prevent freeing during async load
assert self._gpu_block_pool is not None
@@ -666,6 +723,12 @@ class SimpleCPUOffloadScheduler:
so the scheduler can free blocks immediately."""
req_id = request.request_id
# Release any temp CPU hit pin from get_num_new_matched_tokens()
# if request is canceled or preempted before update_state_after_alloc()
pending = self._pending_cpu_hits.pop(req_id, None)
if pending is not None:
self._free_pending_cpu_hit(pending)
# Handle load: defer cleanup if load is in-flight
load_state = self._reqs_to_load.get(req_id)
if load_state is not None:
@@ -692,6 +755,15 @@ class SimpleCPUOffloadScheduler:
) -> tuple[bool, dict[str, Any] | None]:
return self.request_finished(request, block_ids=[])
def _free_pending_cpu_hit(self, pending: tuple) -> None:
"""Release the temporary CPU block pin taken in get_num_new_matched_tokens()."""
cpu_hit_blocks, _ = pending
blocks_to_free = [
blk for grp in cpu_hit_blocks for blk in grp if not blk.is_null
]
if blocks_to_free:
self.cpu_block_pool.free_blocks(blocks_to_free)
def _cleanup_load_request(self, req_id: str) -> None:
"""Release all load resources for a request.