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

Fork 0

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

[SimpleCPUOffloadConnector]: Add support for reset_cache() (#39726)

Signed-off-by: Jonathan Chen <chenleejonathan@gmail.com> Signed-off-by: Jonathan <chenleejonathan@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

Jonathan Chen committed Jun 17, 2026 at 22:47 UTC d57888efa41b317c34b912c21ca36bc20bdd8da1

3 files changed +251 -9

tests/v1/simple_kv_offload/test_scheduler.py

+174

@@ -1354,3 +1354,177 @@ def test_toctou_cpu_hit_evicted_between_phases_no_crash() -> None:

1354 )

1355 assert len(meta_b.load_gpu_blocks) == 2

1356 assert len(meta_b.load_cpu_blocks) == 2

1357 +

1358 +

1359 +# ---------------------------------------------------------------------------

1360 +# Test 12: Reset with pending eager stores waits for completion

1361 +# ---------------------------------------------------------------------------

1362 +def test_reset_pending_eager_stores() -> None:

1363 + """Eager mode: reset() abandons in-flight stores until they complete."""

1364 + fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=16, lazy=False)

1365 + sched = fix.scheduler

1366 + gpu_pool = fix.gpu_block_pool

1367 +

1368 + num_blocks = 2

1369 + req = make_request(num_blocks=num_blocks)

1370 +

1371 + kv_blocks = _alloc_and_register(fix, req, num_blocks)

1372 + sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0)

1373 + block_ids = kv_blocks.get_block_ids()

1374 + sched_out = make_scheduler_output(

1375 + {req.request_id: num_blocks * BLOCK_SIZE},

1376 + new_reqs={req.request_id: block_ids},

1377 + )

1378 +

1379 + meta = sched.build_connector_meta(sched_out)

1380 + assert meta.store_event >= 0

1381 + assert len(sched._store_event_to_blocks) > 0

1382 +

1383 + # GPU blocks should have elevated ref_cnt from touch()

1384 + for bid in meta.store_gpu_blocks:

1385 + assert gpu_pool.blocks[bid].ref_cnt > 0

1386 +

1387 + # Free the request's own block refs (simulates preemption)

1388 + gpu_pool.free_blocks(gpu_pool.blocks[bid] for bid in block_ids[0])

1389 +

1390 + # Reset should keep DMA refs pinned until the worker reports completion.

1391 + assert sched.reset() is False

1392 + assert len(sched._store_event_to_blocks) == 0

1393 + assert len(sched._abandoned_store_event_to_blocks) == 1

1394 + assert len(sched._reqs_to_store) == 0

1395 + assert len(sched._store_event_to_reqs) == 0

1396 +

1397 + num_used = gpu_pool.num_gpu_blocks - gpu_pool.get_num_free_blocks()

1398 + assert num_used > 1

1399 +

1400 + simulate_store_completion(sched, meta.store_event)

1401 + assert len(sched._abandoned_store_event_to_blocks) == 0

1402 +

1403 + # All GPU blocks should now be free (ref_cnt == 0) except null block

1404 + num_used = gpu_pool.num_gpu_blocks - gpu_pool.get_num_free_blocks()

1405 + assert num_used == 1, f"Expected only null block in use, got {num_used}"

1406 +

1407 + # GPU prefix cache reset should now succeed

1408 + assert gpu_pool.reset_prefix_cache() is True

1409 + assert sched.reset() is True

1410 +

1411 +

1412 +# ---------------------------------------------------------------------------

1413 +# Test 13: Reset with pending lazy stores waits for completion

1414 +# ---------------------------------------------------------------------------

1415 +def test_reset_pending_lazy_stores() -> None:

1416 + """Lazy mode: reset() abandons in-flight stores until they complete."""

1417 + fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=8, lazy=True)

1418 + sched = fix.scheduler

1419 + gpu_pool = fix.gpu_block_pool

1420 +

1421 + num_blocks = 2

1422 + req = make_request(num_blocks=num_blocks)

1423 +

