@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / fab07e4d0f

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 3, 09:15 UTC

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

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>

Qiuyang Yue committed May 18, 2026 at 21:22 UTC fab07e4d0f7f266643c6ac0dc944f9f433ef2140

2 files changed +189 -12

tests/v1/simple_kv_offload/test_scheduler.py

+105

@@ -1241,3 +1241,108 @@ def test_partial_gpu_prefix_plus_cpu_load() -> None:

1241 assert bid in ext_block_ids, (

1242 f"Load GPU block {bid} should be an ext_comp block, not a comp or new block"

1243 )

1244 +

1245 +

1246 +# ---------------------------------------------------------------------------

1247 +# Test 11: TOCTOU between Phase A and Phase B (regression for #39702)

1248 +# ---------------------------------------------------------------------------

1249 +def test_toctou_cpu_hit_evicted_between_phases_no_crash() -> None:

1250 + """Regression for vllm-project/vllm#39702.

1251 +

1252 + When ``get_num_new_matched_tokens`` (Phase A) reports a CPU cache hit

1253 + of ``N`` tokens but ``update_state_after_alloc`` (Phase B) runs after

1254 + other requests have caused LRU eviction of those exact blocks, the

1255 + second ``find_longest_cache_hit`` call returns 0 while

1256 + ``num_external_tokens`` is still ``N``, triggering

1257 + ``AssertionError: Expected N hit tokens, got 0``.

1258 +

1259 + Setup: ``num_cpu_blocks=5`` (4 usable; null_block takes 1).

1260 + 1. Store req_a's 2 blocks to CPU. CPU: [a0, a1, _, _].

1261 + 2. Phase A on req_b (same prompt as req_a) reports a 2-block hit.

1262 + Without the fix this does NOT pin a0/a1 — they remain at LRU front.

1263 + 3. Store req_c (2) + req_d (2) — 4 more blocks into 4 slots. With

1264 + a0/a1 unpinned and at LRU front, they get evicted. CPU: [c0, c1,

1265 + d0, d1].

1266 + 4. Phase B on req_b: re-searches the CPU coordinator, finds 0

1267 + cached hashes, asserts.

1268 +

1269 + After the fix, Phase A pins the hit blocks so step 3 evicts req_c's

1270 + blocks instead, and Phase B reads the cached ``(cpu_hit_blocks,

1271 + hit_length)`` tuple without re-searching.

1272 + """

1273 + # 5 total = 4 usable (null_block takes 1)

1274 + fix = make_scheduler(num_cpu_blocks=5, num_gpu_blocks=16, lazy=False)

1275 + sched = fix.scheduler

1276 +

1277 + # --- Step 1: Store req_a's 2 blocks to CPU cache ---

1278 + req_a = make_request(num_blocks=2)

1279 + kv_a = _alloc_and_register(fix, req_a, 2)

1280 + sched.update_state_after_alloc(req_a, kv_a, num_external_tokens=0)

1281 + sched_out_a = make_scheduler_output(

1282 + {req_a.request_id: 2 * BLOCK_SIZE},

1283 + new_reqs={req_a.request_id: kv_a.get_block_ids()},

1284 + )

1285 + meta_a = sched.build_connector_meta(sched_out_a)

1286 + assert meta_a.store_event >= 0

1287 + simulate_store_completion(sched, meta_a.store_event)

1288 +

1289 + # --- Step 2: Phase A — req_b (same prompt as req_a) reports a CPU hit ---

1290 + req_b = Request(

1291 + request_id="req-b-toctou",

1292 + prompt_token_ids=req_a.prompt_token_ids,

1293 + sampling_params=req_a.sampling_params,

1294 + pooling_params=None,

1295 + mm_features=None,

1296 + block_hasher=req_a._block_hasher,

1297 + )

1298 + hit_tokens, is_async = sched.get_num_new_matched_tokens(

1299 + req_b, num_computed_tokens=0

1300 + )

1301 + assert hit_tokens == 2 * BLOCK_SIZE, (

1302 + f"Phase A should report 2 blocks of CPU hit, got {hit_tokens}"

1303 + )

1304 + assert is_async is True

1305 +

1306 + # --- Step 3: TOCTOU window — fill CPU cache so LRU evicts req_a's blocks

1307 + # (in production this corresponds to other concurrent requests landing

1308 + # between Phase A and Phase B for req_b). 4 usable slots, req_a occupies 2;

1309 + # req_c (2) + req_d (2) require evicting 2 LRU blocks. ---

1310 + req_c = make_request(num_blocks=2)

1311 + req_d = make_request(num_blocks=2)

1312 + kv_c = _alloc_and_register(fix, req_c, 2)

1313 + kv_d = _alloc_and_register(fix, req_d, 2)

1314 + sched.update_state_after_alloc(req_c, kv_c, num_external_tokens=0)

