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 13, 09:30 UTC
Implement ARC KV cache eviction policy for CPU offloader (#27039)
Signed-off-by: Alberto Perdomo <aperdomo@redhat.com> Signed-off-by: alberto <aperdomo@redhat.com> Co-authored-by: Or Ozeri <or@ozery.com>
alberto committed
Nov 12, 2025 at 17:51 UTC
bac904565f170ba198c2398a0f627b38f9cb8e18
3 files changed
+565
-5
tests/v1/kv_offload/test_cpu_manager.py
+308
@@ -11,6 +11,7 @@ from vllm.v1.kv_offload.abstract import (
11
OffloadingEvent,
12
PrepareStoreOutput,
13
)
14
+from vllm.v1.kv_offload.arc_manager import ARCOffloadingManager
15
from vllm.v1.kv_offload.backends.cpu import CPUBackend
16
from vllm.v1.kv_offload.lru_manager import LRUOffloadingManager
17
from vllm.v1.kv_offload.mediums import CPULoadStoreSpec
@@ -187,3 +188,310 @@ def test_cpu_manager():
188
expected_stores=({3, 4, 5}, {6, 7, 8}),
189
expected_evictions=({2, 3, 4}, {8}),
190
)
191
+
192
+
193
+def test_arc_manager_basic():
194
+ """
195
+ Tests ARCOffloadingManager basic operations with a CPUBackend.
196
+ Verifies that ARC handles store, load, and lookup operations correctly.
197
+ """
198
+ # initialize a CPU backend with a capacity of 4 blocks
199
+ block_size = 256
200
+ cpu_backend = CPUBackend(block_size=block_size, num_blocks=4)
201
+ arc_manager = ARCOffloadingManager(cpu_backend, enable_events=True)
202
+
203
+ # prepare store [1, 2]
204
+ prepare_store_output = arc_manager.prepare_store(to_hashes([1, 2]))
205
+ verify_store_output(
206
+ prepare_store_output,
207
+ ExpectedPrepareStoreOutput(
208
+ block_hashes_to_store=[1, 2],
209
+ store_block_ids=[0, 1],
210
+ block_hashes_evicted=[],
211
+ ),
212
+ )
213
+
214
+ # lookup [1, 2] -> not ready
215
+ assert arc_manager.lookup(to_hashes([1, 2])) == 0
216
+
217
+ # no events so far
218
+ assert list(arc_manager.take_events()) == []
219
+
220
+ # complete store [1, 2]
221
+ arc_manager.complete_store(to_hashes([1, 2]))
222
+ verify_events(
223
+ arc_manager.take_events(), block_size=block_size, expected_stores=({1, 2},)
224
+ )
225
+
226
+ # lookup [1, 2]
227
+ assert arc_manager.lookup(to_hashes([1])) == 1
228
+ assert arc_manager.lookup(to_hashes([1, 2])) == 2
229
+ assert arc_manager.lookup(to_hashes([1, 2, 3])) == 2
230
+
231
+ # blocks should be in T1 (recent)
232
+ assert len(arc_manager.t1) == 2
233
+ assert len(arc_manager.t2) == 0
234
+
235
+
236
+def test_arc_manager_t1_to_t2_promotion():
237
+ """
238
+ Tests that accessing a block in T1 promotes it to T2 (frequent).
239
+ This is a key feature of ARC's adaptive behavior.
240
+ """
241
+ block_size = 256
242
+ cpu_backend = CPUBackend(block_size=block_size, num_blocks=4)
243
+ arc_manager = ARCOffloadingManager(cpu_backend, enable_events=False)
244
+
245
+ # store and complete block 1
246
+ arc_manager.prepare_store(to_hashes([1]))
247
+ arc_manager.complete_store(to_hashes([1]))
248
+
249
+ # block 1 starts in T1 (recent)
250
+ assert to_hashes([1])[0] in arc_manager.t1
251
+ assert to_hashes([1])[0] not in arc_manager.t2
252
+
253
+ # touch block 1 (simulate second access)
254
+ arc_manager.touch(to_hashes([1]))
255
+
256
+ # block 1 should now be in T2 (frequent)
257
+ assert to_hashes([1])[0] not in arc_manager.t1
258
+ assert to_hashes([1])[0] in arc_manager.t2
259
+
260
+
261
+def test_arc_manager_eviction_with_load():
262
+ """
263
+ Tests ARC eviction behavior similar to LRU test.
264
+ Verifies that blocks being loaded (ref_cnt > 0) cannot be evicted.
265
+ """
266
+ block_size = 256
267
+ cpu_backend = CPUBackend(block_size=block_size, num_blocks=4)
268
+ arc_manager = ARCOffloadingManager(cpu_backend, enable_events=True)
269
+
270
+ # prepare and complete store [1, 2, 3, 4]
271
+ prepare_store_output = arc_manager.prepare_store(to_hashes([1, 2, 3, 4]))
272
+ verify_store_output(
273
+ prepare_store_output,
274
+ ExpectedPrepareStoreOutput(
275
+ block_hashes_to_store=[1, 2, 3, 4],
276
+ store_block_ids=[0, 1, 2, 3],
277
+ block_hashes_evicted=[],
278
+ ),
279
+ )
280
+ arc_manager.complete_store(to_hashes([1, 2, 3, 4]))
281
+
282
+ # prepare load [2, 3] (increases ref_cnt)
283
+ prepare_load_output = arc_manager.prepare_load(to_hashes([2, 3]))
284
+ verify_load_output(prepare_load_output, [1, 2])
285
+
286
+ # prepare store [5, 6, 7] with [2, 3] being loaded
287
+ # should fail because [2, 3] have ref_cnt > 0
288
+ assert arc_manager.prepare_store(to_hashes([5, 6, 7])) is None
289
+
290
+ # complete load [2, 3]
291
+ arc_manager.complete_load(to_hashes([2, 3]))
292
+
293
+ # now prepare store [5, 6, 7] should succeed
294
+ # ARC will evict blocks one at a time from T1 as needed
295
+ prepare_store_output = arc_manager.prepare_store(to_hashes([5, 6, 7]))
296
+ assert prepare_store_output is not None
297
+ # Should successfully evict enough blocks to make room (at least 1)
298
+ assert len(prepare_store_output.block_hashes_evicted) >= 1
299
+
300
+
301
+def test_arc_manager_adaptive_target():
302
+ """
303
+ Tests ARC's adaptive target adjustment via ghost lists.
304
+ When a block in B1 (ghost list) is accessed, target_t1_size increases.
305
+ When a block in B2 is accessed, target_t1_size decreases.
306
+ """
307
+ block_size = 256
308
+ cpu_backend = CPUBackend(block_size=block_size, num_blocks=2)
309
+ arc_manager = ARCOffloadingManager(cpu_backend, enable_events=False)
310
+
311
+ # store blocks 1, 2 (fills cache)
312
+ arc_manager.prepare_store(to_hashes([1, 2]))
313
+ arc_manager.complete_store(to_hashes([1, 2]))
314
+
315
+ initial_target = arc_manager.target_t1_size
316
+
317
+ # store block 3, evicting block 1 (moves to B1 ghost list)
318
+ arc_manager.prepare_store(to_hashes([3]))
319
+ arc_manager.complete_store(to_hashes([3]))
320
+
321
+ # block 1 should be in B1 (ghost list)
322
+ assert to_hashes([1])[0] in arc_manager.b1
323
+
324
+ # touch block 1 (cache miss, but in B1)
325
+ # this should increase target_t1_size (favor recency)
326
+ arc_manager.touch(to_hashes([1]))
327
+
328
+ # target should have increased
329
+ assert arc_manager.target_t1_size > initial_target
330
+
331
+
332
+def test_arc_manager_t1_t2_eviction_policy():
333
+ """
334
+ Tests that ARC evicts from T1 or T2 based on target_t1_size.
335
+ If |T1| >= target_t1_size, evict from T1, otherwise from T2.
336
+ """
337
+ block_size = 256
338
+ cpu_backend = CPUBackend(block_size=block_size, num_blocks=4)
339
+ arc_manager = ARCOffloadingManager(cpu_backend, enable_events=False)
340
+
341
+ # store blocks 1, 2, 3, 4
342
+ arc_manager.prepare_store(to_hashes([1, 2, 3, 4]))
343
+ arc_manager.complete_store(to_hashes([1, 2, 3, 4]))
344
+
345
+ # promote blocks 3, 4 to T2 by touching them
346
+ arc_manager.touch(to_hashes([3, 4]))
347
+
348
+ # now: T1 = {1, 2}, T2 = {3, 4}
349
+ assert len(arc_manager.t1) == 2
350
+ assert len(arc_manager.t2) == 2
351
+
352
+ # set target_t1_size to prefer evicting from T1
353
+ # (when |T1| >= target, evict from T1)
354
+ arc_manager.target_t1_size = 1
355
+
356
+ # store block 5, should evict from T1 (block 1, LRU in T1)
357
+ output = arc_manager.prepare_store(to_hashes([5]))
358
+ assert output is not None
359
+ assert to_hashes([1]) == output.block_hashes_evicted
360
+
361
+ arc_manager.complete_store(to_hashes([5]))
362
+
363
+ # block 1 should be in B1 (ghost list)
364
+ assert to_hashes([1])[0] in arc_manager.b1
365
+ # block 5 should be in T1
366
+ assert to_hashes([5])[0] in arc_manager.t1
367
+
368
+
369
+def test_arc_manager_ghost_list_bounds():
370
+ """
371
+ Tests that ghost lists (B1, B2) don't grow unbounded.
372
+ They should be capped at cache_capacity.
373
+ """
374
+ block_size = 256
375
+ cpu_backend = CPUBackend(block_size=block_size, num_blocks=2)
376
+ arc_manager = ARCOffloadingManager(cpu_backend, enable_events=False)
377
+
378
+ # fill cache with blocks 1, 2
379
+ arc_manager.prepare_store(to_hashes([1, 2]))
380
+ arc_manager.complete_store(to_hashes([1, 2]))
381
+
382
+ # store many blocks to fill ghost lists
383
+ for i in range(3, 20):
384
+ arc_manager.prepare_store(to_hashes([i]))
385
+ arc_manager.complete_store(to_hashes([i]))
386
+
387
+ # ghost lists should not exceed cache_capacity
388
+ assert len(arc_manager.b1) <= arc_manager.cache_capacity
389
+ assert len(arc_manager.b2) <= arc_manager.cache_capacity
390
+
391
+
392
+def test_arc_manager_touch_ordering():
393
+ """
394
+ Tests that touch() correctly updates access patterns.
395
+ Similar to LRU test but verifies T1/T2 ordering.
396
+ """
397
+ block_size = 256
398
+ cpu_backend = CPUBackend(block_size=block_size, num_blocks=4)
399
+ arc_manager = ARCOffloadingManager(cpu_backend, enable_events=True)
400
+
401
+ # store blocks 1, 2, 3, 4
402
+ arc_manager.prepare_store(to_hashes([1, 2, 3, 4]))
403
+ arc_manager.complete_store(to_hashes([1, 2, 3, 4]))
404
+
405
+ # promote 3, 4 to T2
406
+ arc_manager.touch(to_hashes([3, 4]))
407
+
408
+ # T1 = {1, 2}, T2 = {3, 4}
409
+ # touch [1, 3, 4] - should promote 1 to T2, and move 3,4 to end of T2
410
+ arc_manager.touch(to_hashes([1, 3, 4]))
411
+
412
+ # T1 = {2}, T2 = {1, 3, 4} (in that order, with 4 most recent)
413
+ assert len(arc_manager.t1) == 1
414
+ assert len(arc_manager.t2) == 3
415
+
416
+ # store block 5, should evict from T1 (block 2, only one in T1)
417
+ prepare_store_output = arc_manager.prepare_store(to_hashes([5]))
418
+ verify_store_output(
419
+ prepare_store_output,
420
+ ExpectedPrepareStoreOutput(
421
+ block_hashes_to_store=[5],
422
+ store_block_ids=[1], # reuses block 2's storage
423
+ block_hashes_evicted=[2],
424
+ ),
425
+ )
426
+
427
+
428
+def test_arc_manager_failed_store():
429
+ """
430
+ Tests that failed store operations clean up correctly.
431
+ Similar to LRU test but for ARC.
432
+ """
433
+ block_size = 256
434
+ cpu_backend = CPUBackend(block_size=block_size, num_blocks=4)
435
+ arc_manager = ARCOffloadingManager(cpu_backend, enable_events=True)
436
+
437
+ # store blocks 1, 2, 3, 4
438
+ arc_manager.prepare_store(to_hashes([1, 2, 3, 4]))
439
+ arc_manager.complete_store(to_hashes([1, 2, 3, 4]))
440
+
441
+ # prepare store block 5 (will evict block 1)
442
+ prepare_store_output = arc_manager.prepare_store(to_hashes([5]))
443
+ assert prepare_store_output is not None
444
+ assert len(prepare_store_output.block_hashes_evicted) == 1
445
+
446
+ # complete store with failure
447
+ arc_manager.complete_store(to_hashes([5]), success=False)
448
+
449
+ # block 5 should not be in cache
450
+ assert arc_manager.lookup(to_hashes([5])) == 0
451
+ # block 5 should not be in T1 or T2
452
+ assert to_hashes([5])[0] not in arc_manager.t1
453
+ assert to_hashes([5])[0] not in arc_manager.t2
454
+
455
+ # evicted block should still be gone (in B1 ghost list)
456
+ evicted_hash = prepare_store_output.block_hashes_evicted[0]
457
+ assert evicted_hash in arc_manager.b1
458
+
459
+
460
+def test_arc_manager_full_scenario():
461
+ """
462
+ Comprehensive test covering multiple ARC operations in sequence.
463
+ Similar to the full LRU test but adapted for ARC behavior.
464
+ """
465
+ block_size = 256
466
+ cpu_backend = CPUBackend(block_size=block_size, num_blocks=4)
467
+ arc_manager = ARCOffloadingManager(cpu_backend, enable_events=True)
468
+
469
+ # store [1, 2]
470
+ arc_manager.prepare_store(to_hashes([1, 2]))
471
+ arc_manager.complete_store(to_hashes([1, 2]))
472
+
473
+ # store [3, 4, 5] -> evicts [1]
474
+ prepare_store_output = arc_manager.prepare_store(to_hashes([3, 4, 5]))
475
+ assert prepare_store_output is not None
476
+ assert len(prepare_store_output.block_hashes_evicted) == 1
477
+ arc_manager.complete_store(to_hashes([3, 4, 5]))
478
+
479
+ # promote some blocks to T2
480
+ arc_manager.touch(to_hashes([2, 3]))
481
+
482
+ # T1 has {4, 5}, T2 has {2, 3}
483
+ assert len(arc_manager.t1) == 2
484
+ assert len(arc_manager.t2) == 2
485
+
486
+ # store [6] -> should evict from T1 (4 is oldest in T1)
487
+ prepare_store_output = arc_manager.prepare_store(to_hashes([6]))
488
+ assert prepare_store_output is not None
489
+ arc_manager.complete_store(to_hashes([6]))
490
+
491
+ # verify blocks 2, 3 (in T2) are still present
492
+ assert arc_manager.lookup(to_hashes([2])) == 1
493
+ assert arc_manager.lookup(to_hashes([3])) == 1
494
+
495
+ # verify events
496
+ events = list(arc_manager.take_events())
497
+ assert len(events) > 0 # should have store and eviction events
vllm/v1/kv_offload/arc_manager.py
new
+237
@@ -0,0 +1,237 @@
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:
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
+ self.t1.move_to_end(block_hash)
94
+ else:
95
+ self.t2[block_hash] = block
96
+
97
+ elif block_hash in self.t2:
98
+ self.t2.move_to_end(block_hash)
99
+
100
+ elif block_hash in self.b1:
101
+ delta = max(1, len(self.b2) / len(self.b1))
102
+ self.target_t1_size = min(
103
+ self.target_t1_size + delta, self.cache_capacity
104
+ )
105
+ # move to MRU position (end) to keep it fresh in the ghost list
106
+ self.b1.move_to_end(block_hash)
107
+
108
+ elif block_hash in self.b2:
109
+ delta = max(1, len(self.b1) / len(self.b2))
110
+ self.target_t1_size = max(self.target_t1_size - delta, 0)
111
+ # move to MRU position (end) to keep it fresh in the ghost list
112
+ self.b2.move_to_end(block_hash)
113
+
114
+ def complete_load(self, block_hashes: Iterable[BlockHash]):
115
+ for block_hash in block_hashes:
116
+ block = self.t1.get(block_hash) or self.t2.get(block_hash)
117
+ assert block is not None, f"Block {block_hash!r} not found"
118
+ assert block.ref_cnt > 0, f"Block {block_hash!r} ref_cnt is already 0"
119
+
120
+ block.ref_cnt -= 1
121
+
122
+ def prepare_store(
123
+ self, block_hashes: Iterable[BlockHash]
124
+ ) -> PrepareStoreOutput | None:
125
+ block_hashes_to_store = []
126
+ for block_hash in block_hashes:
127
+ if block_hash not in self.t1 and block_hash not in self.t2:
128
+ block_hashes_to_store.append(block_hash)
129
+
130
+ if not block_hashes_to_store:
131
+ return PrepareStoreOutput(
132
+ block_hashes_to_store=[],
133
+ store_spec=self.backend.get_load_store_spec([], []),
134
+ block_hashes_evicted=[],
135
+ )
136
+
137
+ num_blocks_to_evict = (
138
+ len(block_hashes_to_store) - self.backend.get_num_free_blocks()
139
+ )
140
+
141
+ to_evict = []
142
+ while num_blocks_to_evict > 0:
143
+ block_to_evict = None
144
+ if len(self.t1) >= int(self.target_t1_size):
145
+ # try to evict the least recently used (oldest) block from T1
146
+ for block_hash, block in self.t1.items():
147
+ if block.ref_cnt == 0:
148
+ block_to_evict = (block_hash, block)
149
+ eviction_t = self.t1
150
+ eviction_b = self.b1
151
+ break
152
+ if not block_to_evict:
153
+ # try to evict the least recently used (oldest) block from T2
154
+ for block_hash, block in self.t2.items():
155
+ if block.ref_cnt == 0:
156
+ block_to_evict = (block_hash, block)
157
+ eviction_t = self.t2
158
+ eviction_b = self.b2
159
+ break
160
+ else:
161
+ # cannot evict enough blocks, cache is full of in-use items
162
+ return None
163
+
164
+ block_hash, block = block_to_evict
165
+ del eviction_t[block_hash]
166
+ eviction_b[block_hash] = None
167
+ to_evict.append(block_hash)
168
+ self.backend.free(block)
169
+ num_blocks_to_evict -= 1
170
+
171
+ for b in [self.b1, self.b2]:
172
+ for i in range(len(b) - self.cache_capacity):
173
+ b.popitem(last=False)
174
+
175
+ if to_evict and self.events is not None:
176
+ self.events.append(
177
+ OffloadingEvent(
178
+ block_hashes=to_evict,
179
+ block_size=self.backend.block_size,
180
+ medium=self.backend.medium,
181
+ removed=True,
182
+ )
183
+ )
184
+
185
+ blocks = self.backend.allocate_blocks(block_hashes_to_store)
186
+ assert len(blocks) == len(block_hashes_to_store), (
187
+ "Backend did not allocate the expected number of blocks"
188
+ )
189
+
190
+ for block_hash, block in zip(block_hashes_to_store, blocks):
191
+ self.t1[block_hash] = block
192
+
193
+ self.b1.pop(block_hash, None)
194
+ self.b2.pop(block_hash, None)
195
+
196
+ store_spec = self.backend.get_load_store_spec(block_hashes_to_store, blocks)
197
+
198
+ return PrepareStoreOutput(
199
+ block_hashes_to_store=block_hashes_to_store,
200
+ store_spec=store_spec,
201
+ block_hashes_evicted=to_evict,
202
+ )
203
+
204
+ def complete_store(self, block_hashes: Iterable[BlockHash], success: bool = True):
205
+ stored_block_hashes: list[BlockHash] = []
206
+
207
+ if success:
208
+ for block_hash in block_hashes:
209
+ block = self.t1.get(block_hash) or self.t2.get(block_hash)
210
+
211
+ if block is not None and not block.is_ready:
212
+ block.ref_cnt = 0
213
+ stored_block_hashes.append(block_hash)
214
+ else:
215
+ for block_hash in block_hashes:
216
+ block = self.t1.pop(block_hash, None)
217
+
218
+ if block is None:
219
+ block = self.t2.pop(block_hash, None)
220
+
221
+ if block is not None and not block.is_ready:
222
+ self.backend.free(block)
223
+
224
+ if stored_block_hashes and self.events is not None:
225
+ self.events.append(
226
+ OffloadingEvent(
227
+ block_hashes=stored_block_hashes,
228
+ block_size=self.backend.block_size,
229
+ medium=self.backend.medium,
230
+ removed=False,
231
+ )
232
+ )
233
+
234
+ def take_events(self) -> Iterable[OffloadingEvent]:
235
+ if self.events is not None:
236
+ yield from self.events
237
+ self.events.clear()
vllm/v1/kv_offload/cpu.py
+20
-5
@@ -8,6 +8,7 @@ from vllm.config import VllmConfig, get_layers_from_vllm_config
8
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
9
from vllm.platforms import current_platform
10
from vllm.v1.kv_offload.abstract import LoadStoreSpec, OffloadingManager
11
+from vllm.v1.kv_offload.arc_manager import ARCOffloadingManager
12
from vllm.v1.kv_offload.backends.cpu import CPUBackend
13
from vllm.v1.kv_offload.lru_manager import LRUOffloadingManager
14
from vllm.v1.kv_offload.mediums import CPULoadStoreSpec, GPULoadStoreSpec
@@ -33,18 +34,32 @@ class CPUOffloadingSpec(OffloadingSpec):
34
# worker-side
35
self._handler: OffloadingHandler | None = None
36
37
+ self.eviction_policy: str = self.extra_config.get("eviction_policy", "lru")
38
+
39
def get_manager(self) -> OffloadingManager:
40
if not self._manager:
41
kv_events_config = self.vllm_config.kv_events_config
42
enable_events = (
43
kv_events_config is not None and kv_events_config.enable_kv_cache_events
44
)
42
- self._manager = LRUOffloadingManager(
43
- CPUBackend(
44
- block_size=self.offloaded_block_size, num_blocks=self.num_cpu_blocks
45
- ),
46
- enable_events=enable_events,
45
+
46
+ backend = CPUBackend(
47
+ block_size=self.offloaded_block_size, num_blocks=self.num_cpu_blocks
48
)
49
+
50
+ if self.eviction_policy == "lru":
51
+ self._manager = LRUOffloadingManager(
52
+ backend=backend, enable_events=enable_events
53
+ )
54
+ elif self.eviction_policy == "arc":
55
+ self._manager = ARCOffloadingManager(
56
+ backend=backend, enable_events=enable_events
57
+ )
58
+ else:
59
+ raise ValueError(
60
+ f"Unknown eviction policy: {self.eviction_policy}. "
61
+ f"Supported policies: lru, arc"
62
+ )
63
return self._manager
64
65
def get_handlers(