1424 + # Allocate, hash, and free — blocks move to free queue with hashes

1425 + gpu_blocks = _allocate_gpu_blocks(gpu_pool, req, num_blocks, group_id=0)

1426 + gpu_pool.free_blocks(gpu_blocks)

1427 +

1428 + # Push hashed blocks to LRU head

1429 + fillers = _flush_old_blocks_to_lru_head(gpu_pool, num_filler_blocks=5)

1430 +

1431 + # Lazy scanner offloads old hashed blocks

1432 + sched_out = make_scheduler_output({})

1433 + meta = sched.build_connector_meta(sched_out)

1434 + assert meta.store_event >= 0

1435 + assert len(sched._store_event_to_blocks) > 0

1436 +

1437 + gpu_pool.free_blocks(fillers)

1438 +

1439 + # Reset should keep DMA refs pinned until the worker reports completion.

1440 + assert sched.reset() is False

1441 + assert len(sched._store_event_to_blocks) == 0

1442 + assert len(sched._abandoned_store_event_to_blocks) == 1

1443 + assert sched._cursor is None

1444 +

1445 + simulate_store_completion(sched, meta.store_event)

1446 + assert len(sched._abandoned_store_event_to_blocks) == 0

1447 + assert sched.reset() is True

1448 +

1449 + # No CPU cache hits after reset

1450 + req2 = Request(

1451 + request_id="req-after-lazy-reset",

1452 + prompt_token_ids=req.prompt_token_ids,

1453 + sampling_params=req.sampling_params,

1454 + pooling_params=None,

1455 + mm_features=None,

1456 + block_hasher=req._block_hasher,

1457 + )

1458 + hit_tokens, _ = sched.get_num_new_matched_tokens(req2, num_computed_tokens=0)

1459 + assert hit_tokens == 0, "CPU cache should be empty after reset"

1460 +

1461 +

1462 +# ---------------------------------------------------------------------------

1463 +# Test 14: Reset with pending loads waits for completion

1464 +# ---------------------------------------------------------------------------

1465 +def test_reset_pending_loads() -> None:

1466 + """reset() abandons in-flight loads until they complete."""

1467 + fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=16, lazy=False)

1468 + sched = fix.scheduler

1469 + gpu_pool = fix.gpu_block_pool

1470 +

1471 + num_blocks = 2

1472 +

1473 + # First store blocks to CPU

1474 + req = make_request(num_blocks=num_blocks)

1475 + kv_blocks = _alloc_and_register(fix, req, num_blocks)

1476 + sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0)

1477 + block_ids = kv_blocks.get_block_ids()

1478 + sched_out = make_scheduler_output(

1479 + {req.request_id: num_blocks * BLOCK_SIZE},

1480 + new_reqs={req.request_id: block_ids},

1481 + )

1482 + meta = sched.build_connector_meta(sched_out)

1483 + simulate_store_completion(sched, meta.store_event)

1484 +

1485 + # Start a load — CPU cache hit

1486 + req2 = Request(

1487 + request_id="req-load-reset",

1488 + prompt_token_ids=req.prompt_token_ids,

1489 + sampling_params=req.sampling_params,

1490 + pooling_params=None,

1491 + mm_features=None,

1492 + block_hasher=req._block_hasher,

1493 + )

1494 + hit_tokens, is_async = sched.get_num_new_matched_tokens(req2, num_computed_tokens=0)

1495 + assert hit_tokens > 0

1496 +

1497 + gpu_blocks2 = gpu_pool.get_new_blocks(num_blocks)

1498 + kv_blocks2 = KVCacheBlocks(blocks=(gpu_blocks2,))

1499 + sched.update_state_after_alloc(req2, kv_blocks2, num_external_tokens=hit_tokens)

1500 +

1501 + block_ids2 = kv_blocks2.get_block_ids()

1502 + sched_out2 = make_scheduler_output(

1503 + {req2.request_id: 1},

1504 + new_reqs={req2.request_id: block_ids2},

1505 + )

