[KV Offloading] Add free block iterator for CPU offload scheduling (#47849)
Signed-off-by: chaunceyjiang <chaunceyjiang@gmail.com>
Chauncey committed
Jul 9, 2026 at 14:59 UTC
529af888425040e7424b10237ad5909a71558b99
2 files changed
+18
-17
vllm/v1/core/kv_cache_utils.py
+14
@@ -407,6 +407,20 @@ class FreeKVCacheBlockQueue:
407
curr_block = curr_block.next_free_block
408
return ret
409
410
+ def iter_blocks_after(
411
+ self,
412
+ cursor: KVCacheBlock | None,
413
+ ) -> Iterator[KVCacheBlock]:
414
+ """Iterate free blocks in eviction order after the cursor."""
415
+ if cursor is None:
416
+ curr_block = self.fake_free_list_head.next_free_block
417
+ else:
418
+ curr_block = cursor.next_free_block
419
+
420
+ while curr_block is not None and curr_block is not self.fake_free_list_tail:
421
+ yield curr_block
422
+ curr_block = curr_block.next_free_block
423
+
424
425
def need_extra_keys(request: Request) -> bool:
426
"""Check whether the blocks allocated to this request need extra hash keys.
vllm/v1/simple_kv_offload/manager.py
+4
-17
@@ -487,24 +487,14 @@ class SimpleCPUOffloadScheduler:
487
if self._cursor is not None and self._cursor.ref_cnt > 0:
488
self._cursor = None
489
490
- # Determine start node.
491
- if self._cursor is None:
492
- node = free_queue.fake_free_list_head.next_free_block
493
- else:
494
- node = self._cursor.next_free_block
495
-
496
- tail = free_queue.fake_free_list_tail
490
gpu_ids: list[int] = []
491
block_hashes: list[bytes] = []
499
- covered = 0
492
last_visited = self._cursor
493
502
- while (
503
- node is not None
504
- and node is not tail
505
- and covered < self._target_free
506
- and len(gpu_ids) < num_cpu_free
507
- ):
494
+ for covered, node in enumerate(free_queue.iter_blocks_after(self._cursor)):
495
+ if covered >= self._target_free or len(gpu_ids) >= num_cpu_free:
496
+ break
497
+
498
last_visited = node
499
bhash = node.block_hash
500
@@ -516,9 +506,6 @@ class SimpleCPUOffloadScheduler:
506
gpu_ids.append(node.block_id)
507
block_hashes.append(bhash)
508
519
- covered += 1
520
- node = node.next_free_block
521
-
509
self._cursor = last_visited
510
511
# Batch-allocate CPU blocks and stamp hashes.