1315 + sched.update_state_after_alloc(req_d, kv_d, num_external_tokens=0)

1316 + sched_out_pressure = make_scheduler_output(

1317 + {

1318 + req_c.request_id: 2 * BLOCK_SIZE,

1319 + req_d.request_id: 2 * BLOCK_SIZE,

1320 + },

1321 + new_reqs={

1322 + req_c.request_id: kv_c.get_block_ids(),

1323 + req_d.request_id: kv_d.get_block_ids(),

1324 + },

1325 + )

1326 + meta_pressure = sched.build_connector_meta(sched_out_pressure)

1327 + assert meta_pressure.store_event >= 0

1328 + simulate_store_completion(sched, meta_pressure.store_event)

1329 +

1330 + # --- Step 4: Phase B — must not crash ---

1331 + # Before fix: AssertionError: Expected 32 hit tokens, got 0

1332 + # After fix: Phase A pinned the hits → Step 3 evicted req_c instead,

1333 + # Phase B consumes the cached (cpu_hit_blocks, hit_length) tuple.

1334 + gpu_blocks_b = fix.gpu_block_pool.get_new_blocks(2)

1335 + kv_blocks_b = KVCacheBlocks(blocks=(gpu_blocks_b,))

1336 + sched.update_state_after_alloc(req_b, kv_blocks_b, num_external_tokens=hit_tokens)

1337 +

1338 + # --- Step 5: with the fix, the load is queued correctly ---

1339 + sched_out_b = make_scheduler_output(

1340 + {req_b.request_id: 1},

1341 + new_reqs={req_b.request_id: kv_blocks_b.get_block_ids()},

1342 + )

1343 + meta_b = sched.build_connector_meta(sched_out_b)

1344 + assert meta_b.load_event >= 0, (

1345 + "Phase B should queue a load event using the pinned CPU hit blocks"

1346 + )

1347 + assert len(meta_b.load_gpu_blocks) == 2

1348 + assert len(meta_b.load_cpu_blocks) == 2

vllm/v1/simple_kv_offload/manager.py

+84 -12

@@ -131,6 +131,12 @@ class SimpleCPUOffloadScheduler:

131 # the worker reports completions by event index, not request id.

132 self._load_event_to_reqs: dict[int, list[str]] = {}

133

134 + # Pending (cpu_hit_blocks, hit_length) tuples from find_longest_cache_hit,

135 + # kept pinned via touch() while awaiting update_state_after_alloc().

136 + self._pending_cpu_hits: dict[

137 + str, tuple[tuple[list[KVCacheBlock], ...], int]

138 + ] = {}

139 +

140 # Store metadata

141 self._lazy_mode = lazy_offload

142 # Lazy mode: use a cursor to track the last scanned block in the GPU free queue.

@@ -212,7 +218,17 @@ class SimpleCPUOffloadScheduler:

218 def get_num_new_matched_tokens(

219 self, request: "Request", num_computed_tokens: int

220 ) -> tuple[int | None, bool]:

215 - """Return (num_new_tokens, is_async) from consecutive CPU cache hits."""

221 + """Return (num_new_tokens, is_async) from consecutive CPU cache hits.

222 +

223 + Pins found CPU blocks via touch() so they survive LRU eviction until

224 + update_state_after_alloc() consumes them. Any pin from an earlier

225 + call on the same request (e.g. retry after a failed allocate_slots)

226 + is dropped first.

227 + """

228 + stale = self._pending_cpu_hits.pop(request.request_id, None)

229 + if stale is not None:

230 + self._free_pending_cpu_hit(stale)

231 +

232 skipped = num_computed_tokens // self.block_size

233 remaining_hashes = request.block_hashes[skipped:]

234

@@ -223,11 +239,19 @@ class SimpleCPUOffloadScheduler:

239 max_hit_len = request.num_tokens - 1 - num_computed_tokens

240 if max_hit_len <= 0:

241 return 0, False

226 - _, hit_length = self.cpu_coordinator.find_longest_cache_hit(

242 + cpu_hit_blocks, hit_length = self.cpu_coordinator.find_longest_cache_hit(

243 remaining_hashes, max_hit_len

244 )

245

246 if hit_length > 0:

247 + pin_blocks = [

248 + blk for grp in cpu_hit_blocks for blk in grp if not blk.is_null

249 + ]

250 + self.cpu_block_pool.touch(pin_blocks)

251 + self._pending_cpu_hits[request.request_id] = (

252 + cpu_hit_blocks,

253 + hit_length,

254 + )

255 return hit_length, True

256 return 0, False

257

@@ -253,29 +277,60 @@ class SimpleCPUOffloadScheduler:

277 num_stored_blocks=[0] * num_groups,

278 )

279

280 + # Pop the CPU hit cached by get_num_new_matched_tokens(). The

281 + # found blocks were pinned there to survive LRU eviction in the window

282 + # between get_num_new_matched_tokens() and this matching call.