1506 + meta2 = sched.build_connector_meta(sched_out2)

1507 + assert meta2.load_event >= 0

1508 + assert req2.request_id in sched._reqs_to_load

1509 +

1510 + # Free request block refs (simulates preemption)

1511 + gpu_pool.free_blocks(gpu_pool.blocks[bid] for bid in block_ids[0])

1512 + gpu_pool.free_blocks(gpu_pool.blocks[bid] for bid in block_ids2[0])

1513 +

1514 + # Reset should keep load touch refs until the worker reports completion.

1515 + assert sched.reset() is False

1516 + assert len(sched._reqs_to_load) == 0

1517 + assert len(sched._abandoned_reqs_to_load) == 1

1518 + assert len(sched._load_event_to_reqs) == 1

1519 +

1520 + num_used = gpu_pool.num_gpu_blocks - gpu_pool.get_num_free_blocks()

1521 + assert num_used > 1

1522 +

1523 + simulate_load_completion(sched, {req2.request_id})

1524 + assert len(sched._abandoned_reqs_to_load) == 0

1525 + assert len(sched._load_event_to_reqs) == 0

1526 + assert sched.reset() is True

1527 +

1528 + # All GPU blocks free

1529 + num_used = gpu_pool.num_gpu_blocks - gpu_pool.get_num_free_blocks()

1530 + assert num_used == 1, f"Expected only null block in use, got {num_used}"

vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py

+6 -6

@@ -245,10 +245,10 @@ class SimpleCPUOffloadConnector(KVConnectorBase_V1, SupportsHMA):

245 return self.scheduler_manager.take_events()

246 return []

247

248 + # NOTE: Workers are not contacted. In-flight transfers drain naturally,

249 + # and stale completions are ignored by the guarded

250 + # SimpleCPUOffloadScheduler._process_store_event().

251 def reset_cache(self) -> bool | None:

249 - raise NotImplementedError(

250 - "SimpleCPUOffloadConnector does not support reset_cache(). "

251 - "reset_prefix_cache() requires synchronizing all pending "

252 - "CPU offload transfers before clearing GPU prefix cache blocks, "

253 - "which is not yet implemented."

254 - )

252 + if self.scheduler_manager is not None:

253 + return self.scheduler_manager.reset()

254 + return None

vllm/v1/simple_kv_offload/manager.py

+71 -3

@@ -159,10 +159,12 @@ class SimpleCPUOffloadScheduler:

159 else:

160 self._target_free = 0

161 self._store_event_to_blocks: dict[int, TransferMeta] = {}

162 + self._abandoned_store_event_to_blocks: dict[int, TransferMeta] = {}

163 # Eager mode only

164 self._reqs_to_store: dict[str, StoreRequestState] = {}

165 self._store_event_to_reqs: dict[int, list[str]] = {}

166 self._in_flight_store_gpu_blocks: set[int] = set()

167 + self._abandoned_reqs_to_load: dict[str, LoadRequestState] = {}

168

169 # Event counters

170 self._load_event_counter: int = 0

@@ -427,7 +429,10 @@ class SimpleCPUOffloadScheduler:

429 load_event=load_event,

430 load_gpu_blocks=load_gpu,

431 load_cpu_blocks=load_cpu,

430 - load_event_to_reqs=self._load_event_to_reqs,

432 + load_event_to_reqs={

433 + event_idx: list(req_ids)

434 + for event_idx, req_ids in self._load_event_to_reqs.items()

435 + },

436 store_event=store_event,

437 store_gpu_blocks=store_gpu,

438 store_cpu_blocks=store_cpu,

@@ -680,9 +685,17 @@ class SimpleCPUOffloadScheduler:

685

686 def _process_store_event(self, event_idx: int) -> None:

687 """Process a fully-completed store event."""

683 - transfer = self._store_event_to_blocks.pop(event_idx)

688 + transfer = self._store_event_to_blocks.pop(event_idx, None)

689 + if transfer is None:

