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 6, 03:15 UTC
[KV Offload] Refactor CPU offloading: pluggable CachePolicy, remove Backend abstraction, restructure into `cpu/` package (#37874)
Signed-off-by: Ronen Schaffer <ronen.schaffer@ibm.com>
Ronen Schaffer committed
Mar 24, 2026 at 07:02 UTC
e3c6c10cad9b0f9c36b4ca58933316b3e400f1cf
13 files changed
+580
-626
tests/v1/kv_offload/test_cpu_manager.py
+86
-57
@@ -12,9 +12,8 @@ from vllm.v1.kv_offload.abstract import (
12
OffloadingEvent,
13
PrepareStoreOutput,
14
)
15
-from vllm.v1.kv_offload.arc_manager import ARCOffloadingManager
16
-from vllm.v1.kv_offload.backends.cpu import CPUBackend
17
-from vllm.v1.kv_offload.lru_manager import LRUOffloadingManager
15
+from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager
16
+from vllm.v1.kv_offload.cpu.policies.arc import ARCCachePolicy
17
from vllm.v1.kv_offload.mediums import CPULoadStoreSpec
18
19
@@ -79,12 +78,12 @@ def verify_events(
78
assert tuple(stores) == to_hash_sets(expected_stores)
79
80
82
-@pytest.mark.parametrize("manager_class", [LRUOffloadingManager, ARCOffloadingManager])
83
-def test_already_stored_block_not_evicted_during_prepare_store(manager_class):
81
+@pytest.mark.parametrize("eviction_policy", ["lru", "arc"])
82
+def test_already_stored_block_not_evicted_during_prepare_store(eviction_policy):
83
"""
84
Regression test: a block that is already stored must not be evicted
85
by prepare_store() when it needs to make room for new blocks.
87
- Applies to both LRUOffloadingManager and ARCOffloadingManager.
86
+ Applies to both lru and arc policies.
87
88
Scenario:
89
- Store blocks [1, 2] and complete.
@@ -96,8 +95,12 @@ def test_already_stored_block_not_evicted_during_prepare_store(manager_class):
95
- After complete_store([2, 3, 4, 5]), block 2 must still be present.
96
"""
97
block_size = 256
99
- cpu_backend = CPUBackend(block_size=block_size, num_blocks=4)
100
- manager = manager_class(cpu_backend, enable_events=True)
98
+ manager = CPUOffloadingManager(
99
+ block_size=block_size,
100
+ num_blocks=4,
101
+ cache_policy=eviction_policy,
102
+ enable_events=True,
103
+ )
104
105
# store [1, 2] and complete
106
manager.prepare_store(to_hashes([1, 2]))
@@ -129,12 +132,13 @@ def test_already_stored_block_not_evicted_during_prepare_store(manager_class):
132
133
def test_cpu_manager():
134
"""
132
- Tests LRUOffloadingManager with a CPUBackend.
135
+ Tests CPUOffloadingManager with lru policy.
136
"""
137
# initialize a CPU backend with a capacity of 4 blocks
138
block_size = 256
136
- cpu_backend = CPUBackend(block_size=block_size, num_blocks=4)
137
- cpu_manager = LRUOffloadingManager(cpu_backend, enable_events=True)
139
+ cpu_manager = CPUOffloadingManager(
140
+ block_size=block_size, num_blocks=4, cache_policy="lru", enable_events=True
141
+ )
142
143
# prepare store [1, 2]
144
prepare_store_output = cpu_manager.prepare_store(to_hashes([1, 2]))
@@ -241,13 +245,15 @@ def test_cpu_manager():
245
246
def test_arc_manager_basic():
247
"""
244
- Tests ARCOffloadingManager basic operations with a CPUBackend.
248
+ Tests CPUOffloadingManager with arc policy.
249
Verifies that ARC handles store, load, and lookup operations correctly.
250
"""
247
- # initialize a CPU backend with a capacity of 4 blocks
251
block_size = 256
249
- cpu_backend = CPUBackend(block_size=block_size, num_blocks=4)
250
- arc_manager = ARCOffloadingManager(cpu_backend, enable_events=True)
252
+ arc_manager = CPUOffloadingManager(
253
+ block_size=block_size, num_blocks=4, cache_policy="arc", enable_events=True
254
+ )
255
+ arc_policy = arc_manager._policy
256
+ assert isinstance(arc_policy, ARCCachePolicy)
257
258
# prepare store [1, 2]
259
prepare_store_output = arc_manager.prepare_store(to_hashes([1, 2]))
@@ -278,8 +284,8 @@ def test_arc_manager_basic():
284
assert arc_manager.lookup(to_hashes([1, 2, 3])) == 2
285
286
# blocks should be in T1 (recent)
281
- assert len(arc_manager.t1) == 2
282
- assert len(arc_manager.t2) == 0
287
+ assert len(arc_policy.t1) == 2
288
+ assert len(arc_policy.t2) == 0
289
290
291
def test_arc_manager_t1_to_t2_promotion():
@@ -288,23 +294,26 @@ def test_arc_manager_t1_to_t2_promotion():
294
This is a key feature of ARC's adaptive behavior.
295
"""
296
block_size = 256
291
- cpu_backend = CPUBackend(block_size=block_size, num_blocks=4)
292
- arc_manager = ARCOffloadingManager(cpu_backend, enable_events=False)
297
+ arc_manager = CPUOffloadingManager(
298
+ block_size=block_size, num_blocks=4, cache_policy="arc", enable_events=False
299
+ )
300
+ arc_policy = arc_manager._policy
301
+ assert isinstance(arc_policy, ARCCachePolicy)
302
303
# store and complete block 1
304
arc_manager.prepare_store(to_hashes([1]))
305
arc_manager.complete_store(to_hashes([1]))
306
307
# block 1 starts in T1 (recent)
299
- assert to_hashes([1])[0] in arc_manager.t1
300
- assert to_hashes([1])[0] not in arc_manager.t2
308
+ assert to_hashes([1])[0] in arc_policy.t1
309
+ assert to_hashes([1])[0] not in arc_policy.t2
310
311
# touch block 1 (simulate second access)
312
arc_manager.touch(to_hashes([1]))
313
314
# block 1 should now be in T2 (frequent)
306
- assert to_hashes([1])[0] not in arc_manager.t1
307
- assert to_hashes([1])[0] in arc_manager.t2
315
+ assert to_hashes([1])[0] not in arc_policy.t1
316
+ assert to_hashes([1])[0] in arc_policy.t2
317
318
319
def test_arc_manager_eviction_with_load():
@@ -313,8 +322,9 @@ def test_arc_manager_eviction_with_load():
322
Verifies that blocks being loaded (ref_cnt > 0) cannot be evicted.
323
"""
324
block_size = 256
316
- cpu_backend = CPUBackend(block_size=block_size, num_blocks=4)
317
- arc_manager = ARCOffloadingManager(cpu_backend, enable_events=True)
325
+ arc_manager = CPUOffloadingManager(
326
+ block_size=block_size, num_blocks=4, cache_policy="arc", enable_events=True
327
+ )
328
329
# prepare and complete store [1, 2, 3, 4]
330
prepare_store_output = arc_manager.prepare_store(to_hashes([1, 2, 3, 4]))
@@ -354,28 +364,31 @@ def test_arc_manager_adaptive_target():
364
When a block in B2 is accessed, target_t1_size decreases.
365
"""
366
block_size = 256
357
- cpu_backend = CPUBackend(block_size=block_size, num_blocks=2)
358
- arc_manager = ARCOffloadingManager(cpu_backend, enable_events=False)
367
+ arc_manager = CPUOffloadingManager(
368
+ block_size=block_size, num_blocks=2, cache_policy="arc", enable_events=False
369
+ )
370
+ arc_policy = arc_manager._policy
371
+ assert isinstance(arc_policy, ARCCachePolicy)
372
373
# store blocks 1, 2 (fills cache)
374
arc_manager.prepare_store(to_hashes([1, 2]))
375
arc_manager.complete_store(to_hashes([1, 2]))
376
364
- initial_target = arc_manager.target_t1_size
377
+ initial_target = arc_policy.target_t1_size
378
379
# store block 3, evicting block 1 (moves to B1 ghost list)
380
arc_manager.prepare_store(to_hashes([3]))
381
arc_manager.complete_store(to_hashes([3]))
382
383
# block 1 should be in B1 (ghost list)
371
- assert to_hashes([1])[0] in arc_manager.b1
384
+ assert to_hashes([1])[0] in arc_policy.b1
385
386
# touch block 1 (cache miss, but in B1)
387
# this should increase target_t1_size (favor recency)
388
arc_manager.touch(to_hashes([1]))
389
390
# target should have increased
378
- assert arc_manager.target_t1_size > initial_target
391
+ assert arc_policy.target_t1_size > initial_target
392
393
394
def test_arc_manager_t1_t2_eviction_policy():
@@ -384,8 +397,11 @@ def test_arc_manager_t1_t2_eviction_policy():
397
If |T1| >= target_t1_size, evict from T1, otherwise from T2.
398
"""
399
block_size = 256
387
- cpu_backend = CPUBackend(block_size=block_size, num_blocks=4)
388
- arc_manager = ARCOffloadingManager(cpu_backend, enable_events=False)
400
+ arc_manager = CPUOffloadingManager(
401
+ block_size=block_size, num_blocks=4, cache_policy="arc", enable_events=False
402
+ )
403
+ arc_policy = arc_manager._policy
404
+ assert isinstance(arc_policy, ARCCachePolicy)
405
406
# store blocks 1, 2, 3, 4
407
arc_manager.prepare_store(to_hashes([1, 2, 3, 4]))
@@ -395,12 +411,12 @@ def test_arc_manager_t1_t2_eviction_policy():
411
arc_manager.touch(to_hashes([3, 4]))
412
413
# now: T1 = {1, 2}, T2 = {3, 4}
398
- assert len(arc_manager.t1) == 2
399
- assert len(arc_manager.t2) == 2
414
+ assert len(arc_policy.t1) == 2
415
+ assert len(arc_policy.t2) == 2
416
417
# set target_t1_size to prefer evicting from T1
418
# (when |T1| >= target, evict from T1)
403
- arc_manager.target_t1_size = 1
419
+ arc_policy.target_t1_size = 1
420
421
# store block 5, should evict from T1 (block 1, LRU in T1)
422
output = arc_manager.prepare_store(to_hashes([5]))
@@ -410,9 +426,9 @@ def test_arc_manager_t1_t2_eviction_policy():
426
arc_manager.complete_store(to_hashes([5]))
427
428
# block 1 should be in B1 (ghost list)
413
- assert to_hashes([1])[0] in arc_manager.b1
429
+ assert to_hashes([1])[0] in arc_policy.b1
430
# block 5 should be in T1
415
- assert to_hashes([5])[0] in arc_manager.t1
431
+ assert to_hashes([5])[0] in arc_policy.t1
432
433
434
def test_arc_manager_ghost_list_bounds():
@@ -421,8 +437,11 @@ def test_arc_manager_ghost_list_bounds():
437
They should be capped at cache_capacity.
438
"""
439
block_size = 256
424
- cpu_backend = CPUBackend(block_size=block_size, num_blocks=2)
425
- arc_manager = ARCOffloadingManager(cpu_backend, enable_events=False)
440
+ arc_manager = CPUOffloadingManager(
441
+ block_size=block_size, num_blocks=2, cache_policy="arc", enable_events=False
442
+ )
443
+ arc_policy = arc_manager._policy
444
+ assert isinstance(arc_policy, ARCCachePolicy)
445
446
# fill cache with blocks 1, 2
447
arc_manager.prepare_store(to_hashes([1, 2]))
@@ -434,8 +453,8 @@ def test_arc_manager_ghost_list_bounds():
453
arc_manager.complete_store(to_hashes([i]))
454
455
# ghost lists should not exceed cache_capacity
437
- assert len(arc_manager.b1) <= arc_manager.cache_capacity
438
- assert len(arc_manager.b2) <= arc_manager.cache_capacity
456
+ assert len(arc_policy.b1) <= arc_policy.cache_capacity
457
+ assert len(arc_policy.b2) <= arc_policy.cache_capacity
458
459
460
def test_arc_manager_touch_ordering():
@@ -444,8 +463,11 @@ def test_arc_manager_touch_ordering():
463
Similar to LRU test but verifies T1/T2 ordering.
464
"""
465
block_size = 256
447
- cpu_backend = CPUBackend(block_size=block_size, num_blocks=4)
448
- arc_manager = ARCOffloadingManager(cpu_backend, enable_events=True)
466
+ arc_manager = CPUOffloadingManager(
467
+ block_size=block_size, num_blocks=4, cache_policy="arc", enable_events=True
468
+ )
469
+ arc_policy = arc_manager._policy
470
+ assert isinstance(arc_policy, ARCCachePolicy)
471
472
# store blocks 1, 2, 3, 4
473
arc_manager.prepare_store(to_hashes([1, 2, 3, 4]))
@@ -459,8 +481,8 @@ def test_arc_manager_touch_ordering():
481
arc_manager.touch(to_hashes([1, 3, 4]))
482
483
# T1 = {2}, T2 = {1, 3, 4} (in that order, with 4 most recent)
462
- assert len(arc_manager.t1) == 1
463
- assert len(arc_manager.t2) == 3
484
+ assert len(arc_policy.t1) == 1
485
+ assert len(arc_policy.t2) == 3
486
487
# store block 5, should evict from T1 (block 2, only one in T1)
488
prepare_store_output = arc_manager.prepare_store(to_hashes([5]))
@@ -480,8 +502,11 @@ def test_arc_manager_failed_store():
502
Similar to LRU test but for ARC.
503
"""
504
block_size = 256
483
- cpu_backend = CPUBackend(block_size=block_size, num_blocks=4)
484
- arc_manager = ARCOffloadingManager(cpu_backend, enable_events=True)
505
+ arc_manager = CPUOffloadingManager(
506
+ block_size=block_size, num_blocks=4, cache_policy="arc", enable_events=True
507
+ )
508
+ arc_policy = arc_manager._policy
509
+ assert isinstance(arc_policy, ARCCachePolicy)
510
511
# store blocks 1, 2, 3, 4
512
arc_manager.prepare_store(to_hashes([1, 2, 3, 4]))
@@ -498,12 +523,12 @@ def test_arc_manager_failed_store():
523
# block 5 should not be in cache
524
assert arc_manager.lookup(to_hashes([5])) == 0
525
# block 5 should not be in T1 or T2
501
- assert to_hashes([5])[0] not in arc_manager.t1
502
- assert to_hashes([5])[0] not in arc_manager.t2
526
+ assert to_hashes([5])[0] not in arc_policy.t1
527
+ assert to_hashes([5])[0] not in arc_policy.t2
528
529
# evicted block should still be gone (in B1 ghost list)
530
evicted_hash = prepare_store_output.block_hashes_evicted[0]
506
- assert evicted_hash in arc_manager.b1
531
+ assert evicted_hash in arc_policy.b1
532
533
534
def test_arc_manager_full_scenario():
@@ -512,8 +537,11 @@ def test_arc_manager_full_scenario():
537
Similar to the full LRU test but adapted for ARC behavior.
538
"""
539
block_size = 256
515
- cpu_backend = CPUBackend(block_size=block_size, num_blocks=4)
516
- arc_manager = ARCOffloadingManager(cpu_backend, enable_events=True)
540
+ arc_manager = CPUOffloadingManager(
541
+ block_size=block_size, num_blocks=4, cache_policy="arc", enable_events=True
542
+ )
543
+ arc_policy = arc_manager._policy
544
+ assert isinstance(arc_policy, ARCCachePolicy)
545
546
# store [1, 2]
547
arc_manager.prepare_store(to_hashes([1, 2]))
@@ -529,8 +557,8 @@ def test_arc_manager_full_scenario():
557
arc_manager.touch(to_hashes([2, 3]))
558
559
# T1 has {4, 5}, T2 has {2, 3}
532
- assert len(arc_manager.t1) == 2
533
- assert len(arc_manager.t2) == 2
560
+ assert len(arc_policy.t1) == 2
561
+ assert len(arc_policy.t2) == 2
562
563
# store [6] -> should evict from T1 (4 is oldest in T1)
564
prepare_store_output = arc_manager.prepare_store(to_hashes([6]))
@@ -548,11 +576,12 @@ def test_arc_manager_full_scenario():
576
577
def test_filter_reused_manager():
578
"""
551
- Tests FilterReusedOffloadingManager with a CPUBackend.
579
+ Tests FilterReusedOffloadingManager with a CPUOffloadingManager.
580
"""
581
block_size = 256
554
- cpu_backend = CPUBackend(block_size=block_size, num_blocks=4)
555
- lru_manager = LRUOffloadingManager(cpu_backend, enable_events=True)
582
+ lru_manager = CPUOffloadingManager(
583
+ block_size=block_size, num_blocks=4, cache_policy="lru", enable_events=True
584
+ )
585
586
from vllm.v1.kv_offload.reuse_manager import FilterReusedOffloadingManager
587
vllm/v1/kv_offload/arc_manager.py
deleted
-244
@@ -1,244 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-from collections import OrderedDict
4
-from collections.abc import Iterable
5
-
6
-from vllm.v1.core.kv_cache_utils import BlockHash
7
-from vllm.v1.kv_offload.abstract import (
8
- LoadStoreSpec,
9
- OffloadingEvent,
10
- OffloadingManager,
11
- PrepareStoreOutput,
12
-)
13
-from vllm.v1.kv_offload.backend import Backend, BlockStatus
14
-
15
-
16
-class ARCOffloadingManager(OffloadingManager):
17
- """
18
- An OffloadingManager implementing the ARC (Adaptive Replacement Cache)
19
- eviction policy with a pluggable backend.
20
-
21
- Data Structures:
22
- T1: Recent cache containing blocks accessed once.
23
- T2: Frequent cache containing blocks accessed multiple times.
24
- B1/B2: Ghost lists tracking recently evicted blocks from T1/T2.
25
- target_t1_size: Adaptive target size for the T1 partition.
26
-
27
- Algorithm Flow:
28
- 1. Cache lookup (lookup):
29
- Searches T1 and T2 for block hashes and counts consecutive hits
30
- until a miss or non-ready block is encountered.
31
-
32
- 2. Cache touch (touch) - Adaptive Learning:
33
- For each block_hash (in reverse order):
34
- - If in T1: Move to T2 (promotion from recent to frequent).
35
- - If in T2: Move to MRU position (end of queue).
36
- - If in B1 ghost list: Increase target_t1_size.
37
- - If in B2 ghost list: Decrease target_t1_size.
38
-
39
- 3. Block eviction (prepare_store) - Adaptive Replacement:
40
- Determines eviction source based on adaptive target:
41
- - If T1 size > target_t1_size: Evict from T1, add to B1.
42
- - Otherwise: Evict from T2, add to B2.
43
- Finally, bound each ghost list size.
44
-
45
- 4. Block insertion (prepare_store):
46
- New blocks are always inserted into T1 and removed from B1/B2 if
47
- present. Blocks may later be promoted to T2 during touch operations.
48
-
49
- Adaptive Behavior:
50
- The algorithm self-tunes the recency vs. frequency trade-off:
51
- - B1 hit: Recent access patterns matter more → increase T1.
52
- - B2 hit: Frequent access patterns matter more → decrease T1.
53
- """
54
-
55
- def __init__(self, backend: Backend, enable_events: bool = False):
56
- self.backend: Backend = backend
57
- self.target_t1_size: float = 0.0
58
- self.t1: OrderedDict[BlockHash, BlockStatus] = OrderedDict()
59
- self.t2: OrderedDict[BlockHash, BlockStatus] = OrderedDict()
60
- # block_hash -> None (only care about presence)
61
- self.b1: OrderedDict[BlockHash, None] = OrderedDict()
62
- self.b2: OrderedDict[BlockHash, None] = OrderedDict()
63
- self.events: list[OffloadingEvent] | None = [] if enable_events else None
64
- self.cache_capacity: int = self.backend.get_num_free_blocks()
65
-
66
- def lookup(self, block_hashes: Iterable[BlockHash]) -> int | None:
67
- hit_count = 0
68
- for block_hash in block_hashes:
69
- block = self.t1.get(block_hash) or self.t2.get(block_hash)
70
- if block is None or not block.is_ready:
71
- break
72
- hit_count += 1
73
- return hit_count
74
-
75
- def prepare_load(self, block_hashes: Iterable[BlockHash]) -> LoadStoreSpec:
76
- blocks = []
77
- for block_hash in block_hashes:
78
- block = self.t1.get(block_hash) or self.t2.get(block_hash)
79
- assert block is not None, f"Block {block_hash!r} not found in cache"
80
- assert block.is_ready, f"Block {block_hash!r} is not ready for reading"
81
-
82
- block.ref_cnt += 1
83
- blocks.append(block)
84
-
85
- return self.backend.get_load_store_spec(block_hashes, blocks)
86
-
87
- def touch(self, block_hashes: Iterable[BlockHash]):
88
- for block_hash in reversed(list(block_hashes)):
89
- if block_hash in self.t1:
90
- block = self.t1.pop(block_hash)
91
- if not block.is_ready:
92
- # block was just prepared to be stored, not really touched twice
93
- # keep it in T1 and mark as most recently used
94
- self.t1[block_hash] = block
95
- else:
96
- self.t2[block_hash] = block
97
-
98
- elif block_hash in self.t2:
99
- self.t2.move_to_end(block_hash)
100
-
101
- elif block_hash in self.b1:
102
- delta = max(1, len(self.b2) / len(self.b1))
103
- self.target_t1_size = min(
104
- self.target_t1_size + delta, self.cache_capacity
105
- )
106
- # move to MRU position (end) to keep it fresh in the ghost list
107
- self.b1.move_to_end(block_hash)
108
-
109
- elif block_hash in self.b2:
110
- delta = max(1, len(self.b1) / len(self.b2))
111
- self.target_t1_size = max(self.target_t1_size - delta, 0)
112
- # move to MRU position (end) to keep it fresh in the ghost list
113
- self.b2.move_to_end(block_hash)
114
-
115
- def complete_load(self, block_hashes: Iterable[BlockHash]):
116
- for block_hash in block_hashes:
117
- block = self.t1.get(block_hash) or self.t2.get(block_hash)
118
- assert block is not None, f"Block {block_hash!r} not found"
119
- assert block.ref_cnt > 0, f"Block {block_hash!r} ref_cnt is already 0"
120
-
121
- block.ref_cnt -= 1
122
-
123
- def prepare_store(
124
- self, block_hashes: Iterable[BlockHash]
125
- ) -> PrepareStoreOutput | None:
126
- block_hashes_list = list(block_hashes)
127
-
128
- block_hashes_to_store = []
129
- for block_hash in block_hashes_list:
130
- if block_hash not in self.t1 and block_hash not in self.t2:
131
- block_hashes_to_store.append(block_hash)
132
-
133
- if not block_hashes_to_store:
134
- return PrepareStoreOutput(
135
- block_hashes_to_store=[],
136
- store_spec=self.backend.get_load_store_spec([], []),
137
- block_hashes_evicted=[],
138
- )
139
-
140
- num_blocks_to_evict = (
141
- len(block_hashes_to_store) - self.backend.get_num_free_blocks()
142
- )
143
-
144
- to_evict = []
145
- if num_blocks_to_evict > 0:
146
- # Blocks from the original input are excluded from eviction candidates:
147
- # a block that was already stored must remain in the cache after this call.
148
- protected = set(block_hashes_list)
149
- while num_blocks_to_evict > 0:
150
- block_to_evict = None
151
- if len(self.t1) >= int(self.target_t1_size):
152
- # try to evict the least recently used (oldest) block from T1
153
- for block_hash, block in self.t1.items():
154
- if block.ref_cnt == 0 and block_hash not in protected:
155
- block_to_evict = (block_hash, block)
156
- eviction_t = self.t1
157
- eviction_b = self.b1
158
- break
159
- if not block_to_evict:
160
- # try to evict the least recently used (oldest) block from T2
161
- for block_hash, block in self.t2.items():
162
- if block.ref_cnt == 0 and block_hash not in protected:
163
- block_to_evict = (block_hash, block)
164
- eviction_t = self.t2
165
- eviction_b = self.b2
166
- break
167
- else:
168
- # cannot evict enough blocks, cache is full of in-use items
169
- return None
170
-
171
- block_hash, block = block_to_evict
172
- del eviction_t[block_hash]
173
- eviction_b[block_hash] = None
174
- to_evict.append(block_hash)
175
- self.backend.free(block)
176
- num_blocks_to_evict -= 1
177
-
178
- for b in [self.b1, self.b2]:
179
- for i in range(len(b) - self.cache_capacity):
180
- b.popitem(last=False)
181
-
182
- if to_evict and self.events is not None:
183
- self.events.append(
184
- OffloadingEvent(
185
- block_hashes=to_evict,
186
- block_size=self.backend.block_size,
187
- medium=self.backend.medium,
188
- removed=True,
189
- )
190
- )
191
-
192
- blocks = self.backend.allocate_blocks(block_hashes_to_store)
193
- assert len(blocks) == len(block_hashes_to_store), (
194
- "Backend did not allocate the expected number of blocks"
195
- )
196
-
197
- for block_hash, block in zip(block_hashes_to_store, blocks):
198
- self.t1[block_hash] = block
199
-
200
- self.b1.pop(block_hash, None)
201
- self.b2.pop(block_hash, None)
202
-
203
- store_spec = self.backend.get_load_store_spec(block_hashes_to_store, blocks)
204
-
205
- return PrepareStoreOutput(
206
- block_hashes_to_store=block_hashes_to_store,
207
- store_spec=store_spec,
208
- block_hashes_evicted=to_evict,
209
- )
210
-
211
- def complete_store(self, block_hashes: Iterable[BlockHash], success: bool = True):
212
- stored_block_hashes: list[BlockHash] = []
213
-
214
- if success:
215
- for block_hash in block_hashes:
216
- block = self.t1.get(block_hash) or self.t2.get(block_hash)
217
-
218
- if block is not None and not block.is_ready:
219
- block.ref_cnt = 0
220
- stored_block_hashes.append(block_hash)
221
- else:
222
- for block_hash in block_hashes:
223
- block = self.t1.pop(block_hash, None)
224
-
225
- if block is None:
226
- block = self.t2.pop(block_hash, None)
227
-
228
- if block is not None and not block.is_ready:
229
- self.backend.free(block)
230
-
231
- if stored_block_hashes and self.events is not None:
232
- self.events.append(
233
- OffloadingEvent(
234
- block_hashes=stored_block_hashes,
235
- block_size=self.backend.block_size,
236
- medium=self.backend.medium,
237
- removed=False,
238
- )
239
- )
240
-
241
- def take_events(self) -> Iterable[OffloadingEvent]:
242
- if self.events is not None:
243
- yield from self.events
244
- self.events.clear()
vllm/v1/kv_offload/backend.py
deleted
-97
@@ -1,97 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-import ctypes
4
-from abc import ABC, abstractmethod
5
-from collections.abc import Iterable
6
-
7
-from vllm.v1.core.kv_cache_utils import BlockHash
8
-from vllm.v1.kv_offload.abstract import LoadStoreSpec
9
-
10
-
11
-class BlockStatus(ctypes.Structure):
12
- """
13
- Offloading status for a single block of KV data.
14
- Holds the following information:
15
-
16
- ref_cnt - the current number of transfers using this block as a source.
17
- A value of -1 indicates the block is not yet ready to be read.
18
- load_store_spec - backend-specific information on how to actually
19
- read/write the block.
20
- """
21
-
22
- _fields_ = [("ref_cnt", ctypes.c_int32)]
23
-
24
- def __init__(self):
25
- super().__init__()
26
- # initialize block as "not ready" (ref_cnt = -1)
27
- self.ref_cnt = -1
28
-
29
- @property
30
- def is_ready(self) -> bool:
31
- """
32
- Returns whether the block is ready to be read.
33
- """
34
- return self.ref_cnt >= 0
35
-
36
-
37
-class Backend(ABC):
38
- """
39
- An abstract class for allocating and returning specs for writing
40
- KV blocks to some backend.
41
- """
42
-
43
- def __init__(self, block_size: int, medium: str):
44
- self.block_size = block_size
45
- self.medium = medium
46
-
47
- @abstractmethod
48
- def get_num_free_blocks(self):
49
- """
50
- Returns the number of current number of blocks that can be allocated.
51
- """
52
- pass
53
-
54
- @abstractmethod
55
- def allocate_blocks(self, block_hashes: list[BlockHash]) -> list[BlockStatus]:
56
- """
57
- Allocate space for writing blocks.
58
- This method assumes there is enough space for allocation.
59
- It is unsafe to use without checking get_num_free_blocks beforehand.
60
-
61
- Args:
62
- block_hashes: the hashes identifying the blocks to be written.
63
-
64
- Returns:
65
- A list of BlockStatus for the allocated blocks.
66
- The ref_cnt of each returned item will be -1, meaning the block
67
- is not yet ready to be read.
68
- """
69
- pass
70
-
71
- @abstractmethod
72
- def free(self, block: BlockStatus):
73
- """
74
- Free a previously allocated block.
75
- You should only call this function with blocks returned by
76
- allocate_blocks, and only once per each block.
77
-
78
- Args:
79
- block: The block to be freed.
80
- """
81
- pass
82
-
83
- def get_load_store_spec(
84
- self, block_hashes: Iterable[BlockHash], blocks: Iterable[BlockStatus]
85
- ) -> LoadStoreSpec:
86
- """
87
- Get backend-specific information on how to read/write blocks.
88
-
89
- Args:
90
- block_hashes: the list of block hashes identifying the blocks.
91
- blocks: the list of blocks.
92
-
93
- Returns:
94
- A LoadStoreSpec that can be used by a worker
95
- to read/write the blocks.
96
- """
97
- raise NotImplementedError
vllm/v1/kv_offload/backends/cpu.py
deleted
-62
@@ -1,62 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-import ctypes
4
-from collections.abc import Iterable
5
-
6
-from vllm.v1.core.kv_cache_utils import BlockHash
7
-from vllm.v1.kv_offload.abstract import LoadStoreSpec
8
-from vllm.v1.kv_offload.backend import Backend, BlockStatus
9
-from vllm.v1.kv_offload.mediums import CPULoadStoreSpec
10
-
11
-
12
-class CPUBlockStatus(BlockStatus):
13
- _fields_ = BlockStatus._fields_ + [("block_id", ctypes.c_int64)] # type: ignore
14
-
15
- def __init__(self, block_id: int):
16
- super().__init__()
17
- self.block_id = block_id
18
-
19
-
20
-class CPUBackend(Backend):
21
- def __init__(self, block_size: int, num_blocks: int):
22
- super().__init__(block_size=block_size, medium=CPULoadStoreSpec.medium())
23
-
24
- self.num_blocks: int = num_blocks
25
- self.num_allocated_blocks: int = 0
26
- self.allocated_blocks_free_list: list[int] = []
27
-
28
- def get_num_free_blocks(self):
29
- return (
30
- len(self.allocated_blocks_free_list)
31
- + self.num_blocks
32
- - self.num_allocated_blocks
33
- )
34
-
35
- def allocate_blocks(self, block_hashes: list[BlockHash]) -> list[BlockStatus]:
36
- num_fresh_blocks = min(
37
- len(block_hashes), self.num_blocks - self.num_allocated_blocks
38
- )
39
- num_reused_blocks = len(block_hashes) - num_fresh_blocks
40
- assert len(self.allocated_blocks_free_list) >= num_reused_blocks
41
-
42
- # allocate fresh blocks
43
- blocks: list[BlockStatus] = []
44
- for _ in range(num_fresh_blocks):
45
- blocks.append(CPUBlockStatus(self.num_allocated_blocks))
46
- self.num_allocated_blocks += 1
47
-
48
- # allocate reused blocks
49
- for _ in range(num_reused_blocks):
50
- block_id = self.allocated_blocks_free_list.pop()
51
- blocks.append(CPUBlockStatus(block_id))
52
-
53
- return blocks
54
-
55
- def free(self, block: BlockStatus):
56
- assert isinstance(block, CPUBlockStatus)
57
- self.allocated_blocks_free_list.append(block.block_id)
58
-
59
- def get_load_store_spec(
60
- self, block_hashes: Iterable[BlockHash], blocks: Iterable[BlockStatus]
61
- ) -> LoadStoreSpec:
62
- return CPULoadStoreSpec([block.block_id for block in blocks])
vllm/v1/kv_offload/cpu/__init__.py
renamed
vllm/v1/kv_offload/cpu/manager.py
new
+208
@@ -0,0 +1,208 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from collections.abc import Iterable
4
+from typing import Literal
5
+
6
+from vllm.v1.core.kv_cache_utils import BlockHash
7
+from vllm.v1.kv_offload.abstract import (
8
+ LoadStoreSpec,
9
+ OffloadingEvent,
10
+ OffloadingManager,
11
+ PrepareStoreOutput,
12
+)
13
+from vllm.v1.kv_offload.cpu.policies.abstract import BlockStatus, CachePolicy
14
+from vllm.v1.kv_offload.cpu.policies.arc import ARCCachePolicy
15
+from vllm.v1.kv_offload.cpu.policies.lru import LRUCachePolicy
16
+from vllm.v1.kv_offload.mediums import CPULoadStoreSpec
17
+
18
+_CACHE_POLICIES: dict[str, type[CachePolicy]] = {
19
+ "lru": LRUCachePolicy,
20
+ "arc": ARCCachePolicy,
21
+}
22
+
23
+
24
+class CPUOffloadingManager(OffloadingManager):
25
+ """
26
+ An OffloadingManager with a pluggable CachePolicy (LRU or ARC).
27
+
28
+ The manager owns all shared logic: ref-counting, event emission,
29
+ block pool management, and the prepare_store/complete_store skeletons.
30
+ Policy-specific block organization and eviction decisions are delegated
31
+ to the CachePolicy implementation.
32
+ """
33
+
34
+ def __init__(
35
+ self,
36
+ block_size: int,
37
+ num_blocks: int,
38
+ cache_policy: Literal["lru", "arc"] = "lru",
39
+ enable_events: bool = False,
40
+ ):
41
+ self.block_size: int = block_size
42
+ self.medium: str = CPULoadStoreSpec.medium()
43
+ self._num_blocks: int = num_blocks
44
+ self._num_allocated_blocks: int = 0
45
+ self._free_list: list[int] = []
46
+ self.events: list[OffloadingEvent] | None = [] if enable_events else None
47
+ policy_cls = _CACHE_POLICIES.get(cache_policy)
48
+ if policy_cls is None:
49
+ raise ValueError(
50
+ f"Unknown cache policy: {cache_policy!r}. "
51
+ f"Supported: {list(_CACHE_POLICIES)}"
52
+ )
53
+ self._policy: CachePolicy = policy_cls(cache_capacity=num_blocks)
54
+
55
+ # --- block pool ---
56
+
57
+ def _get_num_free_blocks(self) -> int:
58
+ return len(self._free_list) + self._num_blocks - self._num_allocated_blocks
59
+
60
+ def _allocate_blocks(self, block_hashes: list[BlockHash]) -> list[BlockStatus]:
61
+ num_fresh = min(
62
+ len(block_hashes), self._num_blocks - self._num_allocated_blocks
63
+ )
64
+ num_reused = len(block_hashes) - num_fresh
65
+ assert len(self._free_list) >= num_reused
66
+
67
+ # allocate fresh blocks
68
+ blocks: list[BlockStatus] = []
69
+ for _ in range(num_fresh):
70
+ blocks.append(BlockStatus(self._num_allocated_blocks))
71
+ self._num_allocated_blocks += 1
72
+
73
+ # allocate reused blocks
74
+ for _ in range(num_reused):
75
+ blocks.append(BlockStatus(self._free_list.pop()))
76
+ return blocks
77
+
78
+ def _free_block(self, block: BlockStatus) -> None:
79
+ self._free_list.append(block.block_id)
80
+
81
+ def _get_load_store_spec(
82
+ self,
83
+ block_hashes: Iterable[BlockHash],
84
+ blocks: Iterable[BlockStatus],
85
+ ) -> CPULoadStoreSpec:
86
+ return CPULoadStoreSpec([block.block_id for block in blocks])
87
+
88
+ # --- OffloadingManager interface ---
89
+
90
+ def lookup(self, block_hashes: Iterable[BlockHash]) -> int | None:
91
+ hit_count = 0
92
+ for block_hash in block_hashes:
93
+ block = self._policy.get(block_hash)
94
+ if block is None or not block.is_ready:
95
+ break
96
+ hit_count += 1
97
+ return hit_count
98
+
99
+ def prepare_load(self, block_hashes: Iterable[BlockHash]) -> LoadStoreSpec:
100
+ blocks = []
101
+ for block_hash in block_hashes:
102
+ block = self._policy.get(block_hash)
103
+ assert block is not None, f"Block {block_hash!r} not found in cache"
104
+ assert block.is_ready, f"Block {block_hash!r} is not ready for reading"
105
+ block.ref_cnt += 1
106
+ blocks.append(block)
107
+ return self._get_load_store_spec(block_hashes, blocks)
108
+
109
+ def touch(self, block_hashes: Iterable[BlockHash]) -> None:
110
+ self._policy.touch(block_hashes)
111
+
112
+ def complete_load(self, block_hashes: Iterable[BlockHash]) -> None:
113
+ for block_hash in block_hashes:
114
+ block = self._policy.get(block_hash)
115
+ assert block is not None, f"Block {block_hash!r} not found"
116
+ assert block.ref_cnt > 0, f"Block {block_hash!r} ref_cnt is already 0"
117
+ block.ref_cnt -= 1
118
+
119
+ def prepare_store(
120
+ self, block_hashes: Iterable[BlockHash]
121
+ ) -> PrepareStoreOutput | None:
122
+ block_hashes_list = list(block_hashes)
123
+
124
+ # filter out blocks that are already stored
125
+ block_hashes_to_store = [
126
+ bh for bh in block_hashes_list if self._policy.get(bh) is None
127
+ ]
128
+
129
+ if not block_hashes_to_store:
130
+ return PrepareStoreOutput(
131
+ block_hashes_to_store=[],
132
+ store_spec=self._get_load_store_spec([], []),
133
+ block_hashes_evicted=[],
134
+ )
135
+
136
+ num_blocks_to_evict = len(block_hashes_to_store) - self._get_num_free_blocks()
137
+
138
+ to_evict: list[BlockHash] = []
139
+ if num_blocks_to_evict > 0:
140
+ # Blocks from the original input are excluded from eviction candidates:
141
+ # a block that was already stored must remain in the cache after this call.
142
+ protected = set(block_hashes_list)
143
+ evicted = self._policy.evict(num_blocks_to_evict, protected)
144
+ if evicted is None:
145
+ return None
146
+ for block_hash, block in evicted:
147
+ self._free_block(block)
148
+ to_evict.append(block_hash)
149
+
150
+ if to_evict and self.events is not None:
151
+ self.events.append(
152
+ OffloadingEvent(
153
+ block_hashes=to_evict,
154
+ block_size=self.block_size,
155
+ medium=self.medium,
156
+ removed=True,
157
+ )
158
+ )
159
+
160
+ blocks = self._allocate_blocks(block_hashes_to_store)
161
+ assert len(blocks) == len(block_hashes_to_store), (
162
+ "Block pool did not allocate the expected number of blocks"
163
+ )
164
+
165
+ for block_hash, block in zip(block_hashes_to_store, blocks):
166
+ self._policy.insert(block_hash, block)
167
+
168
+ # build store specs for allocated blocks
169
+ store_spec = self._get_load_store_spec(block_hashes_to_store, blocks)
170
+
171
+ return PrepareStoreOutput(
172
+ block_hashes_to_store=block_hashes_to_store,
173
+ store_spec=store_spec,
174
+ block_hashes_evicted=to_evict,
175
+ )
176
+
177
+ def complete_store(
178
+ self, block_hashes: Iterable[BlockHash], success: bool = True
179
+ ) -> None:
180
+ stored_block_hashes: list[BlockHash] = []
181
+
182
+ if success:
183
+ for block_hash in block_hashes:
184
+ block = self._policy.get(block_hash)
185
+ if block is not None and not block.is_ready:
186
+ block.ref_cnt = 0
187
+ stored_block_hashes.append(block_hash)
188
+ else:
189
+ for block_hash in block_hashes:
190
+ block = self._policy.get(block_hash)
191
+ if block is not None and not block.is_ready:
192
+ self._policy.remove(block_hash)
193
+ self._free_block(block)
194
+
195
+ if stored_block_hashes and self.events is not None:
196
+ self.events.append(
197
+ OffloadingEvent(
198
+ block_hashes=stored_block_hashes,
199
+ block_size=self.block_size,
200
+ medium=self.medium,
201
+ removed=False,
202
+ )
203
+ )
204
+
205
+ def take_events(self) -> Iterable[OffloadingEvent]:
206
+ if self.events is not None:
207
+ yield from self.events
208
+ self.events.clear()
vllm/v1/kv_offload/cpu/policies/__init__.py
vllm/v1/kv_offload/cpu/policies/abstract.py
new
+76
@@ -0,0 +1,76 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import ctypes
4
+from abc import ABC, abstractmethod
5
+from collections.abc import Iterable
6
+
7
+from vllm.v1.core.kv_cache_utils import BlockHash
8
+
9
+
10
+class BlockStatus(ctypes.Structure):
11
+ """
12
+ Offloading status for a single block of KV data.
13
+ Holds the following information:
14
+
15
+ ref_cnt - the current number of transfers using this block as a source.
16
+ A value of -1 indicates the block is not yet ready to be read.
17
+ block_id - index of the physical CPU buffer slot.
18
+ """
19
+
20
+ _fields_ = [("ref_cnt", ctypes.c_int32), ("block_id", ctypes.c_int64)]
21
+
22
+ def __init__(self, block_id: int):
23
+ super().__init__()
24
+ # initialize block as "not ready" (ref_cnt = -1)
25
+ self.ref_cnt = -1
26
+ self.block_id = block_id
27
+
28
+ @property
29
+ def is_ready(self) -> bool:
30
+ """
31
+ Returns whether the block is ready to be read.
32
+ """
33
+ return self.ref_cnt >= 0
34
+
35
+
36
+class CachePolicy(ABC):
37
+ """
38
+ Encapsulates both block organization (data structures) and replacement
39
+ decisions (which block to evict). LRU and ARC differ in both dimensions —
40
+ ARC's ghost lists and target_t1_size live at the intersection of storage
41
+ and eviction, so they cannot be separated cleanly.
42
+ """
43
+
44
+ @abstractmethod
45
+ def __init__(self, cache_capacity: int) -> None: ...
46
+
47
+ @abstractmethod
48
+ def get(self, block_hash: BlockHash) -> BlockStatus | None:
49
+ """Find block in data structures. Returns None if not present."""
50
+
51
+ @abstractmethod
52
+ def insert(self, block_hash: BlockHash, block: BlockStatus) -> None:
53
+ """Add a newly allocated block. For ARC: also removes from ghost lists."""
54
+
55
+ @abstractmethod
56
+ def remove(self, block_hash: BlockHash) -> None:
57
+ """Remove a block (used to clean up after a failed store)."""
58
+
59
+ @abstractmethod
60
+ def touch(self, block_hashes: Iterable[BlockHash]) -> None:
61
+ """Mark blocks as recently used."""
62
+
63
+ @abstractmethod
64
+ def evict(
65
+ self, n: int, protected: set[BlockHash]
66
+ ) -> list[tuple[BlockHash, BlockStatus]] | None:
67
+ """
68
+ Evict exactly n blocks, skipping any in protected.
69
+
70
+ Returns a list of (block_hash, block) for the evicted blocks,
71
+ or None if n evictions cannot be satisfied. The operation is atomic:
72
+ if None is returned, no state changes are made.
73
+
74
+ For ARC: ghost list cleanup (trimming to cache_capacity) is performed
75
+ at the end of a successful eviction.
76
+ """
vllm/v1/kv_offload/cpu/policies/arc.py
new
+156
@@ -0,0 +1,156 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from collections import OrderedDict
4
+from collections.abc import Iterable
5
+
6
+from vllm.v1.core.kv_cache_utils import BlockHash
7
+from vllm.v1.kv_offload.cpu.policies.abstract import BlockStatus, CachePolicy
8
+
9
+
10
+class ARCCachePolicy(CachePolicy):
11
+ """
12
+ ARC (Adaptive Replacement Cache) cache policy.
13
+
14
+ Data Structures:
15
+ T1: Recent cache containing blocks accessed once.
16
+ T2: Frequent cache containing blocks accessed multiple times.
17
+ B1/B2: Ghost lists tracking recently evicted blocks from T1/T2.
18
+ target_t1_size: Adaptive target size for the T1 partition.
19
+
20
+ Algorithm Flow:
21
+ 1. Cache lookup (lookup):
22
+ Searches T1 and T2 for block hashes and counts consecutive hits
23
+ until a miss or non-ready block is encountered.
24
+
25
+ 2. Cache touch (touch) - Adaptive Learning:
26
+ For each block_hash (in reverse order):
27
+ - If in T1: Move to T2 (promotion from recent to frequent).
28
+ - If in T2: Move to MRU position (end of queue).
29
+ - If in B1 ghost list: Increase target_t1_size.
30
+ - If in B2 ghost list: Decrease target_t1_size.
31
+
32
+ 3. Block eviction (evict) - Adaptive Replacement:
33
+ Determines eviction source based on adaptive target:
34
+ - If T1 size >= target_t1_size: Evict from T1, add to B1.
35
+ - Otherwise: Evict from T2, add to B2.
36
+ Finally, bound each ghost list size.
37
+
38
+ 4. Block insertion (insert):
39
+ New blocks are always inserted into T1 and removed from B1/B2 if
40
+ present. Blocks may later be promoted to T2 during touch operations.
41
+
42
+ Adaptive Behavior:
43
+ The algorithm self-tunes the recency vs. frequency trade-off:
44
+ - B1 hit: Recent access patterns matter more → increase T1.
45
+ - B2 hit: Frequent access patterns matter more → decrease T1.
46
+ """
47
+
48
+ def __init__(self, cache_capacity: int):
49
+ self.cache_capacity: int = cache_capacity
50
+ self.target_t1_size: float = 0.0
51
+ self.t1: OrderedDict[BlockHash, BlockStatus] = OrderedDict()
52
+ self.t2: OrderedDict[BlockHash, BlockStatus] = OrderedDict()
53
+ # block_hash -> None (only care about presence)
54
+ self.b1: OrderedDict[BlockHash, None] = OrderedDict()
55
+ self.b2: OrderedDict[BlockHash, None] = OrderedDict()
56
+
57
+ def get(self, block_hash: BlockHash) -> BlockStatus | None:
58
+ return self.t1.get(block_hash) or self.t2.get(block_hash)
59
+
60
+ def insert(self, block_hash: BlockHash, block: BlockStatus) -> None:
61
+ self.t1[block_hash] = block
62
+ self.b1.pop(block_hash, None)
63
+ self.b2.pop(block_hash, None)
64
+
65
+ def remove(self, block_hash: BlockHash) -> None:
66
+ if self.t1.pop(block_hash, None) is None:
67
+ self.t2.pop(block_hash, None)
68
+
69
+ def touch(self, block_hashes: Iterable[BlockHash]) -> None:
70
+ for block_hash in reversed(list(block_hashes)):
71
+ if block_hash in self.t1:
72
+ block = self.t1.pop(block_hash)
73
+ if not block.is_ready:
74
+ # block was just prepared to be stored, not really touched
75
+ # twice — keep it in T1 and mark as most recently used
76
+ self.t1[block_hash] = block
77
+ else:
78
+ self.t2[block_hash] = block
79
+
80
+ elif block_hash in self.t2:
81
+ self.t2.move_to_end(block_hash)
82
+
83
+ elif block_hash in self.b1:
84
+ delta = max(1, len(self.b2) / len(self.b1))
85
+ self.target_t1_size = min(
86
+ self.target_t1_size + delta, self.cache_capacity
87
+ )
88
+ # move to MRU position (end) to keep it fresh in the ghost list
89
+ self.b1.move_to_end(block_hash)
90
+
91
+ elif block_hash in self.b2:
92
+ delta = max(1, len(self.b1) / len(self.b2))
93
+ self.target_t1_size = max(self.target_t1_size - delta, 0)
94
+ # move to MRU position (end) to keep it fresh in the ghost list
95
+ self.b2.move_to_end(block_hash)
96
+
97
+ def evict(
98
+ self, n: int, protected: set[BlockHash]
99
+ ) -> list[tuple[BlockHash, BlockStatus]] | None:
100
+ if n == 0:
101
+ return []
102
+
103
+ # Collect candidates atomically: simulate T1 size changes as we select,
104
+ # but do not modify actual data structures until all n are found.
105
+ candidates: list[
106
+ tuple[BlockHash, BlockStatus, bool]
107
+ ] = [] # (hash, block, from_t1)
108
+ already_selected: set[BlockHash] = set()
109
+ virtual_t1_size = len(self.t1)
110
+
111
+ for _ in range(n):
112
+ candidate: tuple[BlockHash, BlockStatus, bool] | None = None
113
+
114
+ if virtual_t1_size >= int(self.target_t1_size):
115
+ for block_hash, block in self.t1.items():
116
+ if (
117
+ block.ref_cnt == 0
118
+ and block_hash not in protected
119
+ and block_hash not in already_selected
120
+ ):
121
+ candidate = (block_hash, block, True)
122
+ virtual_t1_size -= 1
123
+ break
124
+
125
+ if candidate is None:
126
+ for block_hash, block in self.t2.items():
127
+ if (
128
+ block.ref_cnt == 0
129
+ and block_hash not in protected
130
+ and block_hash not in already_selected
131
+ ):
132
+ candidate = (block_hash, block, False)
133
+ break
134
+ if candidate is None:
135
+ return None
136
+
137
+ candidates.append(candidate)
138
+ already_selected.add(candidate[0])
139
+
140
+ # Apply all evictions now that we know n candidates exist.
141
+ result: list[tuple[BlockHash, BlockStatus]] = []
142
+ for block_hash, block, from_t1 in candidates:
143
+ if from_t1:
144
+ del self.t1[block_hash]
145
+ self.b1[block_hash] = None
146
+ else:
147
+ del self.t2[block_hash]
148
+ self.b2[block_hash] = None
149
+ result.append((block_hash, block))
150
+
151
+ # Trim ghost lists to cache_capacity.
152
+ for ghost in (self.b1, self.b2):
153
+ for _ in range(len(ghost) - self.cache_capacity):
154
+ ghost.popitem(last=False)
155
+
156
+ return result
vllm/v1/kv_offload/cpu/policies/lru.py
new
+46
@@ -0,0 +1,46 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from collections import OrderedDict
4
+from collections.abc import Iterable
5
+
6
+from vllm.v1.core.kv_cache_utils import BlockHash
7
+from vllm.v1.kv_offload.cpu.policies.abstract import BlockStatus, CachePolicy
8
+
9
+
10
+class LRUCachePolicy(CachePolicy):
11
+ """LRU cache policy backed by a single OrderedDict."""
12
+
13
+ def __init__(self, cache_capacity: int):
14
+ # cache_capacity unused by LRU but accepted for a uniform constructor
15
+ self.blocks: OrderedDict[BlockHash, BlockStatus] = OrderedDict()
16
+
17
+ def get(self, block_hash: BlockHash) -> BlockStatus | None:
18
+ return self.blocks.get(block_hash)
19
+
20
+ def insert(self, block_hash: BlockHash, block: BlockStatus) -> None:
21
+ self.blocks[block_hash] = block
22
+
23
+ def remove(self, block_hash: BlockHash) -> None:
24
+ del self.blocks[block_hash]
25
+
26
+ def touch(self, block_hashes: Iterable[BlockHash]) -> None:
27
+ for block_hash in reversed(list(block_hashes)):
28
+ if block_hash in self.blocks:
29
+ self.blocks.move_to_end(block_hash)
30
+
31
+ def evict(
32
+ self, n: int, protected: set[BlockHash]
33
+ ) -> list[tuple[BlockHash, BlockStatus]] | None:
34
+ if n == 0:
35
+ return []
36
+ candidates: list[tuple[BlockHash, BlockStatus]] = []
37
+ for block_hash, block in self.blocks.items():
38
+ if block.ref_cnt == 0 and block_hash not in protected:
39
+ candidates.append((block_hash, block))
40
+ if len(candidates) == n:
41
+ break
42
+ if len(candidates) < n:
43
+ return None
44
+ for block_hash, _ in candidates:
45
+ del self.blocks[block_hash]
46
+ return candidates
vllm/v1/kv_offload/cpu/spec.py
renamed
+7
-19
@@ -9,9 +9,7 @@ from vllm.platforms import current_platform
9
from vllm.v1.attention.backend import AttentionBackend
10
from vllm.v1.kv_cache_interface import KVCacheConfig
11
from vllm.v1.kv_offload.abstract import LoadStoreSpec, OffloadingManager
12
-from vllm.v1.kv_offload.arc_manager import ARCOffloadingManager
13
-from vllm.v1.kv_offload.backends.cpu import CPUBackend
14
-from vllm.v1.kv_offload.lru_manager import LRUOffloadingManager
12
+from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager
13
from vllm.v1.kv_offload.mediums import CPULoadStoreSpec, GPULoadStoreSpec
14
from vllm.v1.kv_offload.reuse_manager import FilterReusedOffloadingManager
15
from vllm.v1.kv_offload.spec import OffloadingSpec
@@ -68,23 +66,13 @@ class CPUOffloadingSpec(OffloadingSpec):
66
assert len(self.gpu_block_size) == 1
67
gpu_block_size = self.gpu_block_size[0]
68
offloaded_block_size = gpu_block_size * self.block_size_factor
71
- backend = CPUBackend(
72
- block_size=offloaded_block_size, num_blocks=self.num_blocks
73
- )
69
75
- if self.eviction_policy == "lru":
76
- self._manager = LRUOffloadingManager(
77
- backend=backend, enable_events=enable_events
78
- )
79
- elif self.eviction_policy == "arc":
80
- self._manager = ARCOffloadingManager(
81
- backend=backend, enable_events=enable_events
82
- )
83
- else:
84
- raise ValueError(
85
- f"Unknown eviction policy: {self.eviction_policy}. "
86
- f"Supported policies: lru, arc"
87
- )
70
+ self._manager = CPUOffloadingManager(
71
+ block_size=offloaded_block_size,
72
+ num_blocks=self.num_blocks,
73
+ cache_policy=self.eviction_policy, # type: ignore[arg-type]
74
+ enable_events=enable_events,
75
+ )
76
77
# store_threshold: how many times a block must appear in lookup()
78
# before it is eligible for CPU offloading. Values < 2 disable
vllm/v1/kv_offload/factory.py
+1
-1
@@ -54,5 +54,5 @@ class OffloadingSpecFactory:
54
55
# Register various specs here.
56
OffloadingSpecFactory.register_spec(
57
- "CPUOffloadingSpec", "vllm.v1.kv_offload.cpu", "CPUOffloadingSpec"
57
+ "CPUOffloadingSpec", "vllm.v1.kv_offload.cpu.spec", "CPUOffloadingSpec"
58
)
vllm/v1/kv_offload/lru_manager.py
deleted
-146
@@ -1,146 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-from collections import OrderedDict
4
-from collections.abc import Iterable
5
-
6
-from vllm.v1.core.kv_cache_utils import BlockHash
7
-from vllm.v1.kv_offload.abstract import (
8
- LoadStoreSpec,
9
- OffloadingEvent,
10
- OffloadingManager,
11
- PrepareStoreOutput,
12
-)
13
-from vllm.v1.kv_offload.backend import Backend, BlockStatus
14
-
15
-
16
-class LRUOffloadingManager(OffloadingManager):
17
- """
18
- An OffloadingManager with a pluggable backend, which evicts blocks by LRU.
19
- """
20
-
21
- def __init__(self, backend: Backend, enable_events: bool = False):
22
- self.backend: Backend = backend
23
- # block_hash -> BlockStatus
24
- self.blocks: OrderedDict[BlockHash, BlockStatus] = OrderedDict()
25
- self.events: list[OffloadingEvent] | None = [] if enable_events else None
26
-
27
- def lookup(self, block_hashes: Iterable[BlockHash]) -> int | None:
28
- hit_count = 0
29
- for block_hash in block_hashes:
30
- block = self.blocks.get(block_hash)
31
- if block is None or not block.is_ready:
32
- break
33
- hit_count += 1
34
- return hit_count
35
-
36
- def prepare_load(self, block_hashes: Iterable[BlockHash]) -> LoadStoreSpec:
37
- blocks = []
38
- for block_hash in block_hashes:
39
- block = self.blocks[block_hash]
40
- assert block.is_ready
41
- block.ref_cnt += 1
42
- blocks.append(block)
43
-
44
- return self.backend.get_load_store_spec(block_hashes, blocks)
45
-
46
- def touch(self, block_hashes: Iterable[BlockHash]):
47
- for block_hash in reversed(list(block_hashes)):
48
- if self.blocks.get(block_hash):
49
- self.blocks.move_to_end(block_hash)
50
-
51
- def complete_load(self, block_hashes: Iterable[BlockHash]):
52
- for block_hash in block_hashes:
53
- block = self.blocks[block_hash]
54
- assert block.ref_cnt > 0
55
- block.ref_cnt -= 1
56
-
57
- def prepare_store(
58
- self, block_hashes: Iterable[BlockHash]
59
- ) -> PrepareStoreOutput | None:
60
- block_hashes_list = list(block_hashes)
61
-
62
- # filter out blocks that are already stored
63
- block_hashes_to_store = [
64
- block_hash
65
- for block_hash in block_hashes_list
66
- if block_hash not in self.blocks
67
- ]
68
-
69
- num_blocks_to_evict = (
70
- len(block_hashes_to_store) - self.backend.get_num_free_blocks()
71
- )
72
-
73
- # build list of blocks to evict
74
- to_evict = []
75
- if num_blocks_to_evict > 0:
76
- # Blocks from the original input are excluded from eviction candidates:
77
- # a block that was already stored must remain in the cache after this call.
78
- protected = set(block_hashes_list)
79
- for block_hash, block in self.blocks.items():
80
- if block.ref_cnt == 0 and block_hash not in protected:
81
- to_evict.append(block_hash)
82
- num_blocks_to_evict -= 1
83
- if num_blocks_to_evict == 0:
84
- break
85
- else:
86
- # we could not evict enough blocks
87
- return None
88
-
89
- # evict blocks
90
- for block_hash in to_evict:
91
- self.backend.free(self.blocks.pop(block_hash))
92
-
93
- if to_evict and self.events is not None:
94
- self.events.append(
95
- OffloadingEvent(
96
- block_hashes=to_evict,
97
- block_size=self.backend.block_size,
98
- medium=self.backend.medium,
99
- removed=True,
100
- )
101
- )
102
-
103
- blocks = self.backend.allocate_blocks(block_hashes_to_store)
104
- assert len(blocks) == len(block_hashes_to_store)
105
-
106
- for block_hash, block in zip(block_hashes_to_store, blocks):
107
- self.blocks[block_hash] = block
108
-
109
- # build store specs for allocated blocks
110
- store_spec = self.backend.get_load_store_spec(block_hashes_to_store, blocks)
111
-
112
- return PrepareStoreOutput(
113
- block_hashes_to_store=block_hashes_to_store,
114
- store_spec=store_spec,
115
- block_hashes_evicted=to_evict,
116
- )
117
-
118
- def complete_store(self, block_hashes: Iterable[BlockHash], success: bool = True):
119
- stored_block_hashes: list[BlockHash] = []
120
- if success:
121
- for block_hash in block_hashes:
122
- block = self.blocks[block_hash]
123
- if not block.is_ready:
124
- block.ref_cnt = 0
125
- stored_block_hashes.append(block_hash)
126
- else:
127
- for block_hash in block_hashes:
128
- block = self.blocks[block_hash]
129
- if not block.is_ready:
130
- self.backend.free(block)
131
- del self.blocks[block_hash]
132
-
133
- if stored_block_hashes and self.events is not None:
134
- self.events.append(
135
- OffloadingEvent(
136
- block_hashes=stored_block_hashes,
137
- block_size=self.backend.block_size,
138
- medium=self.backend.medium,
139
- removed=False,
140
- )
141
- )
142
-
143
- def take_events(self) -> Iterable[OffloadingEvent]:
144
- if self.events is not None:
145
- yield from self.events
146
- self.events.clear()