283 + pending = self._pending_cpu_hits.pop(req_id, None)

284 +

285 if num_external_tokens == 0:

286 + if pending is not None:

287 + logger.warning(

288 + "SimpleCPUOffloadScheduler: update_state_after_alloc "

289 + "called for req_id=%s with no external tokens but "

290 + "get_num_new_matched_tokens() unexpectedly recorded "

291 + "a pending CPU hit; releasing the stale pin.",

292 + req_id,

293 + )

294 + self._free_pending_cpu_hit(pending)

295 + return

296 +

297 + if pending is None:

298 + logger.warning(

299 + "SimpleCPUOffloadScheduler: update_state_after_alloc called "

300 + "for req_id=%s with num_external_tokens=%d but no pending "

301 + "CPU hit from get_num_new_matched_tokens(); skipping load.",

302 + req_id,

303 + num_external_tokens,

304 + )

305 return

306

307 + cpu_hit_blocks_full, _ = pending

308 +

309 num_blocks_to_load = num_external_tokens // self.block_size

310 assert num_blocks_to_load > 0

311

312 skipped = sum(blk.block_hash is not None for blk in blocks.blocks[self.fa_gidx])

313 num_computed_tokens = skipped * self.block_size

264 - hashes_to_load = request.block_hashes[skipped : skipped + num_blocks_to_load]

265 -

266 - # Find CPU cached blocks across all groups.

267 - max_hit_len = len(hashes_to_load) * self.block_size

268 - cpu_hit_blocks, hit_length = self.cpu_coordinator.find_longest_cache_hit(

269 - hashes_to_load, max_hit_len

270 - )

271 - assert hit_length == num_external_tokens, (

272 - f"Expected {num_external_tokens} hit tokens, got {hit_length}"

273 - )

314

315 # Build transfer pairs across all groups.

316 total_computed_tokens = num_computed_tokens + num_external_tokens

317 kv_cache_groups = self.cpu_kv_cache_config.kv_cache_groups

318

319 + # The scheduler may have accepted fewer blocks than

320 + # get_num_new_matched_tokens() reported.

321 + # (e.g. due to token budget in test_partial_gpu_prefix_plus_cpu_load).

322 + # Take only the leading N blocks per group matching num_external_tokens;

323 + # the rest will be released along with the temp pin below.

324 + cpu_hit_blocks: list[list[KVCacheBlock]] = []

325 + for g in range(num_groups):

326 + g_block_size = kv_cache_groups[g].kv_cache_spec.block_size

327 + assert num_external_tokens % g_block_size == 0, (

328 + f"num_external_tokens={num_external_tokens} not aligned to "

329 + f"group {g} block_size={g_block_size}"

330 + )

331 + n_take_g = num_external_tokens // g_block_size

332 + cpu_hit_blocks.append(cpu_hit_blocks_full[g][:n_take_g])

333 +

334 gpu_block_ids: list[int] = []

335 cpu_block_ids: list[int] = []

336 cpu_blocks_to_touch: list[KVCacheBlock] = []

@@ -304,6 +359,8 @@ class SimpleCPUOffloadScheduler:

359

360 # Touch CPU blocks to prevent eviction during async load.

361 self.cpu_block_pool.touch(cpu_blocks_to_touch)

362 + # Release the temporary pin held since get_num_new_matched_tokens().

363 + self._free_pending_cpu_hit(pending)

364

365 # Touch GPU blocks to prevent freeing during async load

366 assert self._gpu_block_pool is not None

@@ -666,6 +723,12 @@ class SimpleCPUOffloadScheduler:

723 so the scheduler can free blocks immediately."""

724 req_id = request.request_id

725

726 + # Release any temp CPU hit pin from get_num_new_matched_tokens()

727 + # if request is canceled or preempted before update_state_after_alloc()

728 + pending = self._pending_cpu_hits.pop(req_id, None)

729 + if pending is not None:

730 + self._free_pending_cpu_hit(pending)

731 +

732 # Handle load: defer cleanup if load is in-flight

733 load_state = self._reqs_to_load.get(req_id)

734 if load_state is not None:

@@ -692,6 +755,15 @@ class SimpleCPUOffloadScheduler:

755 ) -> tuple[bool, dict[str, Any] | None]:

756 return self.request_finished(request, block_ids=[])

757

758 + def _free_pending_cpu_hit(self, pending: tuple) -> None:

759 + """Release the temporary CPU block pin taken in get_num_new_matched_tokens()."""

760 + cpu_hit_blocks, _ = pending

761 + blocks_to_free = [

762 + blk for grp in cpu_hit_blocks for blk in grp if not blk.is_null

763 + ]

764 + if blocks_to_free:

765 + self.cpu_block_pool.free_blocks(blocks_to_free)

766 +

767 def _cleanup_load_request(self, req_id: str) -> None:

768 """Release all load resources for a request.

769