690 + transfer = self._abandoned_store_event_to_blocks.pop(event_idx, None)

691 + if transfer is None:

692 + return # guard stale events from before a reset() call

693 + self._release_transfer_refs(transfer)

694 + return

695 +

696 if not self._lazy_mode:

697 self._in_flight_store_gpu_blocks.difference_update(transfer.gpu_block_ids)

698 +

699 self._process_store_completion(transfer.gpu_block_ids, transfer.cpu_block_ids)

700 logger.debug(

701 "Store event %d completed: cached %d blocks to CPU",

@@ -725,9 +738,22 @@ class SimpleCPUOffloadScheduler:

738 self._gpu_block_pool.blocks[bid] for bid in gpu_block_ids

739 )

740

741 + def _release_transfer_refs(self, transfer: TransferMeta) -> None:

742 + """Release transfer refs without making copied data cacheable."""

743 + cpu_blocks = [self.cpu_block_pool.blocks[bid] for bid in transfer.cpu_block_ids]

744 + for cpu_block in cpu_blocks:

745 + cpu_block.reset_hash()

746 + self.cpu_block_pool.free_blocks(cpu_blocks)

747 + assert self._gpu_block_pool is not None

748 + self._gpu_block_pool.free_blocks(

749 + self._gpu_block_pool.blocks[bid] for bid in transfer.gpu_block_ids

750 + )

751 +

752 def has_pending_stores(self) -> bool:

753 """Return True if there are in-flight store transfers."""

730 - return bool(self._store_event_to_blocks)

754 + return bool(

755 + self._store_event_to_blocks or self._abandoned_store_event_to_blocks

756 + )

757

758 def request_finished(

759 self,

@@ -787,6 +813,8 @@ class SimpleCPUOffloadScheduler:

813 and frees CPU/GPU touch refs.

814 """

815 state = self._reqs_to_load.pop(req_id, None)

816 + if state is None:

817 + state = self._abandoned_reqs_to_load.pop(req_id, None)

818 if state is None:

819 return

820 # Remove from load event mapping (only this req, not whole event)

@@ -830,3 +858,43 @@ class SimpleCPUOffloadScheduler:

858

859 def take_events(self) -> Iterable[KVCacheEvent]:

860 return self.cpu_block_pool.take_events()

861 +

862 + def reset(self) -> bool:

863 + """Abandon pending transfers and reset the CPU cache when safe.

864 +

865 + Worker-side DMA may still be using blocks after reset is requested.

866 + Keep those block refs pinned until the existing completion path reports

867 + the transfer finished, then release refs without caching abandoned

868 + store results.

869 + """

870 +

871 + self._abandoned_store_event_to_blocks.update(self._store_event_to_blocks)

872 + self._store_event_to_blocks.clear()

873 + self._in_flight_store_gpu_blocks.clear()

874 +

875 + # Loads that have not been sent to the worker cannot have running DMA.

876 + # In-flight loads stay pinned and are cleaned up on completion.

877 + for req_id in list(self._reqs_to_load):

878 + state = self._reqs_to_load.pop(req_id)

879 + if state.load_event is None:

880 + self._reqs_to_load[req_id] = state

881 + self._cleanup_load_request(req_id)

882 + else:

883 + self._abandoned_reqs_to_load[req_id] = state

884 +

885 + self._reqs_to_store.clear()

886 + self._store_event_to_reqs.clear()

887 + self._store_event_pending_counts = {

888 + event_idx: count

889 + for event_idx, count in self._store_event_pending_counts.items()

890 + if event_idx in self._abandoned_store_event_to_blocks

891 + }

892 + self._cursor = None

893 + # NOTE: _load_event_counter / _store_event_counter are not

894 + # reset as they are monotonic and must stay ahead of the workers

895 + # high-water marks to avoid event index collisions

896 +

897 + if self._abandoned_store_event_to_blocks or self._abandoned_reqs_to_load:

898 + return False

899 +

900 + return self.cpu_block_pool.reset_prefix_cache()