@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 4353c9cb4a
Fork 0

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 4, 23:45 UTC

[KV Offload] Pass request context (#39185)

Signed-off-by: omerpaz95 <omerpaz95@gmail.com>

omerpaz95 committed Apr 19, 2026 at 08:54 UTC 4353c9cb4ac8d9818303c2e61282062aff4074c3
7 files changed +170 -81
tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py
+39 -15
@@ -32,8 +32,8 @@ def test_offloading_connector(request_runner, async_scheduling: bool):
32 # 3 blocks, store just the middle block (skip first and last)
33 # blocks = [0, 1, 2], [3, 4, 5], [6, 7, 8]
34 runner.new_request(token_ids=[0] * offloaded_block_size * 3)
35 - runner.manager.prepare_store.side_effect = lambda keys: generate_store_output(
36 - list(keys)[1:2]
35 + runner.manager.prepare_store.side_effect = (
36 + lambda keys, req_context: generate_store_output(list(keys)[1:2])
37 )
38 runner.run(decoded_tokens=[0])
39
@@ -45,18 +45,22 @@ def test_offloading_connector(request_runner, async_scheduling: bool):
45 runner.manager.prepare_store.assert_not_called()
46
47 # +1 token -> single block, fail prepare_store
48 - runner.manager.prepare_store.side_effect = lambda keys: None
48 + runner.manager.prepare_store.side_effect = lambda keys, req_context: None
49 runner.run(decoded_tokens=[0])
50 runner.manager.prepare_store.assert_called()
51
52 # 1 more block (+ token for async scheduling)
53 # now set block_hashes_to_store = []
54 - runner.manager.prepare_store.side_effect = lambda keys: generate_store_output([])
54 + runner.manager.prepare_store.side_effect = (
55 + lambda keys, req_context: generate_store_output([])
56 + )
57 runner.run(decoded_tokens=[0] * (offloaded_block_size + 1))
58
59 # 1 more block (+ token for kicking off offloading)
60 # now check touch was called with all 6 blocks
59 - runner.manager.prepare_store.side_effect = lambda keys: generate_store_output(keys)
61 + runner.manager.prepare_store.side_effect = (
62 + lambda keys, req_context: generate_store_output(keys)
63 + )
64 runner.run(
65 decoded_tokens=[0] * (offloaded_block_size + 1),
66 expected_stored_gpu_block_indexes=(15, 16, 17),
@@ -89,13 +93,17 @@ def test_offloading_connector(request_runner, async_scheduling: bool):
93 runner.new_request(
94 token_ids=[0] * gpu_block_size + [1] * (offloaded_block_size - gpu_block_size)
95 )
92 - runner.manager.prepare_store.side_effect = lambda keys: generate_store_output([])
96 + runner.manager.prepare_store.side_effect = (
97 + lambda keys, req_context: generate_store_output([])
98 + )
99 runner.run(decoded_tokens=[EOS_TOKEN_ID])
100 runner.manager.lookup.assert_not_called()
101
102 # single block lookup with no hits
103 runner.new_request(token_ids=[1] * offloaded_block_size)
98 - runner.manager.prepare_store.side_effect = lambda keys: generate_store_output([])
104 + runner.manager.prepare_store.side_effect = (
105 + lambda keys, req_context: generate_store_output([])
106 + )
107 runner.run(decoded_tokens=[EOS_TOKEN_ID])
108 runner.manager.lookup.assert_called()
109 assert len(list(runner.manager.lookup.call_args.args[0])) == 1
@@ -103,7 +111,9 @@ def test_offloading_connector(request_runner, async_scheduling: bool):
111 # single block lookup with a hit
112 runner.scheduler.reset_prefix_cache()
113 runner.new_request(token_ids=[0] * offloaded_block_size)
106 - runner.manager.prepare_store.side_effect = lambda keys: generate_store_output([])
114 + runner.manager.prepare_store.side_effect = (
115 + lambda keys, req_context: generate_store_output([])
116 + )
117 runner.manager.lookup.return_value = 1
118 runner.run(
119 decoded_tokens=[EOS_TOKEN_ID], expected_loaded_gpu_block_indexes=(0, 1, 2)
@@ -113,7 +123,9 @@ def test_offloading_connector(request_runner, async_scheduling: bool):
123 runner.new_request(
124 token_ids=[0] * offloaded_block_size * 2 + [1] * offloaded_block_size
125 )
116 - runner.manager.prepare_store.side_effect = lambda keys: generate_store_output([])
126 + runner.manager.prepare_store.side_effect = (
127 + lambda keys, req_context: generate_store_output([])
128 + )
129 runner.manager.lookup.return_value = 1
130 runner.run(
131 decoded_tokens=[EOS_TOKEN_ID], expected_loaded_gpu_block_indexes=(3, 4, 5)
@@ -164,14 +176,18 @@ def test_request_preemption(request_runner, async_scheduling: bool):
176 # 2 blocks, store all, without flushing
177 # blocks = [0, 1, 2], [3, 4, 5]
178 runner.new_request(token_ids=[0] * offloaded_block_size * 2)
167 - runner.manager.prepare_store.side_effect = lambda keys: generate_store_output(keys)
179 + runner.manager.prepare_store.side_effect = (
180 + lambda keys, req_context: generate_store_output(keys)
181 + )
182 runner.run(
183 decoded_tokens=[0],
184 complete_transfers=False,
185 )
186
187 # decode 2 more blocks - 1 gpu block, storing [6, 7, 8] (no flush)
174 - runner.manager.prepare_store.side_effect = lambda keys: generate_store_output(keys)
188 + runner.manager.prepare_store.side_effect = (
189 + lambda keys, req_context: generate_store_output(keys)
190 + )
191 runner.run(
192 decoded_tokens=[0] * (2 * offloaded_block_size - gpu_block_size),
193 complete_transfers=False,
@@ -195,7 +211,9 @@ def test_request_preemption(request_runner, async_scheduling: bool):
211 # request should now return from preemption
212 # re-load [0, ..., 8] from the CPU and store [9, 10, 11]
213 runner.manager.lookup.return_value = 3
198 - runner.manager.prepare_store.side_effect = lambda keys: generate_store_output(keys)
214 + runner.manager.prepare_store.side_effect = (
215 + lambda keys, req_context: generate_store_output(keys)
216 + )
217 runner.run(
218 decoded_tokens=[0] * gpu_block_size,
219 expected_loaded_gpu_block_indexes=(0, 1, 2, 3, 4, 5, 6, 7, 8),
@@ -222,7 +240,9 @@ def test_concurrent_lookups_of_the_same_prefix(request_runner, async_scheduling:
240
241 # store 1 blocks
242 runner.new_request(token_ids=[0] * offloaded_block_size)
225 - runner.manager.prepare_store.side_effect = lambda keys: generate_store_output(keys)
243 + runner.manager.prepare_store.side_effect = (
244 + lambda keys, req_context: generate_store_output(keys)
245 + )
246 runner.run(
247 decoded_tokens=[EOS_TOKEN_ID],
248 expected_stored_gpu_block_indexes=(0, 1, 2),
@@ -253,7 +273,9 @@ def test_concurrent_lookups_of_the_same_prefix(request_runner, async_scheduling:
273 assert transfer_jobs == list(runner.offloading_spec.handler.transfer_specs)
274
275 # complete transfers
256 - runner.manager.prepare_store.side_effect = lambda keys: generate_store_output([])
276 + runner.manager.prepare_store.side_effect = (
277 + lambda keys, req_context: generate_store_output([])
278 + )
279 runner.run(
280 decoded_tokens=[EOS_TOKEN_ID],
281 expected_loaded_gpu_block_indexes=(0, 1, 2),
@@ -278,7 +300,9 @@ def test_abort_loading_requests(request_runner, async_scheduling: bool):
300
301 # store 1 blocks
302 runner.new_request(token_ids=[0] * offloaded_block_size)
281 - runner.manager.prepare_store.side_effect = lambda keys: generate_store_output(keys)
303 + runner.manager.prepare_store.side_effect = (
304 + lambda keys, req_context: generate_store_output(keys)
305 + )
306 runner.run(
307 decoded_tokens=[EOS_TOKEN_ID],
308 expected_stored_gpu_block_indexes=(0, 1, 2),
tests/v1/kv_connector/unit/offloading_connector/utils.py
+1 -1
@@ -115,7 +115,7 @@ class MockOffloadingSpec(OffloadingSpec):
115
116 self.manager = MagicMock(spec=OffloadingManager)
117 self.manager.lookup.return_value = 0
118 - self.manager.prepare_load = lambda keys: MockLoadStoreSpec(keys)
118 + self.manager.prepare_load = lambda keys, req_context: MockLoadStoreSpec(keys)
119 self.handler = MockOffloadingHandler()
120
121 def get_manager(self) -> OffloadingManager:
tests/v1/kv_offload/test_cpu_manager.py
+68 -49
@@ -11,6 +11,7 @@ from vllm.v1.kv_offload.abstract import (
11 OffloadingEvent,
12 OffloadKey,
13 PrepareStoreOutput,
14 + ReqContext,
15 make_offload_key,
16 )
17 from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager
@@ -19,6 +20,14 @@ from vllm.v1.kv_offload.mediums import CPULoadStoreSpec
20 from vllm.v1.kv_offload.reuse_manager import FilterReusedOffloadingManager
21
22
23 +def make_req_context(kv_transfer_params: dict | None = None) -> ReqContext:
24 + """Create a ReqContext as production code would, from a request's params."""
25 + return ReqContext(kv_transfer_params=kv_transfer_params)
26 +
27 +
28 +_EMPTY_REQ_CTX = make_req_context()
29 +
30 +
31 @dataclass
32 class ExpectedPrepareStoreOutput:
33 keys_to_store: list[int]
@@ -103,7 +112,7 @@ def test_already_stored_block_not_evicted_during_prepare_store(eviction_policy):
112 )
113
114 # store [1, 2] and complete
106 - manager.prepare_store(to_keys([1, 2]))
115 + manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX)
116 manager.complete_store(to_keys([1, 2]))
117
118 # touch [1] to make block 2 the LRU candidate
@@ -113,7 +122,7 @@ def test_already_stored_block_not_evicted_during_prepare_store(eviction_policy):
122 # - block 2 is already stored -> filtered out of keys_to_store
123 # - block 2 must NOT be evicted even though it is the LRU candidate
124 # - block 1 (ID 0) is evicted instead; new blocks [3,4,5] get IDs 2,3,0
116 - prepare_store_output = manager.prepare_store(to_keys([2, 3, 4, 5]))
125 + prepare_store_output = manager.prepare_store(to_keys([2, 3, 4, 5]), _EMPTY_REQ_CTX)
126 verify_store_output(
127 prepare_store_output,
128 ExpectedPrepareStoreOutput(
@@ -127,7 +136,7 @@ def test_already_stored_block_not_evicted_during_prepare_store(eviction_policy):
136 manager.complete_store(to_keys([2, 3, 4, 5]))
137
138 # block 2 must still be present in the cache
130 - assert manager.lookup(to_keys([2])) == 1
139 + assert manager.lookup(to_keys([2]), _EMPTY_REQ_CTX) == 1
140
141
142 def test_cpu_manager():
@@ -140,7 +149,7 @@ def test_cpu_manager():
149 )
150
151 # prepare store [1, 2]
143 - prepare_store_output = cpu_manager.prepare_store(to_keys([1, 2]))
152 + prepare_store_output = cpu_manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX)
153 verify_store_output(
154 prepare_store_output,
155 ExpectedPrepareStoreOutput(
@@ -151,7 +160,7 @@ def test_cpu_manager():
160 )
161
162 # lookup [1, 2] -> not ready
154 - assert cpu_manager.lookup(to_keys([1, 2])) == 0
163 + assert cpu_manager.lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) == 0
164
165 # no events so far
166 assert list(cpu_manager.take_events()) == []
@@ -161,12 +170,14 @@ def test_cpu_manager():
170 verify_events(cpu_manager.take_events(), expected_stores=({1, 2},))
171
172 # lookup [1, 2]
164 - assert cpu_manager.lookup(to_keys([1])) == 1
165 - assert cpu_manager.lookup(to_keys([1, 2])) == 2
166 - assert cpu_manager.lookup(to_keys([1, 2, 3])) == 2
173 + assert cpu_manager.lookup(to_keys([1]), _EMPTY_REQ_CTX) == 1
174 + assert cpu_manager.lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) == 2
175 + assert cpu_manager.lookup(to_keys([1, 2, 3]), _EMPTY_REQ_CTX) == 2
176
177 # prepare store [2, 3, 4, 5] -> evicts [1]
169 - prepare_store_output = cpu_manager.prepare_store(to_keys([2, 3, 4, 5]))
178 + prepare_store_output = cpu_manager.prepare_store(
179 + to_keys([2, 3, 4, 5]), _EMPTY_REQ_CTX
180 + )
181 verify_store_output(
182 prepare_store_output,
183 ExpectedPrepareStoreOutput(
@@ -180,23 +191,23 @@ def test_cpu_manager():
191 verify_events(cpu_manager.take_events(), expected_evictions=({1},))
192
193 # prepare store with no space
183 - assert cpu_manager.prepare_store(to_keys([1, 6])) is None
194 + assert cpu_manager.prepare_store(to_keys([1, 6]), _EMPTY_REQ_CTX) is None
195
196 # complete store [2, 3, 4, 5]
197 cpu_manager.complete_store(to_keys([2, 3, 4, 5]))
198
199 # prepare load [2, 3]
189 - prepare_load_output = cpu_manager.prepare_load(to_keys([2, 3]))
200 + prepare_load_output = cpu_manager.prepare_load(to_keys([2, 3]), _EMPTY_REQ_CTX)
201 verify_load_output(prepare_load_output, [1, 2])
202
203 # prepare store with no space ([2, 3] is being loaded)
193 - assert cpu_manager.prepare_store(to_keys([6, 7, 8])) is None
204 + assert cpu_manager.prepare_store(to_keys([6, 7, 8]), _EMPTY_REQ_CTX) is None
205
206 # complete load [2, 3]
207 cpu_manager.complete_load(to_keys([2, 3]))
208
209 # prepare store [6, 7, 8] -> evicts [2, 3, 4] (oldest)
199 - prepare_store_output = cpu_manager.prepare_store(to_keys([6, 7, 8]))
210 + prepare_store_output = cpu_manager.prepare_store(to_keys([6, 7, 8]), _EMPTY_REQ_CTX)
211 verify_store_output(
212 prepare_store_output,
213 ExpectedPrepareStoreOutput(
@@ -213,7 +224,7 @@ def test_cpu_manager():
224 cpu_manager.touch(to_keys([5, 6, 7]))
225
226 # prepare store [7, 9] -> evicts [8] (oldest following previous touch)
216 - prepare_store_output = cpu_manager.prepare_store(to_keys([9]))
227 + prepare_store_output = cpu_manager.prepare_store(to_keys([9]), _EMPTY_REQ_CTX)
228 verify_store_output(
229 prepare_store_output,
230 ExpectedPrepareStoreOutput(
@@ -227,8 +238,8 @@ def test_cpu_manager():
238 cpu_manager.complete_store(to_keys([7, 9]), success=False)
239
240 # assert [7] is still stored, but [9] is not
230 - assert cpu_manager.lookup(to_keys([7])) == 1
231 - assert cpu_manager.lookup(to_keys([9])) == 0
241 + assert cpu_manager.lookup(to_keys([7]), _EMPTY_REQ_CTX) == 1
242 + assert cpu_manager.lookup(to_keys([9]), _EMPTY_REQ_CTX) == 0
243
244 verify_events(
245 cpu_manager.take_events(),
@@ -260,7 +271,9 @@ class TestARCPolicy:
271 cpu_manager, arc_policy = self._make_manager()
272
273 # prepare store [1, 2]
263 - prepare_store_output = cpu_manager.prepare_store(to_keys([1, 2]))
274 + prepare_store_output = cpu_manager.prepare_store(
275 + to_keys([1, 2]), _EMPTY_REQ_CTX
276 + )
277 verify_store_output(
278 prepare_store_output,
279 ExpectedPrepareStoreOutput(
@@ -271,7 +284,7 @@ class TestARCPolicy:
284 )
285
286 # lookup [1, 2] -> not ready
274 - assert cpu_manager.lookup(to_keys([1, 2])) == 0
287 + assert cpu_manager.lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) == 0
288
289 # no events so far
290 assert list(cpu_manager.take_events()) == []
@@ -281,9 +294,9 @@ class TestARCPolicy:
294 verify_events(cpu_manager.take_events(), expected_stores=({1, 2},))
295
296 # lookup [1, 2]
284 - assert cpu_manager.lookup(to_keys([1])) == 1
285 - assert cpu_manager.lookup(to_keys([1, 2])) == 2
286 - assert cpu_manager.lookup(to_keys([1, 2, 3])) == 2
297 + assert cpu_manager.lookup(to_keys([1]), _EMPTY_REQ_CTX) == 1
298 + assert cpu_manager.lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) == 2
299 + assert cpu_manager.lookup(to_keys([1, 2, 3]), _EMPTY_REQ_CTX) == 2
300
301 # blocks should be in T1 (recent)
302 assert len(arc_policy.t1) == 2
@@ -297,7 +310,7 @@ class TestARCPolicy:
310 cpu_manager, arc_policy = self._make_manager(enable_events=False)
311
312 # store and complete block 1
300 - cpu_manager.prepare_store(to_keys([1]))
313 + cpu_manager.prepare_store(to_keys([1]), _EMPTY_REQ_CTX)
314 cpu_manager.complete_store(to_keys([1]))
315
316 # block 1 starts in T1 (recent)
@@ -319,7 +332,9 @@ class TestARCPolicy:
332 cpu_manager, _ = self._make_manager()
333
334 # prepare and complete store [1, 2, 3, 4]
322 - prepare_store_output = cpu_manager.prepare_store(to_keys([1, 2, 3, 4]))
335 + prepare_store_output = cpu_manager.prepare_store(
336 + to_keys([1, 2, 3, 4]), _EMPTY_REQ_CTX
337 + )
338 verify_store_output(
339 prepare_store_output,
340 ExpectedPrepareStoreOutput(
@@ -331,19 +346,21 @@ class TestARCPolicy:
346 cpu_manager.complete_store(to_keys([1, 2, 3, 4]))
347
348 # prepare load [2, 3] (increases ref_cnt)
334 - prepare_load_output = cpu_manager.prepare_load(to_keys([2, 3]))
349 + prepare_load_output = cpu_manager.prepare_load(to_keys([2, 3]), _EMPTY_REQ_CTX)
350 verify_load_output(prepare_load_output, [1, 2])
351
352 # prepare store [5, 6, 7] with [2, 3] being loaded
353 # should fail because [2, 3] have ref_cnt > 0
339 - assert cpu_manager.prepare_store(to_keys([5, 6, 7])) is None
354 + assert cpu_manager.prepare_store(to_keys([5, 6, 7]), _EMPTY_REQ_CTX) is None
355
356 # complete load [2, 3]
357 cpu_manager.complete_load(to_keys([2, 3]))
358
359 # now prepare store [5, 6, 7] should succeed
360 # ARC will evict blocks one at a time from T1 as needed
346 - prepare_store_output = cpu_manager.prepare_store(to_keys([5, 6, 7]))
361 + prepare_store_output = cpu_manager.prepare_store(
362 + to_keys([5, 6, 7]), _EMPTY_REQ_CTX
363 + )
364 assert prepare_store_output is not None
365 # Should successfully evict enough blocks to make room (at least 1)
366 assert len(prepare_store_output.evicted_keys) >= 1
@@ -357,13 +374,13 @@ class TestARCPolicy:
374 cpu_manager, arc_policy = self._make_manager(num_blocks=2, enable_events=False)
375
376 # store blocks 1, 2 (fills cache)
360 - cpu_manager.prepare_store(to_keys([1, 2]))
377 + cpu_manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX)
378 cpu_manager.complete_store(to_keys([1, 2]))
379
380 initial_target = arc_policy.target_t1_size
381
382 # store block 3, evicting block 1 (moves to B1 ghost list)
366 - cpu_manager.prepare_store(to_keys([3]))
383 + cpu_manager.prepare_store(to_keys([3]), _EMPTY_REQ_CTX)
384 cpu_manager.complete_store(to_keys([3]))
385
386 # block 1 should be in B1 (ghost list)
@@ -384,7 +401,7 @@ class TestARCPolicy:
401 cpu_manager, arc_policy = self._make_manager(enable_events=False)
402
403 # store blocks 1, 2, 3, 4
387 - cpu_manager.prepare_store(to_keys([1, 2, 3, 4]))
404 + cpu_manager.prepare_store(to_keys([1, 2, 3, 4]), _EMPTY_REQ_CTX)
405 cpu_manager.complete_store(to_keys([1, 2, 3, 4]))
406
407 # promote blocks 3, 4 to T2 by touching them
@@ -399,7 +416,7 @@ class TestARCPolicy:
416 arc_policy.target_t1_size = 1
417
418 # store block 5, should evict from T1 (block 1, LRU in T1)
402 - output = cpu_manager.prepare_store(to_keys([5]))
419 + output = cpu_manager.prepare_store(to_keys([5]), _EMPTY_REQ_CTX)
420 assert output is not None
421 assert to_keys([1]) == output.evicted_keys
422
@@ -418,12 +435,12 @@ class TestARCPolicy:
435 cpu_manager, arc_policy = self._make_manager(num_blocks=2, enable_events=False)
436
437 # fill cache with blocks 1, 2
421 - cpu_manager.prepare_store(to_keys([1, 2]))
438 + cpu_manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX)
439 cpu_manager.complete_store(to_keys([1, 2]))
440
441 # store many blocks to fill ghost lists
442 for i in range(3, 20):
426 - cpu_manager.prepare_store(to_keys([i]))
443 + cpu_manager.prepare_store(to_keys([i]), _EMPTY_REQ_CTX)
444 cpu_manager.complete_store(to_keys([i]))
445
446 # ghost lists should not exceed cache_capacity
@@ -438,7 +455,7 @@ class TestARCPolicy:
455 cpu_manager, arc_policy = self._make_manager()
456
457 # store blocks 1, 2, 3, 4
441 - cpu_manager.prepare_store(to_keys([1, 2, 3, 4]))
458 + cpu_manager.prepare_store(to_keys([1, 2, 3, 4]), _EMPTY_REQ_CTX)
459 cpu_manager.complete_store(to_keys([1, 2, 3, 4]))
460
461 # promote 3, 4 to T2
@@ -453,7 +470,7 @@ class TestARCPolicy:
470 assert len(arc_policy.t2) == 3
471
472 # store block 5, should evict from T1 (block 2, only one in T1)
456 - prepare_store_output = cpu_manager.prepare_store(to_keys([5]))
473 + prepare_store_output = cpu_manager.prepare_store(to_keys([5]), _EMPTY_REQ_CTX)
474 verify_store_output(
475 prepare_store_output,
476 ExpectedPrepareStoreOutput(
@@ -471,11 +488,11 @@ class TestARCPolicy:
488 cpu_manager, arc_policy = self._make_manager()
489
490 # store blocks 1, 2, 3, 4
474 - cpu_manager.prepare_store(to_keys([1, 2, 3, 4]))
491 + cpu_manager.prepare_store(to_keys([1, 2, 3, 4]), _EMPTY_REQ_CTX)
492 cpu_manager.complete_store(to_keys([1, 2, 3, 4]))
493
494 # prepare store block 5 (will evict block 1)
478 - prepare_store_output = cpu_manager.prepare_store(to_keys([5]))
495 + prepare_store_output = cpu_manager.prepare_store(to_keys([5]), _EMPTY_REQ_CTX)
496 assert prepare_store_output is not None
497 assert len(prepare_store_output.evicted_keys) == 1
498
@@ -483,7 +500,7 @@ class TestARCPolicy:
500 cpu_manager.complete_store(to_keys([5]), success=False)
501
502 # block 5 should not be in cache
486 - assert cpu_manager.lookup(to_keys([5])) == 0
503 + assert cpu_manager.lookup(to_keys([5]), _EMPTY_REQ_CTX) == 0
504 # block 5 should not be in T1 or T2
505 assert to_keys([5])[0] not in arc_policy.t1
506 assert to_keys([5])[0] not in arc_policy.t2
@@ -500,11 +517,13 @@ class TestARCPolicy:
517 cpu_manager, arc_policy = self._make_manager()
518
519 # store [1, 2]
503 - cpu_manager.prepare_store(to_keys([1, 2]))
520 + cpu_manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX)
521 cpu_manager.complete_store(to_keys([1, 2]))
522
523 # store [3, 4, 5] -> evicts [1]
507 - prepare_store_output = cpu_manager.prepare_store(to_keys([3, 4, 5]))
524 + prepare_store_output = cpu_manager.prepare_store(
525 + to_keys([3, 4, 5]), _EMPTY_REQ_CTX
526 + )
527 assert prepare_store_output is not None
528 assert len(prepare_store_output.evicted_keys) == 1
529 cpu_manager.complete_store(to_keys([3, 4, 5]))
@@ -517,13 +536,13 @@ class TestARCPolicy:
536 assert len(arc_policy.t2) == 2
537
538 # store [6] -> should evict from T1 (4 is oldest in T1)
520 - prepare_store_output = cpu_manager.prepare_store(to_keys([6]))
539 + prepare_store_output = cpu_manager.prepare_store(to_keys([6]), _EMPTY_REQ_CTX)
540 assert prepare_store_output is not None
541 cpu_manager.complete_store(to_keys([6]))
542
543 # verify blocks 2, 3 (in T2) are still present
525 - assert cpu_manager.lookup(to_keys([2])) == 1
526 - assert cpu_manager.lookup(to_keys([3])) == 1
544 + assert cpu_manager.lookup(to_keys([2]), _EMPTY_REQ_CTX) == 1
545 + assert cpu_manager.lookup(to_keys([3]), _EMPTY_REQ_CTX) == 1
546
547 # verify events
548 events = list(cpu_manager.take_events())
@@ -543,34 +562,34 @@ def test_filter_reused_manager():
562 )
563
564 # Lookup [1, 2] -> 1st time, added to tracker but not eligible for store yet
546 - assert manager.lookup(to_keys([1, 2])) == 0
565 + assert manager.lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) == 0
566
567 # prepare store [1, 2] -> should be filtered
549 - prepare_store_output = manager.prepare_store(to_keys([1, 2]))
568 + prepare_store_output = manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX)
569 assert prepare_store_output is not None
570 assert prepare_store_output.keys_to_store == []
571
572 # Lookup [1] -> 2nd time, eligible now
554 - assert manager.lookup(to_keys([1])) == 0
573 + assert manager.lookup(to_keys([1]), _EMPTY_REQ_CTX) == 0
574
575 # prepare store [1, 2] -> [1] should be eligible, [2] should be filtered
557 - prepare_store_output = manager.prepare_store(to_keys([1, 2]))
576 + prepare_store_output = manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX)
577 assert prepare_store_output is not None
578 assert prepare_store_output.keys_to_store == to_keys([1])
579
580 # Lookup [3, 4] -> 1st time
581 # (evicts [2] from tracker since max_size is 3 and tracker has [1])
563 - assert manager.lookup(to_keys([3, 4])) == 0
582 + assert manager.lookup(to_keys([3, 4]), _EMPTY_REQ_CTX) == 0
583 # Verify [2] was evicted from the tracker (tracker now has: [1], [3], [4])
584 assert to_keys([2])[0] not in manager.counts
585
586 # Lookup [2] again -> (this adds [2] back to the tracker as 1st time)
568 - assert manager.lookup(to_keys([2])) == 0
587 + assert manager.lookup(to_keys([2]), _EMPTY_REQ_CTX) == 0
588 # Verify [2] was re-added with count=1 (not eligible yet)
589 assert manager.counts.get(to_keys([2])[0]) == 1
590
591 # prepare store [2] -> should still be filtered out since count was reset
573 - prepare_store_output = manager.prepare_store(to_keys([2]))
592 + prepare_store_output = manager.prepare_store(to_keys([2]), _EMPTY_REQ_CTX)
593 assert prepare_store_output is not None
594 assert prepare_store_output.keys_to_store == []
595
vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py
+11 -3
@@ -19,6 +19,7 @@ from vllm.v1.core.sched.output import SchedulerOutput
19 from vllm.v1.kv_offload.abstract import (
20 OffloadingManager,
21 OffloadKey,
22 + ReqContext,
23 get_offload_block_hash,
24 make_offload_key,
25 )
@@ -74,6 +75,7 @@ class RequestOffloadState:
75 config: SchedulerOffloadConfig
76 req: Request
77 group_states: tuple[RequestGroupState, ...] = field(init=False)
78 + req_context: ReqContext = field(init=False)
79 # number of hits in the GPU cache
80 num_locally_computed_tokens: int = 0
81
@@ -81,6 +83,7 @@ class RequestOffloadState:
83 self.group_states = tuple(
84 RequestGroupState() for _ in self.config.kv_group_configs
85 )
86 + self.req_context = ReqContext(kv_transfer_params=self.req.kv_transfer_params)
87
88 def update_offload_keys(self) -> None:
89 for group_config, group_state in zip(
@@ -181,7 +184,10 @@ class OffloadingConnectorScheduler:
184 return 0, False
185
186 start_block_idx = num_computed_tokens // group_config.offloaded_block_size
184 - hits = self.manager.lookup(offload_keys[start_block_idx:])
187 + hits = self.manager.lookup(
188 + offload_keys[start_block_idx:],
189 + req_status.req_context,
190 + )
191 if hits is None:
192 # indicates a lookup that should be tried later
193 return None, False
@@ -249,7 +255,7 @@ class OffloadingConnectorScheduler:
255 assert len(request.block_hashes) // self.config.block_size_factor >= num_blocks
256 offload_keys = group_state.offload_keys[start_block_idx:num_blocks]
257
252 - src_spec = self.manager.prepare_load(offload_keys)
258 + src_spec = self.manager.prepare_load(offload_keys, req_status.req_context)
259 dst_spec = GPULoadStoreSpec(
260 block_ids[num_computed_gpu_blocks:],
261 group_sizes=(num_pending_gpu_blocks,),
@@ -304,7 +310,9 @@ class OffloadingConnectorScheduler:
310 assert len(req.block_hashes) >= num_gpu_blocks
311
312 new_offload_keys = group_state.offload_keys[start_block_idx:num_blocks]
307 - store_output = self.manager.prepare_store(new_offload_keys)
313 + store_output = self.manager.prepare_store(
314 + new_offload_keys, req_status.req_context
315 + )
316 if store_output is None:
317 logger.warning(
318 "Request %s: cannot store %s blocks", req_id, num_new_blocks
vllm/v1/kv_offload/abstract.py
+24 -4
@@ -30,7 +30,7 @@ The class provides the following primitives:
30 from abc import ABC, abstractmethod
31 from collections.abc import Iterable
32 from dataclasses import dataclass
33 -from typing import NewType
33 +from typing import Any, NewType
34
35 # `OffloadKey` identifies an offloaded block. It combines a block hash with
36 # its KV cache group index, encoded as raw bytes to avoid tuple GC overhead.
@@ -53,6 +53,11 @@ def get_offload_group_idx(key: OffloadKey) -> int:
53 return int.from_bytes(key[-4:], "big", signed=False)
54
55
56 +@dataclass
57 +class ReqContext:
58 + kv_transfer_params: dict[str, Any] | None = None
59 +
60 +
61 class LoadStoreSpec(ABC):
62 """
63 Abstract metadata that encapsulates information allowing a worker
@@ -86,13 +91,18 @@ class OffloadingEvent:
91
92 class OffloadingManager(ABC):
93 @abstractmethod
89 - def lookup(self, keys: Iterable[OffloadKey]) -> int | None:
94 + def lookup(
95 + self,
96 + keys: Iterable[OffloadKey],
97 + req_context: ReqContext,
98 + ) -> int | None:
99 """
100 Finds the length of the maximal series of blocks, starting from the
101 first one, that are all offloaded.
102
103 Args:
104 keys: the keys identifying the blocks to lookup.
105 + req_context: per-request context (e.g. kv_transfer_params).
106
107 Returns:
108 An integer representing the maximal number of blocks that
@@ -103,7 +113,11 @@ class OffloadingManager(ABC):
113 pass
114
115 @abstractmethod
106 - def prepare_load(self, keys: Iterable[OffloadKey]) -> LoadStoreSpec:
116 + def prepare_load(
117 + self,
118 + keys: Iterable[OffloadKey],
119 + req_context: ReqContext,
120 + ) -> LoadStoreSpec:
121 """
122 Prepare the given blocks to be read.
123 The given blocks will be protected from eviction until
@@ -112,6 +126,7 @@ class OffloadingManager(ABC):
126
127 Args:
128 keys: the keys identifying the blocks.
129 + req_context: per-request context (e.g. kv_transfer_params).
130
131 Returns:
132 A LoadStoreSpec that can be used by a worker to locate and load
@@ -139,7 +154,11 @@ class OffloadingManager(ABC):
154 return
155
156 @abstractmethod
142 - def prepare_store(self, keys: Iterable[OffloadKey]) -> PrepareStoreOutput | None:
157 + def prepare_store(
158 + self,
159 + keys: Iterable[OffloadKey],
160 + req_context: ReqContext,
161 + ) -> PrepareStoreOutput | None:
162 """
163 Prepare the given blocks to be offloaded.
164 The given blocks will be protected from eviction until
@@ -147,6 +166,7 @@ class OffloadingManager(ABC):
166
167 Args:
168 keys: the keys identifying the blocks.
169 + req_context: per-request context (e.g. kv_transfer_params).
170
171 Returns:
172 A PrepareStoreOutput indicating which blocks need storing,
vllm/v1/kv_offload/cpu/manager.py
+16 -3
@@ -9,6 +9,7 @@ from vllm.v1.kv_offload.abstract import (
9 OffloadingManager,
10 OffloadKey,
11 PrepareStoreOutput,
12 + ReqContext,
13 )
14 from vllm.v1.kv_offload.cpu.policies.abstract import BlockStatus, CachePolicy
15 from vllm.v1.kv_offload.cpu.policies.arc import ARCCachePolicy
@@ -83,7 +84,11 @@ class CPUOffloadingManager(OffloadingManager):
84
85 # --- OffloadingManager interface ---
86
86 - def lookup(self, keys: Iterable[OffloadKey]) -> int | None:
87 + def lookup(
88 + self,
89 + keys: Iterable[OffloadKey],
90 + req_context: ReqContext,
91 + ) -> int | None:
92 hit_count = 0
93 for key in keys:
94 block = self._policy.get(key)
@@ -92,7 +97,11 @@ class CPUOffloadingManager(OffloadingManager):
97 hit_count += 1
98 return hit_count
99
95 - def prepare_load(self, keys: Iterable[OffloadKey]) -> LoadStoreSpec:
100 + def prepare_load(
101 + self,
102 + keys: Iterable[OffloadKey],
103 + req_context: ReqContext,
104 + ) -> LoadStoreSpec:
105 blocks = []
106 for key in keys:
107 block = self._policy.get(key)
@@ -112,7 +121,11 @@ class CPUOffloadingManager(OffloadingManager):
121 assert block.ref_cnt > 0, f"Block {key!r} ref_cnt is already 0"
122 block.ref_cnt -= 1
123
115 - def prepare_store(self, keys: Iterable[OffloadKey]) -> PrepareStoreOutput | None:
124 + def prepare_store(
125 + self,
126 + keys: Iterable[OffloadKey],
127 + req_context: ReqContext,
128 + ) -> PrepareStoreOutput | None:
129 keys_list = list(keys)
130
131 # filter out blocks that are already stored
vllm/v1/kv_offload/reuse_manager.py
+11 -6
@@ -16,6 +16,7 @@ from vllm.v1.kv_offload.abstract import (
16 OffloadingManager,
17 OffloadKey,
18 PrepareStoreOutput,
19 + ReqContext,
20 )
21
22
@@ -65,7 +66,7 @@ class FilterReusedOffloadingManager(OffloadingManager):
66 # Intercepted methods
67 # ------------------------------------------------------------------
68
68 - def lookup(self, keys: Iterable[OffloadKey]) -> int | None:
69 + def lookup(self, keys: Iterable[OffloadKey], req_context: ReqContext) -> int | None:
70 """Record each key, then delegate lookup to backing manager."""
71 keys = list(keys)
72 for key in keys:
@@ -76,9 +77,11 @@ class FilterReusedOffloadingManager(OffloadingManager):
77 if len(self.counts) >= self.max_tracker_size:
78 self.counts.popitem(last=False) # evict LRU
79 self.counts[key] = 1
79 - return self._backing.lookup(keys)
80 + return self._backing.lookup(keys, req_context)
81
81 - def prepare_store(self, keys: Iterable[OffloadKey]) -> PrepareStoreOutput | None:
82 + def prepare_store(
83 + self, keys: Iterable[OffloadKey], req_context: ReqContext
84 + ) -> PrepareStoreOutput | None:
85 """Filter out blocks below threshold, then delegate to backing.
86
87 Filtering is evaluated *before* calling the backing manager's
@@ -93,14 +96,16 @@ class FilterReusedOffloadingManager(OffloadingManager):
96 # Passing an empty list is intentional and safe — CPUOffloadingManager
97 # handles it correctly, returning a PrepareStoreOutput with empty lists.
98 # Delegate to the backing manager with only the eligible keys.
96 - return self._backing.prepare_store(eligible)
99 + return self._backing.prepare_store(eligible, req_context)
100
101 # ------------------------------------------------------------------
102 # Delegated methods
103 # ------------------------------------------------------------------
104
102 - def prepare_load(self, keys: Iterable[OffloadKey]) -> LoadStoreSpec:
103 - return self._backing.prepare_load(keys)
105 + def prepare_load(
106 + self, keys: Iterable[OffloadKey], req_context: ReqContext
107 + ) -> LoadStoreSpec:
108 + return self._backing.prepare_load(keys, req_context)
109
110 def touch(self, keys: Iterable[OffloadKey]) -> None:
111 return self._backing.touch(keys)