A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 3, 22:45 UTC
[KV Offload] Pass ReqContext to touch(), complete_load(), and complete_store() (#41366)
Signed-off-by: Ronen Schaffer <ronen.schaffer@ibm.com>
Ronen Schaffer committed
May 10, 2026 at 15:09 UTC
e175192d33fdeefcd6689c25a1f7dbe1a3f043ef
5 files changed
+70
-50
tests/v1/kv_offload/cpu/test_manager.py
+32
-32
@@ -117,10 +117,10 @@ def test_already_stored_block_not_evicted_during_prepare_store(eviction_policy):
117
118
# store [1, 2] and complete
119
manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX)
120
- manager.complete_store(to_keys([1, 2]))
120
+ manager.complete_store(to_keys([1, 2]), _EMPTY_REQ_CTX)
121
122
# touch [1] to make block 2 the LRU candidate
123
- manager.touch(to_keys([1]))
123
+ manager.touch(to_keys([1]), _EMPTY_REQ_CTX)
124
125
# prepare_store([2, 3, 4, 5]):
126
# - block 2 is already stored -> filtered out of keys_to_store
@@ -137,7 +137,7 @@ def test_already_stored_block_not_evicted_during_prepare_store(eviction_policy):
137
)
138
139
# complete_store must not silently drop block 2
140
- manager.complete_store(to_keys([2, 3, 4, 5]))
140
+ manager.complete_store(to_keys([2, 3, 4, 5]), _EMPTY_REQ_CTX)
141
142
# block 2 must still be present in the cache
143
assert manager.lookup(to_key(2), _EMPTY_REQ_CTX) is True
@@ -171,7 +171,7 @@ def test_cpu_manager():
171
assert list(cpu_manager.take_events()) == []
172
173
# complete store [1, 2]
174
- cpu_manager.complete_store(to_keys([1, 2]))
174
+ cpu_manager.complete_store(to_keys([1, 2]), _EMPTY_REQ_CTX)
175
verify_events(cpu_manager.take_events(), expected_stores=({1, 2},))
176
177
# lookup [1, 2]
@@ -199,7 +199,7 @@ def test_cpu_manager():
199
assert cpu_manager.prepare_store(to_keys([1, 6]), _EMPTY_REQ_CTX) is None
200
201
# complete store [2, 3, 4, 5]
202
- cpu_manager.complete_store(to_keys([2, 3, 4, 5]))
202
+ cpu_manager.complete_store(to_keys([2, 3, 4, 5]), _EMPTY_REQ_CTX)
203
204
# lookup (now that we have [2, 3, 4, 5])
205
assert cpu_manager.lookup(to_key(1), _EMPTY_REQ_CTX) is False
@@ -217,7 +217,7 @@ def test_cpu_manager():
217
assert cpu_manager.prepare_store(to_keys([6, 7, 8]), _EMPTY_REQ_CTX) is None
218
219
# complete load [2, 3]
220
- cpu_manager.complete_load(to_keys([2, 3]))
220
+ cpu_manager.complete_load(to_keys([2, 3]), _EMPTY_REQ_CTX)
221
222
# prepare store [6, 7, 8] -> evicts [2, 3, 4] (oldest)
223
prepare_store_output = cpu_manager.prepare_store(to_keys([6, 7, 8]), _EMPTY_REQ_CTX)
@@ -231,10 +231,10 @@ def test_cpu_manager():
231
)
232
233
# complete store [6, 7, 8]
234
- cpu_manager.complete_store(to_keys([6, 7, 8]))
234
+ cpu_manager.complete_store(to_keys([6, 7, 8]), _EMPTY_REQ_CTX)
235
236
# touch [5, 6, 7] (move to end of LRU order)
237
- cpu_manager.touch(to_keys([5, 6, 7]))
237
+ cpu_manager.touch(to_keys([5, 6, 7]), _EMPTY_REQ_CTX)
238
239
# prepare store [7, 9] -> evicts [8] (oldest following previous touch)
240
prepare_store_output = cpu_manager.prepare_store(to_keys([9]), _EMPTY_REQ_CTX)
@@ -248,7 +248,7 @@ def test_cpu_manager():
248
)
249
250
# complete store [7, 9] with failure
251
- cpu_manager.complete_store(to_keys([7, 9]), success=False)
251
+ cpu_manager.complete_store(to_keys([7, 9]), _EMPTY_REQ_CTX, success=False)
252
253
# assert [7] is still stored, but [9] is not
254
assert cpu_manager.lookup(to_key(7), _EMPTY_REQ_CTX) is True
@@ -304,7 +304,7 @@ class TestARCPolicy:
304
assert list(cpu_manager.take_events()) == []
305
306
# complete store [1, 2]
307
- cpu_manager.complete_store(to_keys([1, 2]))
307
+ cpu_manager.complete_store(to_keys([1, 2]), _EMPTY_REQ_CTX)
308
verify_events(cpu_manager.take_events(), expected_stores=({1, 2},))
309
310
# lookup [1, 2]
@@ -325,14 +325,14 @@ class TestARCPolicy:
325
326
# store and complete block 1
327
cpu_manager.prepare_store(to_keys([1]), _EMPTY_REQ_CTX)
328
- cpu_manager.complete_store(to_keys([1]))
328
+ cpu_manager.complete_store(to_keys([1]), _EMPTY_REQ_CTX)
329
330
# block 1 starts in T1 (recent)
331
assert to_keys([1])[0] in arc_policy.t1
332
assert to_keys([1])[0] not in arc_policy.t2
333
334
# touch block 1 (simulate second access)
335
- cpu_manager.touch(to_keys([1]))
335
+ cpu_manager.touch(to_keys([1]), _EMPTY_REQ_CTX)
336
337
# block 1 should now be in T2 (frequent)
338
assert to_keys([1])[0] not in arc_policy.t1
@@ -357,7 +357,7 @@ class TestARCPolicy:
357
evicted_keys=[],
358
),
359
)
360
- cpu_manager.complete_store(to_keys([1, 2, 3, 4]))
360
+ cpu_manager.complete_store(to_keys([1, 2, 3, 4]), _EMPTY_REQ_CTX)
361
362
# prepare load [2, 3] (increases ref_cnt)
363
prepare_load_output = cpu_manager.prepare_load(to_keys([2, 3]), _EMPTY_REQ_CTX)
@@ -368,7 +368,7 @@ class TestARCPolicy:
368
assert cpu_manager.prepare_store(to_keys([5, 6, 7]), _EMPTY_REQ_CTX) is None
369
370
# complete load [2, 3]
371
- cpu_manager.complete_load(to_keys([2, 3]))
371
+ cpu_manager.complete_load(to_keys([2, 3]), _EMPTY_REQ_CTX)
372
373
# now prepare store [5, 6, 7] should succeed
374
# ARC will evict blocks one at a time from T1 as needed
@@ -389,20 +389,20 @@ class TestARCPolicy:
389
390
# store blocks 1, 2 (fills cache)
391
cpu_manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX)
392
- cpu_manager.complete_store(to_keys([1, 2]))
392
+ cpu_manager.complete_store(to_keys([1, 2]), _EMPTY_REQ_CTX)
393
394
initial_target = arc_policy.target_t1_size
395
396
# store block 3, evicting block 1 (moves to B1 ghost list)
397
cpu_manager.prepare_store(to_keys([3]), _EMPTY_REQ_CTX)
398
- cpu_manager.complete_store(to_keys([3]))
398
+ cpu_manager.complete_store(to_keys([3]), _EMPTY_REQ_CTX)
399
400
# block 1 should be in B1 (ghost list)
401
assert to_keys([1])[0] in arc_policy.b1
402
403
# touch block 1 (cache miss, but in B1)
404
# this should increase target_t1_size (favor recency)
405
- cpu_manager.touch(to_keys([1]))
405
+ cpu_manager.touch(to_keys([1]), _EMPTY_REQ_CTX)
406
407
# target should have increased
408
assert arc_policy.target_t1_size > initial_target
@@ -416,10 +416,10 @@ class TestARCPolicy:
416
417
# store blocks 1, 2, 3, 4
418
cpu_manager.prepare_store(to_keys([1, 2, 3, 4]), _EMPTY_REQ_CTX)
419
- cpu_manager.complete_store(to_keys([1, 2, 3, 4]))
419
+ cpu_manager.complete_store(to_keys([1, 2, 3, 4]), _EMPTY_REQ_CTX)
420
421
# promote blocks 3, 4 to T2 by touching them
422
- cpu_manager.touch(to_keys([3, 4]))
422
+ cpu_manager.touch(to_keys([3, 4]), _EMPTY_REQ_CTX)
423
424
# now: T1 = {1, 2}, T2 = {3, 4}
425
assert len(arc_policy.t1) == 2
@@ -434,7 +434,7 @@ class TestARCPolicy:
434
assert output is not None
435
assert to_keys([1]) == output.evicted_keys
436
437
- cpu_manager.complete_store(to_keys([5]))
437
+ cpu_manager.complete_store(to_keys([5]), _EMPTY_REQ_CTX)
438
439
# block 1 should be in B1 (ghost list)
440
assert to_keys([1])[0] in arc_policy.b1
@@ -450,12 +450,12 @@ class TestARCPolicy:
450
451
# fill cache with blocks 1, 2
452
cpu_manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX)
453
- cpu_manager.complete_store(to_keys([1, 2]))
453
+ cpu_manager.complete_store(to_keys([1, 2]), _EMPTY_REQ_CTX)
454
455
# store many blocks to fill ghost lists
456
for i in range(3, 20):
457
cpu_manager.prepare_store(to_keys([i]), _EMPTY_REQ_CTX)
458
- cpu_manager.complete_store(to_keys([i]))
458
+ cpu_manager.complete_store(to_keys([i]), _EMPTY_REQ_CTX)
459
460
# ghost lists should not exceed cache_capacity
461
assert len(arc_policy.b1) <= arc_policy.cache_capacity
@@ -470,14 +470,14 @@ class TestARCPolicy:
470
471
# store blocks 1, 2, 3, 4
472
cpu_manager.prepare_store(to_keys([1, 2, 3, 4]), _EMPTY_REQ_CTX)
473
- cpu_manager.complete_store(to_keys([1, 2, 3, 4]))
473
+ cpu_manager.complete_store(to_keys([1, 2, 3, 4]), _EMPTY_REQ_CTX)
474
475
# promote 3, 4 to T2
476
- cpu_manager.touch(to_keys([3, 4]))
476
+ cpu_manager.touch(to_keys([3, 4]), _EMPTY_REQ_CTX)
477
478
# T1 = {1, 2}, T2 = {3, 4}
479
# touch [1, 3, 4] - should promote 1 to T2, and move 3,4 to end of T2
480
- cpu_manager.touch(to_keys([1, 3, 4]))
480
+ cpu_manager.touch(to_keys([1, 3, 4]), _EMPTY_REQ_CTX)
481
482
# T1 = {2}, T2 = {1, 3, 4} (in that order, with 4 most recent)
483
assert len(arc_policy.t1) == 1
@@ -503,7 +503,7 @@ class TestARCPolicy:
503
504
# store blocks 1, 2, 3, 4
505
cpu_manager.prepare_store(to_keys([1, 2, 3, 4]), _EMPTY_REQ_CTX)
506
- cpu_manager.complete_store(to_keys([1, 2, 3, 4]))
506
+ cpu_manager.complete_store(to_keys([1, 2, 3, 4]), _EMPTY_REQ_CTX)
507
508
# prepare store block 5 (will evict block 1)
509
prepare_store_output = cpu_manager.prepare_store(to_keys([5]), _EMPTY_REQ_CTX)
@@ -511,7 +511,7 @@ class TestARCPolicy:
511
assert len(prepare_store_output.evicted_keys) == 1
512
513
# complete store with failure
514
- cpu_manager.complete_store(to_keys([5]), success=False)
514
+ cpu_manager.complete_store(to_keys([5]), _EMPTY_REQ_CTX, success=False)
515
516
# block 5 should not be in cache
517
assert cpu_manager.lookup(to_key(5), _EMPTY_REQ_CTX) is False
@@ -532,7 +532,7 @@ class TestARCPolicy:
532
533
# store [1, 2]
534
cpu_manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX)
535
- cpu_manager.complete_store(to_keys([1, 2]))
535
+ cpu_manager.complete_store(to_keys([1, 2]), _EMPTY_REQ_CTX)
536
537
# store [3, 4, 5] -> evicts [1]
538
prepare_store_output = cpu_manager.prepare_store(
@@ -540,10 +540,10 @@ class TestARCPolicy:
540
)
541
assert prepare_store_output is not None
542
assert len(prepare_store_output.evicted_keys) == 1
543
- cpu_manager.complete_store(to_keys([3, 4, 5]))
543
+ cpu_manager.complete_store(to_keys([3, 4, 5]), _EMPTY_REQ_CTX)
544
545
# promote some blocks to T2
546
- cpu_manager.touch(to_keys([2, 3]))
546
+ cpu_manager.touch(to_keys([2, 3]), _EMPTY_REQ_CTX)
547
548
# T1 has {4, 5}, T2 has {2, 3}
549
assert len(arc_policy.t1) == 2
@@ -552,7 +552,7 @@ class TestARCPolicy:
552
# store [6] -> should evict from T1 (4 is oldest in T1)
553
prepare_store_output = cpu_manager.prepare_store(to_keys([6]), _EMPTY_REQ_CTX)
554
assert prepare_store_output is not None
555
- cpu_manager.complete_store(to_keys([6]))
555
+ cpu_manager.complete_store(to_keys([6]), _EMPTY_REQ_CTX)
556
557
# verify blocks 2, 3 (in T2) are still present
558
assert cpu_manager.lookup(to_key(2), _EMPTY_REQ_CTX) is True
@@ -609,4 +609,4 @@ def test_filter_reused_manager():
609
assert prepare_store_output is not None
610
assert prepare_store_output.keys_to_store == []
611
612
- manager.complete_store(to_keys([1]))
612
+ manager.complete_store(to_keys([1]), _EMPTY_REQ_CTX)
vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py
+8
-6
@@ -291,7 +291,7 @@ class OffloadingConnectorScheduler:
291
self.config.kv_group_configs, req_status.group_states
292
):
293
if group_config.sliding_window_size_in_blocks is None:
294
- self.manager.touch(group_state.offload_keys)
294
+ self.manager.touch(group_state.offload_keys, req_status.req_context)
295
else:
296
# we aim to keep just blocks that are necessary to hit
297
# the original request (+ decoded blocks)
@@ -300,7 +300,10 @@ class OffloadingConnectorScheduler:
300
group_state.num_hit_blocks
301
- group_config.sliding_window_size_in_blocks,
302
)
303
- self.manager.touch(group_state.offload_keys[blocks_to_skip:])
303
+ self.manager.touch(
304
+ group_state.offload_keys[blocks_to_skip:],
305
+ req_status.req_context,
306
+ )
307
308
def _lookup(self, req_status: RequestOffloadState) -> int | None:
309
"""
@@ -802,14 +805,13 @@ class OffloadingConnectorScheduler:
805
continue
806
assert job_status.pending_count == 0
807
808
+ req_status = self._req_status[job_status.req_id]
809
if job_status.is_store:
806
- self.manager.complete_store(job_status.keys)
810
+ self.manager.complete_store(job_status.keys, req_status.req_context)
811
else:
808
- self.manager.complete_load(job_status.keys)
812
+ self.manager.complete_load(job_status.keys, req_status.req_context)
813
if self._blocks_being_loaded:
814
self._blocks_being_loaded.difference_update(job_status.keys)
811
-
812
- req_status = self._req_status[job_status.req_id]
815
if self._block_id_to_pending_jobs:
816
# Sliding window blocks are tracked from store creation
817
# and must be cleaned up unconditionally.
vllm/v1/kv_offload/base.py
+11
-3
@@ -147,22 +147,24 @@ class OffloadingManager(ABC):
147
"""
148
pass
149
150
- def touch(self, keys: Collection[OffloadKey]):
150
+ def touch(self, keys: Collection[OffloadKey], req_context: ReqContext):
151
"""
152
Mark the given blocks as recently used.
153
This could in practice mean moving them to the end of an LRU list.
154
155
Args:
156
keys: the keys identifying the blocks.
157
+ req_context: per-request context (e.g. kv_transfer_params).
158
"""
159
return
160
160
- def complete_load(self, keys: Collection[OffloadKey]):
161
+ def complete_load(self, keys: Collection[OffloadKey], req_context: ReqContext):
162
"""
163
Marks previous blocks that were prepared to load as done loading.
164
165
Args:
166
keys: the keys identifying the blocks.
167
+ req_context: per-request context (e.g. kv_transfer_params).
168
"""
169
return
170
@@ -189,7 +191,12 @@ class OffloadingManager(ABC):
191
"""
192
pass
193
192
- def complete_store(self, keys: Collection[OffloadKey], success: bool = True):
194
+ def complete_store(
195
+ self,
196
+ keys: Collection[OffloadKey],
197
+ req_context: ReqContext,
198
+ success: bool = True,
199
+ ):
200
"""
201
Marks blocks which were previously prepared to be stored, as stored.
202
Following this call, the blocks become loadable.
@@ -198,6 +205,7 @@ class OffloadingManager(ABC):
205
206
Args:
207
keys: the keys identifying the blocks.
208
+ req_context: per-request context (e.g. kv_transfer_params).
209
success: whether the blocks were stored successfully.
210
"""
211
return
vllm/v1/kv_offload/cpu/manager.py
+8
-3
@@ -106,10 +106,12 @@ class CPUOffloadingManager(OffloadingManager):
106
blocks.append(block)
107
return self._get_load_store_spec(keys, blocks)
108
109
- def touch(self, keys: Collection[OffloadKey]) -> None:
109
+ def touch(self, keys: Collection[OffloadKey], req_context: ReqContext) -> None:
110
self._policy.touch(keys)
111
112
- def complete_load(self, keys: Collection[OffloadKey]) -> None:
112
+ def complete_load(
113
+ self, keys: Collection[OffloadKey], req_context: ReqContext
114
+ ) -> None:
115
for key in keys:
116
block = self._policy.get(key)
117
assert block is not None, f"Block {key!r} not found"
@@ -172,7 +174,10 @@ class CPUOffloadingManager(OffloadingManager):
174
)
175
176
def complete_store(
175
- self, keys: Collection[OffloadKey], success: bool = True
177
+ self,
178
+ keys: Collection[OffloadKey],
179
+ req_context: ReqContext,
180
+ success: bool = True,
181
) -> None:
182
stored_keys: list[OffloadKey] = []
183
vllm/v1/kv_offload/reuse_manager.py
+11
-6
@@ -105,16 +105,21 @@ class FilterReusedOffloadingManager(OffloadingManager):
105
) -> LoadStoreSpec:
106
return self._backing.prepare_load(keys, req_context)
107
108
- def touch(self, keys: Collection[OffloadKey]) -> None:
109
- return self._backing.touch(keys)
108
+ def touch(self, keys: Collection[OffloadKey], req_context: ReqContext) -> None:
109
+ return self._backing.touch(keys, req_context)
110
111
- def complete_load(self, keys: Collection[OffloadKey]) -> None:
112
- return self._backing.complete_load(keys)
111
+ def complete_load(
112
+ self, keys: Collection[OffloadKey], req_context: ReqContext
113
+ ) -> None:
114
+ return self._backing.complete_load(keys, req_context)
115
116
def complete_store(
115
- self, keys: Collection[OffloadKey], success: bool = True
117
+ self,
118
+ keys: Collection[OffloadKey],
119
+ req_context: ReqContext,
120
+ success: bool = True,
121
) -> None:
117
- return self._backing.complete_store(keys, success)
122
+ return self._backing.complete_store(keys, req_context, success)
123
124
def take_events(self) -> Iterable[OffloadingEvent]:
125
return self._backing.take_events()