diff --git a/tests/v1/simple_kv_offload/test_scheduler.py b/tests/v1/simple_kv_offload/test_scheduler.py index d1bdaf3231..104b0184db 100644 --- a/tests/v1/simple_kv_offload/test_scheduler.py +++ b/tests/v1/simple_kv_offload/test_scheduler.py @@ -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 diff --git a/vllm/v1/simple_kv_offload/manager.py b/vllm/v1/simple_kv_offload/manager.py index 0d06460b30..340cb3c724 100644 --- a/vllm/v1/simple_kv_offload/manager.py +++ b/vllm/v1/simple_kv_offload/manager.py @@ -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.