@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 512c5eb455
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 5, 11:00 UTC

[kv_offload+HMA][5/N]: Track group block hashes and block IDs (#37109)

Signed-off-by: Or Ozeri <oro@il.ibm.com>

Or Ozeri committed Apr 8, 2026 at 19:50 UTC 512c5eb455594b7366b6ae5c27e934bdcc5eed6e
11 files changed +564 -497
tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py
+18 -41
@@ -6,6 +6,7 @@ import pytest
6
7 from tests.v1.kv_connector.unit.offloading_connector.utils import (
8 generate_store_output,
9 + to_keys,
10 )
11 from tests.v1.kv_connector.unit.utils import EOS_TOKEN_ID
12 from vllm.distributed.kv_events import BlockRemoved, BlockStored
@@ -31,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)
34 - runner.manager.prepare_store.side_effect = (
35 - lambda block_hashes: generate_store_output(list(block_hashes)[1:2])
35 + runner.manager.prepare_store.side_effect = lambda keys: generate_store_output(
36 + list(keys)[1:2]
37 )
38 runner.run(decoded_tokens=[0])
39
@@ -44,22 +45,18 @@ 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
47 - runner.manager.prepare_store.side_effect = lambda block_hashes: None
48 + runner.manager.prepare_store.side_effect = lambda keys: 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 = []
53 - runner.manager.prepare_store.side_effect = (
54 - lambda block_hashes: generate_store_output([])
55 - )
54 + runner.manager.prepare_store.side_effect = lambda keys: generate_store_output([])
55 runner.run(decoded_tokens=[0] * (offloaded_block_size + 1))
56
57 # 1 more block (+ token for kicking off offloading)
58 # now check touch was called with all 6 blocks
60 - runner.manager.prepare_store.side_effect = (
61 - lambda block_hashes: generate_store_output(block_hashes)
62 - )
59 + runner.manager.prepare_store.side_effect = lambda keys: generate_store_output(keys)
60 runner.run(
61 decoded_tokens=[0] * (offloaded_block_size + 1),
62 expected_stored_gpu_block_indexes=(15, 16, 17),
@@ -92,17 +89,13 @@ def test_offloading_connector(request_runner, async_scheduling: bool):
89 runner.new_request(
90 token_ids=[0] * gpu_block_size + [1] * (offloaded_block_size - gpu_block_size)
91 )
95 - runner.manager.prepare_store.side_effect = (
96 - lambda block_hashes: generate_store_output([])
97 - )
92 + runner.manager.prepare_store.side_effect = lambda keys: generate_store_output([])
93 runner.run(decoded_tokens=[EOS_TOKEN_ID])
94 runner.manager.lookup.assert_not_called()
95
96 # single block lookup with no hits
97 runner.new_request(token_ids=[1] * offloaded_block_size)
103 - runner.manager.prepare_store.side_effect = (
104 - lambda block_hashes: generate_store_output([])
105 - )
98 + runner.manager.prepare_store.side_effect = lambda keys: generate_store_output([])
99 runner.run(decoded_tokens=[EOS_TOKEN_ID])
100 runner.manager.lookup.assert_called()
101 assert len(list(runner.manager.lookup.call_args.args[0])) == 1
@@ -110,9 +103,7 @@ def test_offloading_connector(request_runner, async_scheduling: bool):
103 # single block lookup with a hit
104 runner.scheduler.reset_prefix_cache()
105 runner.new_request(token_ids=[0] * offloaded_block_size)
113 - runner.manager.prepare_store.side_effect = (
114 - lambda block_hashes: generate_store_output([])
115 - )
106 + runner.manager.prepare_store.side_effect = lambda keys: generate_store_output([])
107 runner.manager.lookup.return_value = 1
108 runner.run(
109 decoded_tokens=[EOS_TOKEN_ID], expected_loaded_gpu_block_indexes=(0, 1, 2)
@@ -122,9 +113,7 @@ def test_offloading_connector(request_runner, async_scheduling: bool):
113 runner.new_request(
114 token_ids=[0] * offloaded_block_size * 2 + [1] * offloaded_block_size
115 )
125 - runner.manager.prepare_store.side_effect = (
126 - lambda block_hashes: generate_store_output([])
127 - )
116 + runner.manager.prepare_store.side_effect = lambda keys: generate_store_output([])
117 runner.manager.lookup.return_value = 1
118 runner.run(
119 decoded_tokens=[EOS_TOKEN_ID], expected_loaded_gpu_block_indexes=(3, 4, 5)
@@ -136,10 +125,10 @@ def test_offloading_connector(request_runner, async_scheduling: bool):
125
126 def take_events() -> Iterable[OffloadingEvent]:
127 yield OffloadingEvent(
139 - block_hashes=to_hashes([1, 2, 3]), block_size=16, medium="A", removed=False
128 + keys=to_keys([1, 2, 3]), block_size=16, medium="A", removed=False
129 )
130 yield OffloadingEvent(
142 - block_hashes=to_hashes([4, 5, 6]), block_size=32, medium="B", removed=True
131 + keys=to_keys([4, 5, 6]), block_size=32, medium="B", removed=True
132 )
133
134 runner.manager.take_events.side_effect = take_events
@@ -179,18 +168,14 @@ def test_request_preemption(request_runner, async_scheduling: bool):
168 # 2 blocks, store all, without flushing
169 # blocks = [0, 1, 2], [3, 4, 5]
170 runner.new_request(token_ids=[0] * offloaded_block_size * 2)
182 - runner.manager.prepare_store.side_effect = (
183 - lambda block_hashes: generate_store_output(block_hashes)
184 - )
171 + runner.manager.prepare_store.side_effect = lambda keys: generate_store_output(keys)
172 runner.run(
173 decoded_tokens=[0],
174 complete_transfers=False,
175 )
176
177 # decode 2 more blocks - 1 gpu block, storing [6, 7, 8] (no flush)
191 - runner.manager.prepare_store.side_effect = (
192 - lambda block_hashes: generate_store_output(block_hashes)
193 - )
178 + runner.manager.prepare_store.side_effect = lambda keys: generate_store_output(keys)
179 runner.run(
180 decoded_tokens=[0] * (2 * offloaded_block_size - gpu_block_size),
181 complete_transfers=False,
@@ -214,9 +199,7 @@ def test_request_preemption(request_runner, async_scheduling: bool):
199 # request should now return from preemption
200 # re-load [0, ..., 8] from the CPU and store [9, 10, 11]
201 runner.manager.lookup.return_value = 3
217 - runner.manager.prepare_store.side_effect = (
218 - lambda block_hashes: generate_store_output(block_hashes)
219 - )
202 + runner.manager.prepare_store.side_effect = lambda keys: generate_store_output(keys)
203 runner.run(
204 decoded_tokens=[0] * gpu_block_size,
205 expected_loaded_gpu_block_indexes=(0, 1, 2, 3, 4, 5, 6, 7, 8),
@@ -243,9 +226,7 @@ def test_concurrent_lookups_of_the_same_prefix(request_runner, async_scheduling:
226
227 # store 1 blocks
228 runner.new_request(token_ids=[0] * offloaded_block_size)
246 - runner.manager.prepare_store.side_effect = (
247 - lambda block_hashes: generate_store_output(block_hashes)
248 - )
229 + runner.manager.prepare_store.side_effect = lambda keys: generate_store_output(keys)
230 runner.run(
231 decoded_tokens=[EOS_TOKEN_ID],
232 expected_stored_gpu_block_indexes=(0, 1, 2),
@@ -276,9 +257,7 @@ def test_concurrent_lookups_of_the_same_prefix(request_runner, async_scheduling:
257 assert transfer_jobs == list(runner.offloading_spec.handler.transfer_specs)
258
259 # complete transfers
279 - runner.manager.prepare_store.side_effect = (
280 - lambda block_hashes: generate_store_output([])
281 - )
260 + runner.manager.prepare_store.side_effect = lambda keys: generate_store_output([])
261 runner.run(
262 decoded_tokens=[EOS_TOKEN_ID],
263 expected_loaded_gpu_block_indexes=(0, 1, 2),
@@ -303,9 +282,7 @@ def test_abort_loading_requests(request_runner, async_scheduling: bool):
282
283 # store 1 blocks
284 runner.new_request(token_ids=[0] * offloaded_block_size)
306 - runner.manager.prepare_store.side_effect = (
307 - lambda block_hashes: generate_store_output(block_hashes)
308 - )
285 + runner.manager.prepare_store.side_effect = lambda keys: generate_store_output(keys)
286 runner.run(
287 decoded_tokens=[EOS_TOKEN_ID],
288 expected_stored_gpu_block_indexes=(0, 1, 2),
tests/v1/kv_connector/unit/offloading_connector/utils.py
+22 -17
@@ -27,7 +27,6 @@ from vllm.forward_context import ForwardContext
27 from vllm.utils.hashing import sha256
28 from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
29 from vllm.v1.core.kv_cache_utils import (
30 - BlockHash,
30 get_request_block_hasher,
31 init_none_hash,
32 )
@@ -41,7 +40,9 @@ from vllm.v1.kv_cache_interface import (
40 from vllm.v1.kv_offload.abstract import (
41 LoadStoreSpec,
42 OffloadingManager,
43 + OffloadKey,
44 PrepareStoreOutput,
45 + make_offload_key,
46 )
47 from vllm.v1.kv_offload.mediums import GPULoadStoreSpec
48 from vllm.v1.kv_offload.spec import OffloadingSpec
@@ -55,16 +56,20 @@ from vllm.v1.request import Request
56 from vllm.v1.structured_output import StructuredOutputManager
57
58
59 +def to_keys(int_ids: list[int]) -> list[OffloadKey]:
60 + return [make_offload_key(str(i).encode(), 0) for i in int_ids]
61 +
62 +
63 class MockLoadStoreSpec(LoadStoreSpec):
59 - def __init__(self, block_hashes: Iterable[BlockHash]):
60 - self.block_hashes: list[BlockHash] = list(block_hashes)
64 + def __init__(self, offload_keys: Iterable[OffloadKey]):
65 + self.offload_keys: list[OffloadKey] = list(offload_keys)
66
67 @staticmethod
68 def medium() -> str:
69 return "Mock"
70
71 def __repr__(self) -> str:
67 - return repr(self.block_hashes)
72 + return repr(self.offload_keys)
73
74
75 class MockOffloadingHandler(OffloadingHandler):
@@ -110,9 +115,7 @@ class MockOffloadingSpec(OffloadingSpec):
115
116 self.manager = MagicMock(spec=OffloadingManager)
117 self.manager.lookup.return_value = 0
113 - self.manager.prepare_load = lambda block_hashes: (
114 - MockLoadStoreSpec(block_hashes)
115 - )
118 + self.manager.prepare_load = lambda keys: MockLoadStoreSpec(keys)
119 self.handler = MockOffloadingHandler()
120
121 def get_manager(self) -> OffloadingManager:
@@ -231,8 +234,10 @@ class RequestRunner:
234 assert isinstance(manager, MagicMock)
235 self.manager: MagicMock = manager
236
234 - assert connector_scheduler.gpu_block_size == gpu_block_size
235 - assert connector_scheduler.offloaded_block_size == offloaded_block_size
237 + assert len(connector_scheduler.config.kv_group_configs) == 1
238 + kv_group_config = connector_scheduler.config.kv_group_configs[0]
239 + assert kv_group_config.gpu_block_size == gpu_block_size
240 + assert kv_group_config.offloaded_block_size == offloaded_block_size
241
242 # extract OffloadingSpec of worker_connector
243 connector_worker = self.worker_connector.connector_worker
@@ -307,11 +312,11 @@ class RequestRunner:
312 for block_id in gpu_spec.block_ids:
313 gpu_block_indices.append(self.gpu_block_index[block_id.item()])
314
310 - # list of (block_hash, sub_block_offset)
315 + # list of (offload_key, sub_block_offset)
316 offload_addresses: list[Any] = []
312 - for block_hash in offload_spec.block_hashes:
317 + for offload_key in offload_spec.offload_keys:
318 for sub_block_idx in range(block_size_factor):
314 - offload_addresses.append((block_hash, sub_block_idx))
319 + offload_addresses.append((offload_key, sub_block_idx))
320
321 if store:
322 assert len(gpu_block_indices) == len(offload_addresses)
@@ -510,10 +515,10 @@ def request_runner():
515 yield runner_factory # pass factory to the test
516
517
513 -def generate_store_output(block_hashes: Iterable[BlockHash]):
514 - block_hashes = list(block_hashes)
518 +def generate_store_output(keys: Iterable[OffloadKey]):
519 + keys = list(keys)
520 return PrepareStoreOutput(
516 - block_hashes_to_store=list(block_hashes),
517 - store_spec=MockLoadStoreSpec(block_hashes),
518 - block_hashes_evicted=[],
521 + keys_to_store=list(keys),
522 + store_spec=MockLoadStoreSpec(keys),
523 + evicted_keys=[],
524 )
tests/v1/kv_offload/test_cpu_manager.py
+137 -134
@@ -6,11 +6,12 @@ from dataclasses import dataclass
6 import numpy as np
7 import pytest
8
9 -from vllm.v1.core.kv_cache_utils import BlockHash
9 from vllm.v1.kv_offload.abstract import (
10 LoadStoreSpec,
11 OffloadingEvent,
12 + OffloadKey,
13 PrepareStoreOutput,
14 + make_offload_key,
15 )
16 from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager
17 from vllm.v1.kv_offload.cpu.policies.arc import ARCCachePolicy
@@ -20,13 +21,13 @@ from vllm.v1.kv_offload.reuse_manager import FilterReusedOffloadingManager
21
22 @dataclass
23 class ExpectedPrepareStoreOutput:
23 - block_hashes_to_store: list[int]
24 + keys_to_store: list[int]
25 store_block_ids: list[int]
25 - block_hashes_evicted: list[int]
26 + evicted_keys: list[int]
27
28
28 -def to_hashes(int_hashes: list[int]) -> list[BlockHash]:
29 - return [BlockHash(str(i).encode()) for i in int_hashes]
29 +def to_keys(int_ids: list[int]) -> list[OffloadKey]:
30 + return [make_offload_key(str(i).encode(), 0) for i in int_ids]
31
32
33 def verify_store_output(
@@ -34,11 +35,11 @@ def verify_store_output(
35 expected_prepare_store_output: ExpectedPrepareStoreOutput,
36 ):
37 assert prepare_store_output is not None
37 - assert prepare_store_output.block_hashes_to_store == to_hashes(
38 - expected_prepare_store_output.block_hashes_to_store
38 + assert prepare_store_output.keys_to_store == to_keys(
39 + expected_prepare_store_output.keys_to_store
40 )
40 - assert prepare_store_output.block_hashes_evicted == to_hashes(
41 - expected_prepare_store_output.block_hashes_evicted
41 + assert prepare_store_output.evicted_keys == to_keys(
42 + expected_prepare_store_output.evicted_keys
43 )
44 store_spec = prepare_store_output.store_spec
45 assert isinstance(store_spec, CPULoadStoreSpec)
@@ -62,21 +63,23 @@ def verify_events(
63 expected_stores: tuple[set[int], ...] = (),
64 expected_evictions: tuple[set[int], ...] = (),
65 ):
65 - stores: list[set[BlockHash]] = []
66 - evictions: list[set[BlockHash]] = []
66 + stores: list[set[OffloadKey]] = []
67 + evictions: list[set[OffloadKey]] = []
68 for event in events:
69 assert event.medium == CPULoadStoreSpec.medium()
70 assert event.block_size == block_size
71 if event.removed:
71 - evictions.append(set(event.block_hashes))
72 + evictions.append(set(event.keys))
73 else:
73 - stores.append(set(event.block_hashes))
74 + stores.append(set(event.keys))
75
75 - def to_hash_sets(int_sets: tuple[set[int], ...]) -> tuple[set[BlockHash], ...]:
76 - return tuple([set(to_hashes(list(int_set))) for int_set in int_sets])
76 + def to_key_sets(
77 + int_sets: tuple[set[int], ...],
78 + ) -> tuple[set[OffloadKey], ...]:
79 + return tuple([set(to_keys(list(int_set))) for int_set in int_sets])
80
78 - assert tuple(evictions) == to_hash_sets(expected_evictions)
79 - assert tuple(stores) == to_hash_sets(expected_stores)
81 + assert tuple(evictions) == to_key_sets(expected_evictions)
82 + assert tuple(stores) == to_key_sets(expected_stores)
83
84
85 @pytest.mark.parametrize("eviction_policy", ["lru", "arc"])
@@ -104,31 +107,31 @@ def test_already_stored_block_not_evicted_during_prepare_store(eviction_policy):
107 )
108
109 # store [1, 2] and complete
107 - manager.prepare_store(to_hashes([1, 2]))
108 - manager.complete_store(to_hashes([1, 2]))
110 + manager.prepare_store(to_keys([1, 2]))
111 + manager.complete_store(to_keys([1, 2]))
112
113 # touch [1] to make block 2 the LRU candidate
111 - manager.touch(to_hashes([1]))
114 + manager.touch(to_keys([1]))
115
116 # prepare_store([2, 3, 4, 5]):
114 - # - block 2 is already stored → filtered out of block_hashes_to_store
117 + # - block 2 is already stored -> filtered out of keys_to_store
118 # - block 2 must NOT be evicted even though it is the LRU candidate
119 # - block 1 (ID 0) is evicted instead; new blocks [3,4,5] get IDs 2,3,0
117 - prepare_store_output = manager.prepare_store(to_hashes([2, 3, 4, 5]))
120 + prepare_store_output = manager.prepare_store(to_keys([2, 3, 4, 5]))
121 verify_store_output(
122 prepare_store_output,
123 ExpectedPrepareStoreOutput(
121 - block_hashes_to_store=[3, 4, 5],
124 + keys_to_store=[3, 4, 5],
125 store_block_ids=[2, 3, 0],
123 - block_hashes_evicted=[1], # block 1 evicted, not block 2
126 + evicted_keys=[1], # block 1 evicted, not block 2
127 ),
128 )
129
130 # complete_store must not silently drop block 2
128 - manager.complete_store(to_hashes([2, 3, 4, 5]))
131 + manager.complete_store(to_keys([2, 3, 4, 5]))
132
133 # block 2 must still be present in the cache
131 - assert manager.lookup(to_hashes([2])) == 1
134 + assert manager.lookup(to_keys([2])) == 1
135
136
137 def test_cpu_manager():
@@ -142,41 +145,41 @@ def test_cpu_manager():
145 )
146
147 # prepare store [1, 2]
145 - prepare_store_output = cpu_manager.prepare_store(to_hashes([1, 2]))
148 + prepare_store_output = cpu_manager.prepare_store(to_keys([1, 2]))
149 verify_store_output(
150 prepare_store_output,
151 ExpectedPrepareStoreOutput(
149 - block_hashes_to_store=[1, 2],
152 + keys_to_store=[1, 2],
153 store_block_ids=[0, 1],
151 - block_hashes_evicted=[],
154 + evicted_keys=[],
155 ),
156 )
157
158 # lookup [1, 2] -> not ready
156 - assert cpu_manager.lookup(to_hashes([1, 2])) == 0
159 + assert cpu_manager.lookup(to_keys([1, 2])) == 0
160
161 # no events so far
162 assert list(cpu_manager.take_events()) == []
163
164 # complete store [1, 2]
162 - cpu_manager.complete_store(to_hashes([1, 2]))
165 + cpu_manager.complete_store(to_keys([1, 2]))
166 verify_events(
167 cpu_manager.take_events(), block_size=block_size, expected_stores=({1, 2},)
168 )
169
170 # lookup [1, 2]
168 - assert cpu_manager.lookup(to_hashes([1])) == 1
169 - assert cpu_manager.lookup(to_hashes([1, 2])) == 2
170 - assert cpu_manager.lookup(to_hashes([1, 2, 3])) == 2
171 + assert cpu_manager.lookup(to_keys([1])) == 1
172 + assert cpu_manager.lookup(to_keys([1, 2])) == 2
173 + assert cpu_manager.lookup(to_keys([1, 2, 3])) == 2
174
175 # prepare store [2, 3, 4, 5] -> evicts [1]
173 - prepare_store_output = cpu_manager.prepare_store(to_hashes([2, 3, 4, 5]))
176 + prepare_store_output = cpu_manager.prepare_store(to_keys([2, 3, 4, 5]))
177 verify_store_output(
178 prepare_store_output,
179 ExpectedPrepareStoreOutput(
177 - block_hashes_to_store=[3, 4, 5],
180 + keys_to_store=[3, 4, 5],
181 store_block_ids=[2, 3, 0],
179 - block_hashes_evicted=[1],
182 + evicted_keys=[1],
183 ),
184 )
185
@@ -186,55 +189,55 @@ def test_cpu_manager():
189 )
190
191 # prepare store with no space
189 - assert cpu_manager.prepare_store(to_hashes([1, 6])) is None
192 + assert cpu_manager.prepare_store(to_keys([1, 6])) is None
193
194 # complete store [2, 3, 4, 5]
192 - cpu_manager.complete_store(to_hashes([2, 3, 4, 5]))
195 + cpu_manager.complete_store(to_keys([2, 3, 4, 5]))
196
197 # prepare load [2, 3]
195 - prepare_load_output = cpu_manager.prepare_load(to_hashes([2, 3]))
198 + prepare_load_output = cpu_manager.prepare_load(to_keys([2, 3]))
199 verify_load_output(prepare_load_output, [1, 2])
200
201 # prepare store with no space ([2, 3] is being loaded)
199 - assert cpu_manager.prepare_store(to_hashes([6, 7, 8])) is None
202 + assert cpu_manager.prepare_store(to_keys([6, 7, 8])) is None
203
204 # complete load [2, 3]
202 - cpu_manager.complete_load(to_hashes([2, 3]))
205 + cpu_manager.complete_load(to_keys([2, 3]))
206
207 # prepare store [6, 7, 8] -> evicts [2, 3, 4] (oldest)
205 - prepare_store_output = cpu_manager.prepare_store(to_hashes([6, 7, 8]))
208 + prepare_store_output = cpu_manager.prepare_store(to_keys([6, 7, 8]))
209 verify_store_output(
210 prepare_store_output,
211 ExpectedPrepareStoreOutput(
209 - block_hashes_to_store=[6, 7, 8],
212 + keys_to_store=[6, 7, 8],
213 store_block_ids=[3, 2, 1],
211 - block_hashes_evicted=[2, 3, 4],
214 + evicted_keys=[2, 3, 4],
215 ),
216 )
217
218 # complete store [6, 7, 8]
216 - cpu_manager.complete_store(to_hashes([6, 7, 8]))
219 + cpu_manager.complete_store(to_keys([6, 7, 8]))
220
221 # touch [5, 6, 7] (move to end of LRU order)
219 - cpu_manager.touch(to_hashes([5, 6, 7]))
222 + cpu_manager.touch(to_keys([5, 6, 7]))
223
224 # prepare store [7, 9] -> evicts [8] (oldest following previous touch)
222 - prepare_store_output = cpu_manager.prepare_store(to_hashes([9]))
225 + prepare_store_output = cpu_manager.prepare_store(to_keys([9]))
226 verify_store_output(
227 prepare_store_output,
228 ExpectedPrepareStoreOutput(
226 - block_hashes_to_store=[9],
229 + keys_to_store=[9],
230 store_block_ids=[1],
228 - block_hashes_evicted=[8],
231 + evicted_keys=[8],
232 ),
233 )
234
235 # complete store [7, 9] with failure
233 - cpu_manager.complete_store(to_hashes([7, 9]), success=False)
236 + cpu_manager.complete_store(to_keys([7, 9]), success=False)
237
238 # assert [7] is still stored, but [9] is not
236 - assert cpu_manager.lookup(to_hashes([7])) == 1
237 - assert cpu_manager.lookup(to_hashes([9])) == 0
239 + assert cpu_manager.lookup(to_keys([7])) == 1
240 + assert cpu_manager.lookup(to_keys([9])) == 0
241
242 verify_events(
243 cpu_manager.take_events(),
@@ -268,32 +271,32 @@ class TestARCPolicy:
271 cpu_manager, arc_policy = self._make_manager()
272
273 # prepare store [1, 2]
271 - prepare_store_output = cpu_manager.prepare_store(to_hashes([1, 2]))
274 + prepare_store_output = cpu_manager.prepare_store(to_keys([1, 2]))
275 verify_store_output(
276 prepare_store_output,
277 ExpectedPrepareStoreOutput(
275 - block_hashes_to_store=[1, 2],
278 + keys_to_store=[1, 2],
279 store_block_ids=[0, 1],
277 - block_hashes_evicted=[],
280 + evicted_keys=[],
281 ),
282 )
283
284 # lookup [1, 2] -> not ready
282 - assert cpu_manager.lookup(to_hashes([1, 2])) == 0
285 + assert cpu_manager.lookup(to_keys([1, 2])) == 0
286
287 # no events so far
288 assert list(cpu_manager.take_events()) == []
289
290 # complete store [1, 2]
288 - cpu_manager.complete_store(to_hashes([1, 2]))
291 + cpu_manager.complete_store(to_keys([1, 2]))
292 verify_events(
293 cpu_manager.take_events(), block_size=256, expected_stores=({1, 2},)
294 )
295
296 # lookup [1, 2]
294 - assert cpu_manager.lookup(to_hashes([1])) == 1
295 - assert cpu_manager.lookup(to_hashes([1, 2])) == 2
296 - assert cpu_manager.lookup(to_hashes([1, 2, 3])) == 2
297 + assert cpu_manager.lookup(to_keys([1])) == 1
298 + assert cpu_manager.lookup(to_keys([1, 2])) == 2
299 + assert cpu_manager.lookup(to_keys([1, 2, 3])) == 2
300
301 # blocks should be in T1 (recent)
302 assert len(arc_policy.t1) == 2
@@ -307,19 +310,19 @@ class TestARCPolicy:
310 cpu_manager, arc_policy = self._make_manager(enable_events=False)
311
312 # store and complete block 1
310 - cpu_manager.prepare_store(to_hashes([1]))
311 - cpu_manager.complete_store(to_hashes([1]))
313 + cpu_manager.prepare_store(to_keys([1]))
314 + cpu_manager.complete_store(to_keys([1]))
315
316 # block 1 starts in T1 (recent)
314 - assert to_hashes([1])[0] in arc_policy.t1
315 - assert to_hashes([1])[0] not in arc_policy.t2
317 + assert to_keys([1])[0] in arc_policy.t1
318 + assert to_keys([1])[0] not in arc_policy.t2
319
320 # touch block 1 (simulate second access)
318 - cpu_manager.touch(to_hashes([1]))
321 + cpu_manager.touch(to_keys([1]))
322
323 # block 1 should now be in T2 (frequent)
321 - assert to_hashes([1])[0] not in arc_policy.t1
322 - assert to_hashes([1])[0] in arc_policy.t2
324 + assert to_keys([1])[0] not in arc_policy.t1
325 + assert to_keys([1])[0] in arc_policy.t2
326
327 def test_eviction_with_load(self):
328 """
@@ -329,34 +332,34 @@ class TestARCPolicy:
332 cpu_manager, _ = self._make_manager()
333
334 # prepare and complete store [1, 2, 3, 4]
332 - prepare_store_output = cpu_manager.prepare_store(to_hashes([1, 2, 3, 4]))
335 + prepare_store_output = cpu_manager.prepare_store(to_keys([1, 2, 3, 4]))
336 verify_store_output(
337 prepare_store_output,
338 ExpectedPrepareStoreOutput(
336 - block_hashes_to_store=[1, 2, 3, 4],
339 + keys_to_store=[1, 2, 3, 4],
340 store_block_ids=[0, 1, 2, 3],
338 - block_hashes_evicted=[],
341 + evicted_keys=[],
342 ),
343 )
341 - cpu_manager.complete_store(to_hashes([1, 2, 3, 4]))
344 + cpu_manager.complete_store(to_keys([1, 2, 3, 4]))
345
346 # prepare load [2, 3] (increases ref_cnt)
344 - prepare_load_output = cpu_manager.prepare_load(to_hashes([2, 3]))
347 + prepare_load_output = cpu_manager.prepare_load(to_keys([2, 3]))
348 verify_load_output(prepare_load_output, [1, 2])
349
350 # prepare store [5, 6, 7] with [2, 3] being loaded
351 # should fail because [2, 3] have ref_cnt > 0
349 - assert cpu_manager.prepare_store(to_hashes([5, 6, 7])) is None
352 + assert cpu_manager.prepare_store(to_keys([5, 6, 7])) is None
353
354 # complete load [2, 3]
352 - cpu_manager.complete_load(to_hashes([2, 3]))
355 + cpu_manager.complete_load(to_keys([2, 3]))
356
357 # now prepare store [5, 6, 7] should succeed
358 # ARC will evict blocks one at a time from T1 as needed
356 - prepare_store_output = cpu_manager.prepare_store(to_hashes([5, 6, 7]))
359 + prepare_store_output = cpu_manager.prepare_store(to_keys([5, 6, 7]))
360 assert prepare_store_output is not None
361 # Should successfully evict enough blocks to make room (at least 1)
359 - assert len(prepare_store_output.block_hashes_evicted) >= 1
362 + assert len(prepare_store_output.evicted_keys) >= 1
363
364 def test_adaptive_target(self):
365 """
@@ -367,21 +370,21 @@ class TestARCPolicy:
370 cpu_manager, arc_policy = self._make_manager(num_blocks=2, enable_events=False)
371
372 # store blocks 1, 2 (fills cache)
370 - cpu_manager.prepare_store(to_hashes([1, 2]))
371 - cpu_manager.complete_store(to_hashes([1, 2]))
373 + cpu_manager.prepare_store(to_keys([1, 2]))
374 + cpu_manager.complete_store(to_keys([1, 2]))
375
376 initial_target = arc_policy.target_t1_size
377
378 # store block 3, evicting block 1 (moves to B1 ghost list)
376 - cpu_manager.prepare_store(to_hashes([3]))
377 - cpu_manager.complete_store(to_hashes([3]))
379 + cpu_manager.prepare_store(to_keys([3]))
380 + cpu_manager.complete_store(to_keys([3]))
381
382 # block 1 should be in B1 (ghost list)
380 - assert to_hashes([1])[0] in arc_policy.b1
383 + assert to_keys([1])[0] in arc_policy.b1
384
385 # touch block 1 (cache miss, but in B1)
386 # this should increase target_t1_size (favor recency)
384 - cpu_manager.touch(to_hashes([1]))
387 + cpu_manager.touch(to_keys([1]))
388
389 # target should have increased
390 assert arc_policy.target_t1_size > initial_target
@@ -394,11 +397,11 @@ class TestARCPolicy:
397 cpu_manager, arc_policy = self._make_manager(enable_events=False)
398
399 # store blocks 1, 2, 3, 4
397 - cpu_manager.prepare_store(to_hashes([1, 2, 3, 4]))
398 - cpu_manager.complete_store(to_hashes([1, 2, 3, 4]))
400 + cpu_manager.prepare_store(to_keys([1, 2, 3, 4]))
401 + cpu_manager.complete_store(to_keys([1, 2, 3, 4]))
402
403 # promote blocks 3, 4 to T2 by touching them
401 - cpu_manager.touch(to_hashes([3, 4]))
404 + cpu_manager.touch(to_keys([3, 4]))
405
406 # now: T1 = {1, 2}, T2 = {3, 4}
407 assert len(arc_policy.t1) == 2
@@ -409,16 +412,16 @@ class TestARCPolicy:
412 arc_policy.target_t1_size = 1
413
414 # store block 5, should evict from T1 (block 1, LRU in T1)
412 - output = cpu_manager.prepare_store(to_hashes([5]))
415 + output = cpu_manager.prepare_store(to_keys([5]))
416 assert output is not None
414 - assert to_hashes([1]) == output.block_hashes_evicted
417 + assert to_keys([1]) == output.evicted_keys
418
416 - cpu_manager.complete_store(to_hashes([5]))
419 + cpu_manager.complete_store(to_keys([5]))
420
421 # block 1 should be in B1 (ghost list)
419 - assert to_hashes([1])[0] in arc_policy.b1
422 + assert to_keys([1])[0] in arc_policy.b1
423 # block 5 should be in T1
421 - assert to_hashes([5])[0] in arc_policy.t1
424 + assert to_keys([5])[0] in arc_policy.t1
425
426 def test_ghost_list_bounds(self):
427 """
@@ -428,13 +431,13 @@ class TestARCPolicy:
431 cpu_manager, arc_policy = self._make_manager(num_blocks=2, enable_events=False)
432
433 # fill cache with blocks 1, 2
431 - cpu_manager.prepare_store(to_hashes([1, 2]))
432 - cpu_manager.complete_store(to_hashes([1, 2]))
434 + cpu_manager.prepare_store(to_keys([1, 2]))
435 + cpu_manager.complete_store(to_keys([1, 2]))
436
437 # store many blocks to fill ghost lists
438 for i in range(3, 20):
436 - cpu_manager.prepare_store(to_hashes([i]))
437 - cpu_manager.complete_store(to_hashes([i]))
439 + cpu_manager.prepare_store(to_keys([i]))
440 + cpu_manager.complete_store(to_keys([i]))
441
442 # ghost lists should not exceed cache_capacity
443 assert len(arc_policy.b1) <= arc_policy.cache_capacity
@@ -448,28 +451,28 @@ class TestARCPolicy:
451 cpu_manager, arc_policy = self._make_manager()
452
453 # store blocks 1, 2, 3, 4
451 - cpu_manager.prepare_store(to_hashes([1, 2, 3, 4]))
452 - cpu_manager.complete_store(to_hashes([1, 2, 3, 4]))
454 + cpu_manager.prepare_store(to_keys([1, 2, 3, 4]))
455 + cpu_manager.complete_store(to_keys([1, 2, 3, 4]))
456
457 # promote 3, 4 to T2
455 - cpu_manager.touch(to_hashes([3, 4]))
458 + cpu_manager.touch(to_keys([3, 4]))
459
460 # T1 = {1, 2}, T2 = {3, 4}
461 # touch [1, 3, 4] - should promote 1 to T2, and move 3,4 to end of T2
459 - cpu_manager.touch(to_hashes([1, 3, 4]))
462 + cpu_manager.touch(to_keys([1, 3, 4]))
463
464 # T1 = {2}, T2 = {1, 3, 4} (in that order, with 4 most recent)
465 assert len(arc_policy.t1) == 1
466 assert len(arc_policy.t2) == 3
467
468 # store block 5, should evict from T1 (block 2, only one in T1)
466 - prepare_store_output = cpu_manager.prepare_store(to_hashes([5]))
469 + prepare_store_output = cpu_manager.prepare_store(to_keys([5]))
470 verify_store_output(
471 prepare_store_output,
472 ExpectedPrepareStoreOutput(
470 - block_hashes_to_store=[5],
473 + keys_to_store=[5],
474 store_block_ids=[1], # reuses block 2's storage
472 - block_hashes_evicted=[2],
475 + evicted_keys=[2],
476 ),
477 )
478
@@ -481,25 +484,25 @@ class TestARCPolicy:
484 cpu_manager, arc_policy = self._make_manager()
485
486 # store blocks 1, 2, 3, 4
484 - cpu_manager.prepare_store(to_hashes([1, 2, 3, 4]))
485 - cpu_manager.complete_store(to_hashes([1, 2, 3, 4]))
487 + cpu_manager.prepare_store(to_keys([1, 2, 3, 4]))
488 + cpu_manager.complete_store(to_keys([1, 2, 3, 4]))
489
490 # prepare store block 5 (will evict block 1)
488 - prepare_store_output = cpu_manager.prepare_store(to_hashes([5]))
491 + prepare_store_output = cpu_manager.prepare_store(to_keys([5]))
492 assert prepare_store_output is not None
490 - assert len(prepare_store_output.block_hashes_evicted) == 1
493 + assert len(prepare_store_output.evicted_keys) == 1
494
495 # complete store with failure
493 - cpu_manager.complete_store(to_hashes([5]), success=False)
496 + cpu_manager.complete_store(to_keys([5]), success=False)
497
498 # block 5 should not be in cache
496 - assert cpu_manager.lookup(to_hashes([5])) == 0
499 + assert cpu_manager.lookup(to_keys([5])) == 0
500 # block 5 should not be in T1 or T2
498 - assert to_hashes([5])[0] not in arc_policy.t1
499 - assert to_hashes([5])[0] not in arc_policy.t2
501 + assert to_keys([5])[0] not in arc_policy.t1
502 + assert to_keys([5])[0] not in arc_policy.t2
503
504 # evicted block should still be gone (in B1 ghost list)
502 - evicted_hash = prepare_store_output.block_hashes_evicted[0]
505 + evicted_hash = prepare_store_output.evicted_keys[0]
506 assert evicted_hash in arc_policy.b1
507
508 def test_full_scenario(self):
@@ -510,30 +513,30 @@ class TestARCPolicy:
513 cpu_manager, arc_policy = self._make_manager()
514
515 # store [1, 2]
513 - cpu_manager.prepare_store(to_hashes([1, 2]))
514 - cpu_manager.complete_store(to_hashes([1, 2]))
516 + cpu_manager.prepare_store(to_keys([1, 2]))
517 + cpu_manager.complete_store(to_keys([1, 2]))
518
519 # store [3, 4, 5] -> evicts [1]
517 - prepare_store_output = cpu_manager.prepare_store(to_hashes([3, 4, 5]))
520 + prepare_store_output = cpu_manager.prepare_store(to_keys([3, 4, 5]))
521 assert prepare_store_output is not None
519 - assert len(prepare_store_output.block_hashes_evicted) == 1
520 - cpu_manager.complete_store(to_hashes([3, 4, 5]))
522 + assert len(prepare_store_output.evicted_keys) == 1
523 + cpu_manager.complete_store(to_keys([3, 4, 5]))
524
525 # promote some blocks to T2
523 - cpu_manager.touch(to_hashes([2, 3]))
526 + cpu_manager.touch(to_keys([2, 3]))
527
528 # T1 has {4, 5}, T2 has {2, 3}
529 assert len(arc_policy.t1) == 2
530 assert len(arc_policy.t2) == 2
531
532 # store [6] -> should evict from T1 (4 is oldest in T1)
530 - prepare_store_output = cpu_manager.prepare_store(to_hashes([6]))
533 + prepare_store_output = cpu_manager.prepare_store(to_keys([6]))
534 assert prepare_store_output is not None
532 - cpu_manager.complete_store(to_hashes([6]))
535 + cpu_manager.complete_store(to_keys([6]))
536
537 # verify blocks 2, 3 (in T2) are still present
535 - assert cpu_manager.lookup(to_hashes([2])) == 1
536 - assert cpu_manager.lookup(to_hashes([3])) == 1
538 + assert cpu_manager.lookup(to_keys([2])) == 1
539 + assert cpu_manager.lookup(to_keys([3])) == 1
540
541 # verify events
542 events = list(cpu_manager.take_events())
@@ -554,35 +557,35 @@ def test_filter_reused_manager():
557 )
558
559 # Lookup [1, 2] -> 1st time, added to tracker but not eligible for store yet
557 - assert manager.lookup(to_hashes([1, 2])) == 0
560 + assert manager.lookup(to_keys([1, 2])) == 0
561
562 # prepare store [1, 2] -> should be filtered
560 - prepare_store_output = manager.prepare_store(to_hashes([1, 2]))
563 + prepare_store_output = manager.prepare_store(to_keys([1, 2]))
564 assert prepare_store_output is not None
562 - assert prepare_store_output.block_hashes_to_store == []
565 + assert prepare_store_output.keys_to_store == []
566
567 # Lookup [1] -> 2nd time, eligible now
565 - assert manager.lookup(to_hashes([1])) == 0
568 + assert manager.lookup(to_keys([1])) == 0
569
570 # prepare store [1, 2] -> [1] should be eligible, [2] should be filtered
568 - prepare_store_output = manager.prepare_store(to_hashes([1, 2]))
571 + prepare_store_output = manager.prepare_store(to_keys([1, 2]))
572 assert prepare_store_output is not None
570 - assert prepare_store_output.block_hashes_to_store == to_hashes([1])
573 + assert prepare_store_output.keys_to_store == to_keys([1])
574
575 # Lookup [3, 4] -> 1st time
576 # (evicts [2] from tracker since max_size is 3 and tracker has [1])
574 - assert manager.lookup(to_hashes([3, 4])) == 0
577 + assert manager.lookup(to_keys([3, 4])) == 0
578 # Verify [2] was evicted from the tracker (tracker now has: [1], [3], [4])
576 - assert to_hashes([2])[0] not in manager.counts
579 + assert to_keys([2])[0] not in manager.counts
580
581 # Lookup [2] again -> (this adds [2] back to the tracker as 1st time)
579 - assert manager.lookup(to_hashes([2])) == 0
582 + assert manager.lookup(to_keys([2])) == 0
583 # Verify [2] was re-added with count=1 (not eligible yet)
581 - assert manager.counts.get(to_hashes([2])[0]) == 1
584 + assert manager.counts.get(to_keys([2])[0]) == 1
585
586 # prepare store [2] -> should still be filtered out since count was reset
584 - prepare_store_output = manager.prepare_store(to_hashes([2]))
587 + prepare_store_output = manager.prepare_store(to_keys([2]))
588 assert prepare_store_output is not None
586 - assert prepare_store_output.block_hashes_to_store == []
589 + assert prepare_store_output.keys_to_store == []
590
588 - manager.complete_store(to_hashes([1]))
591 + manager.complete_store(to_keys([1]))
vllm/distributed/kv_transfer/kv_connector/utils.py
+1 -1
@@ -301,7 +301,7 @@ def kv_postprocess_blksize_and_layout_on_receive(cache, indices, block_size_rati
301
302 def yield_req_data(
303 scheduler_output,
304 -) -> Iterator[tuple[str, tuple[list[int], ...], bool]]:
304 +) -> Iterator[tuple[str, tuple[list[int], ...] | None, bool]]:
305 """
306 Yields:
307 (req_id, new_block_id_groups, preempted)
vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py
+195 -118
@@ -2,8 +2,9 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 from collections import defaultdict
4 from collections.abc import Iterable
5 +from dataclasses import dataclass, field
6 from itertools import islice
6 -from typing import Any
7 +from typing import Any, NamedTuple
8
9 from vllm.distributed.kv_events import BlockRemoved, BlockStored, KVCacheEvent
10 from vllm.distributed.kv_transfer.kv_connector.utils import yield_req_data
@@ -14,9 +15,13 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.common import (
15 )
16 from vllm.logger import init_logger
17 from vllm.v1.core.kv_cache_manager import KVCacheBlocks
17 -from vllm.v1.core.kv_cache_utils import BlockHash
18 from vllm.v1.core.sched.output import SchedulerOutput
19 -from vllm.v1.kv_offload.abstract import OffloadingManager
19 +from vllm.v1.kv_offload.abstract import (
20 + OffloadingManager,
21 + OffloadKey,
22 + get_offload_block_hash,
23 + make_offload_key,
24 +)
25 from vllm.v1.kv_offload.mediums import GPULoadStoreSpec
26 from vllm.v1.kv_offload.spec import OffloadingSpec
27 from vllm.v1.kv_offload.worker.worker import TransferSpec
@@ -26,46 +31,103 @@ from vllm.v1.request import Request
31 logger = init_logger(__name__)
32
33
34 +class GroupOffloadConfig(NamedTuple):
35 + group_idx: int
36 + gpu_block_size: int
37 + offloaded_block_size: int
38 + hash_block_size_factor: int
39 +
40 +
41 +class SchedulerOffloadConfig(NamedTuple):
42 + kv_group_configs: tuple[GroupOffloadConfig, ...]
43 + block_size_factor: int
44 +
45 + @classmethod
46 + def from_spec(cls, spec: OffloadingSpec) -> "SchedulerOffloadConfig":
47 + return cls(
48 + kv_group_configs=tuple(
49 + GroupOffloadConfig(
50 + group_idx=idx,
51 + gpu_block_size=gpu_block_size,
52 + offloaded_block_size=gpu_block_size * spec.block_size_factor,
53 + hash_block_size_factor=(
54 + (gpu_block_size * spec.block_size_factor)
55 + // spec.hash_block_size
56 + ),
57 + )
58 + for idx, gpu_block_size in enumerate(spec.gpu_block_size)
59 + ),
60 + block_size_factor=spec.block_size_factor,
61 + )
62 +
63 +
64 +@dataclass
65 +class RequestGroupState:
66 + offload_keys: list[OffloadKey] = field(default_factory=list)
67 + block_ids: list[int] = field(default_factory=list)
68 + # index of next block (of size offloaded_block_size) to offload
69 + next_stored_block_idx: int = 0
70 +
71 +
72 +@dataclass(slots=True)
73 +class RequestOffloadState:
74 + config: SchedulerOffloadConfig
75 + req: Request
76 + group_states: tuple[RequestGroupState, ...] = field(init=False)
77 + # number of hits in the GPU cache
78 + num_locally_computed_tokens: int = 0
79 +
80 + def __post_init__(self) -> None:
81 + self.group_states = tuple(
82 + RequestGroupState() for _ in self.config.kv_group_configs
83 + )
84 +
85 + def update_offload_keys(self) -> None:
86 + for group_config, group_state in zip(
87 + self.config.kv_group_configs, self.group_states
88 + ):
89 + for req_block_hash in islice(
90 + self.req.block_hashes,
91 + group_config.hash_block_size_factor * len(group_state.offload_keys)
92 + + group_config.hash_block_size_factor
93 + - 1,
94 + None,
95 + group_config.hash_block_size_factor,
96 + ):
97 + group_state.offload_keys.append(
98 + make_offload_key(req_block_hash, group_config.group_idx)
99 + )
100 +
101 + def update_block_id_groups(
102 + self, new_block_id_groups: tuple[list[int], ...] | None
103 + ) -> None:
104 + if new_block_id_groups is None:
105 + return
106 +
107 + assert len(new_block_id_groups) == len(self.group_states)
108 + for group_state, new_blocks in zip(self.group_states, new_block_id_groups):
109 + group_state.block_ids.extend(new_blocks)
110 +
111 +
112 class OffloadingConnectorScheduler:
113 """Implementation of Scheduler side methods"""
114
115 def __init__(self, spec: OffloadingSpec):
33 - assert len(spec.gpu_block_size) == 1
34 - self.gpu_block_size = spec.gpu_block_size[0]
35 - self.offloaded_block_size = self.gpu_block_size * spec.block_size_factor
36 - self.block_size_factor = spec.block_size_factor
116 + self.config = SchedulerOffloadConfig.from_spec(spec)
117 self.manager: OffloadingManager = spec.get_manager()
118
39 - self._requests: dict[ReqId, Request] = {}
40 - # list of GPU block IDs per request
41 - self._request_block_ids: dict[ReqId, list[int]] = {}
119 + self._req_status: dict[ReqId, RequestOffloadState] = {}
120 # requests to load for the current scheduler step
121 self._reqs_to_load: dict[ReqId, TransferSpec] = {}
44 - # request blocks are stored in order
45 - # index of next block (of size offloaded_block_size) to offload
46 - self._next_stored_block_idx: dict[ReqId, int] = {}
122 # if GPU prefix caching is enabled,
123 # track loaded blocks to avoid redundant loads
49 - self._blocks_being_loaded: set[BlockHash] | None = (
124 + self._blocks_being_loaded: set[OffloadKey] | None = (
125 set() if spec.vllm_config.cache_config.enable_prefix_caching else None
126 )
127
53 - # request ID -> set(block hashes being stored/load)
54 - self._reqs_being_stored = defaultdict[ReqId, set[BlockHash]](set)
55 - self._reqs_being_loaded = defaultdict[ReqId, set[BlockHash]](set)
56 -
57 - def _get_block_hashes(
58 - self,
59 - req: Request,
60 - start_idx: int = 0,
61 - end_idx: int | None = None,
62 - ) -> Iterable[BlockHash]:
63 - return islice(
64 - req.block_hashes,
65 - self.block_size_factor * start_idx + self.block_size_factor - 1,
66 - self.block_size_factor * end_idx if end_idx else None,
67 - self.block_size_factor,
68 - )
128 + # request ID -> set(offload keys being stored/loaded)
129 + self._reqs_being_stored = defaultdict[ReqId, set[OffloadKey]](set)
130 + self._reqs_being_loaded = defaultdict[ReqId, set[OffloadKey]](set)
131
132 def get_num_new_matched_tokens(
133 self, request: Request, num_computed_tokens: int
@@ -89,22 +151,37 @@ class OffloadingConnectorScheduler:
151 - `True` if tokens will be loaded asynchronously
152 (between scheduler steps).
153 """
92 - num_blocks = request.num_tokens // self.offloaded_block_size
154 + if req_status := self._req_status.get(request.request_id):
155 + # make sure block IDs are cleared
156 + for group_state in req_status.group_states:
157 + group_state.block_ids.clear()
158 + else:
159 + req_status = RequestOffloadState(config=self.config, req=request)
160 + req_status.update_offload_keys()
161 + self._req_status[request.request_id] = req_status
162
94 - assert len(request.block_hashes) // self.block_size_factor == num_blocks
95 - block_hashes = self._get_block_hashes(request)
163 + req_status.num_locally_computed_tokens = num_computed_tokens
164
97 - self.manager.touch(block_hashes)
165 + # Below assertions will be removed once this function supports HMA
166 + assert len(self.config.kv_group_configs) == 1
167 + assert len(req_status.group_states) == 1
168 + group_config = self.config.kv_group_configs[0]
169 + group_state = req_status.group_states[0]
170
99 - full_block_tokens = self.offloaded_block_size * num_blocks
100 - if full_block_tokens - num_computed_tokens < self.offloaded_block_size:
171 + num_blocks = request.num_tokens // group_config.offloaded_block_size
172 +
173 + assert len(request.block_hashes) // self.config.block_size_factor == num_blocks
174 + offload_keys = group_state.offload_keys
175 +
176 + self.manager.touch(offload_keys)
177 +
178 + full_block_tokens = group_config.offloaded_block_size * num_blocks
179 + if full_block_tokens - num_computed_tokens < group_config.offloaded_block_size:
180 # we can load less than a block, skip
181 return 0, False
182
104 - start_block_idx = num_computed_tokens // self.offloaded_block_size
105 - hits = self.manager.lookup(
106 - self._get_block_hashes(request, start_idx=start_block_idx)
107 - )
183 + start_block_idx = num_computed_tokens // group_config.offloaded_block_size
184 + hits = self.manager.lookup(offload_keys[start_block_idx:])
185 if hits is None:
186 # indicates a lookup that should be tried later
187 return None, False
@@ -112,7 +189,8 @@ class OffloadingConnectorScheduler:
189 return 0, False
190
191 num_hit_tokens = (
115 - self.offloaded_block_size * (start_block_idx + hits) - num_computed_tokens
192 + group_config.offloaded_block_size * (start_block_idx + hits)
193 + - num_computed_tokens
194 )
195 logger.debug(
196 "Request %s hit %s offloaded tokens after %s GPU hit tokens",
@@ -120,147 +198,147 @@ class OffloadingConnectorScheduler:
198 num_hit_tokens,
199 num_computed_tokens,
200 )
123 - if num_hit_tokens < self.offloaded_block_size:
201 + if num_hit_tokens < group_config.offloaded_block_size:
202 return 0, False
203
126 - if self._blocks_being_loaded:
127 - block_hashes = self._get_block_hashes(
128 - request, start_idx=start_block_idx, end_idx=start_block_idx + hits
204 + if self._blocks_being_loaded and any(
205 + key in self._blocks_being_loaded
206 + for key in offload_keys[start_block_idx : start_block_idx + hits]
207 + ):
208 + # hit blocks are being loaded, delay request
209 + logger.debug(
210 + "Delaying request %s since some of its blocks are already being loaded",
211 + request.request_id,
212 )
130 -
131 - if any(
132 - block_hash in self._blocks_being_loaded for block_hash in block_hashes
133 - ):
134 - # hit blocks are being loaded, delay request
135 - logger.debug(
136 - "Delaying request %s since some of its blocks are already"
137 - " being loaded",
138 - request.request_id,
139 - )
140 - return None, False
213 + return None, False
214
215 return num_hit_tokens, True
216
217 def update_state_after_alloc(
218 self, request: Request, blocks: KVCacheBlocks, num_external_tokens: int
219 ):
147 - self._requests[request.request_id] = request
148 - # the block ids are updated in _get_reqs_to_store
149 - self._request_block_ids[request.request_id] = []
150 -
220 if num_external_tokens == 0:
221 return
222
223 + req_status = self._req_status[request.request_id]
224 block_groups = blocks.get_block_ids()
225 +
226 + # Below assertions will be removed once this function supports HMA
227 + assert len(self.config.kv_group_configs) == 1
228 + assert len(req_status.group_states) == 1
229 + assert len(block_groups) == 1
230 block_ids = block_groups[0]
231 + group_config = self.config.kv_group_configs[0]
232 + group_state = req_status.group_states[0]
233
234 num_computed_gpu_blocks = sum(
235 block.block_hash is not None for block in blocks.blocks[0]
236 )
160 - num_computed_tokens = num_computed_gpu_blocks * self.gpu_block_size
237 + num_computed_tokens = num_computed_gpu_blocks * group_config.gpu_block_size
238 full_block_tokens = num_computed_tokens + num_external_tokens
162 - assert full_block_tokens % self.offloaded_block_size == 0
239 + assert full_block_tokens % group_config.offloaded_block_size == 0
240
241 num_pending_gpu_blocks = len(block_ids) - num_computed_gpu_blocks
165 - assert num_external_tokens == num_pending_gpu_blocks * self.gpu_block_size
242 + assert (
243 + num_external_tokens == num_pending_gpu_blocks * group_config.gpu_block_size
244 + )
245
167 - start_block_idx = num_computed_tokens // self.offloaded_block_size
168 - num_blocks = full_block_tokens // self.offloaded_block_size
246 + start_block_idx = num_computed_tokens // group_config.offloaded_block_size
247 + num_blocks = full_block_tokens // group_config.offloaded_block_size
248
170 - assert len(request.block_hashes) // self.block_size_factor >= num_blocks
171 - block_hashes = self._get_block_hashes(
172 - request, start_idx=start_block_idx, end_idx=num_blocks
173 - )
249 + assert len(request.block_hashes) // self.config.block_size_factor >= num_blocks
250 + offload_keys = group_state.offload_keys[start_block_idx:num_blocks]
251
175 - src_spec = self.manager.prepare_load(block_hashes)
252 + src_spec = self.manager.prepare_load(offload_keys)
253 dst_spec = GPULoadStoreSpec(
254 block_ids[num_computed_gpu_blocks:],
255 group_sizes=(num_pending_gpu_blocks,),
256 block_indices=(num_computed_gpu_blocks,),
257 )
258
182 - block_hashes = self._get_block_hashes(
183 - request, start_idx=start_block_idx, end_idx=num_blocks
184 - )
185 -
259 self._reqs_to_load[request.request_id] = (src_spec, dst_spec)
260 req_blocks_being_loaded = self._reqs_being_loaded[request.request_id]
188 - req_blocks_being_loaded.update(block_hashes)
189 - self._next_stored_block_idx[request.request_id] = num_blocks
261 + req_blocks_being_loaded.update(offload_keys)
262 + group_state.next_stored_block_idx = num_blocks
263
264 if self._blocks_being_loaded is not None:
265 self._blocks_being_loaded.update(req_blocks_being_loaded)
266
267 def _get_reqs_to_store(self, scheduler_output: SchedulerOutput):
268 + # Below assertion will be removed once this function supports HMA
269 + assert len(self.config.kv_group_configs) == 1
270 + group_config = self.config.kv_group_configs[0]
271 +
272 reqs_to_store: dict[ReqId, TransferSpec] = {}
273 # iterate over both new and cached requests
274 for req_id, new_block_id_groups, preempted in yield_req_data(scheduler_output):
275 + req_status = self._req_status[req_id]
276 + req_status.update_offload_keys()
277 +
278 if preempted:
199 - self._request_block_ids[req_id] = []
279 + for group_state in req_status.group_states:
280 + group_state.block_ids.clear()
281
282 if new_block_id_groups:
202 - new_block_ids = new_block_id_groups[0]
203 - self._request_block_ids[req_id] += new_block_ids
283 + req_status.update_block_id_groups(new_block_id_groups)
284 +
285 + # Below assertion will be removed once this function supports HMA
286 + assert len(req_status.group_states) == 1
287 + group_state = req_status.group_states[0]
288
205 - block_ids = self._request_block_ids[req_id]
289 + block_ids = group_state.block_ids
290
207 - req = self._requests[req_id]
291 + req = req_status.req
292 new_tokens = scheduler_output.num_scheduled_tokens[req_id]
293 expected_tokens = req.num_computed_tokens + new_tokens
294 # with async scheduling, some tokens may be missing
295 total_tokens = min(expected_tokens, req.num_tokens)
212 - num_blocks = total_tokens // self.offloaded_block_size
213 - start_block_idx = self._next_stored_block_idx.get(req_id, 0)
296 + num_blocks = total_tokens // group_config.offloaded_block_size
297 + start_block_idx = group_state.next_stored_block_idx
298 num_new_blocks = num_blocks - start_block_idx
299
300 if num_new_blocks <= 0:
301 continue
302
219 - num_gpu_blocks = num_blocks * self.block_size_factor
303 + num_gpu_blocks = num_blocks * self.config.block_size_factor
304 assert len(req.block_hashes) >= num_gpu_blocks
305
222 - new_block_hashes = self._get_block_hashes(
223 - req, start_idx=start_block_idx, end_idx=num_blocks
224 - )
225 - store_output = self.manager.prepare_store(new_block_hashes)
306 + new_offload_keys = group_state.offload_keys[start_block_idx:num_blocks]
307 + store_output = self.manager.prepare_store(new_offload_keys)
308 if store_output is None:
309 logger.warning(
310 "Request %s: cannot store %s blocks", req_id, num_new_blocks
311 )
312 continue
313
232 - self._next_stored_block_idx[req_id] = num_blocks
314 + group_state.next_stored_block_idx = num_blocks
315
234 - if not store_output.block_hashes_to_store:
316 + if not store_output.keys_to_store:
317 continue
236 - block_hashes_to_store = set(store_output.block_hashes_to_store)
318 + keys_to_store = set(store_output.keys_to_store)
319
238 - block_hashes = self._get_block_hashes(req, end_idx=num_blocks)
239 - self.manager.touch(block_hashes)
320 + self.manager.touch(group_state.offload_keys[:num_blocks])
321
241 - new_block_hashes = self._get_block_hashes(
242 - req, start_idx=start_block_idx, end_idx=num_blocks
243 - )
322 dst_spec = store_output.store_spec
323 src_block_ids: list[int] = []
246 - for idx, blk_hash in enumerate(new_block_hashes):
247 - if blk_hash not in block_hashes_to_store:
324 + for idx, key in enumerate(new_offload_keys):
325 + if key not in keys_to_store:
326 continue
327 offloaded_block_idx = start_block_idx + idx
250 - gpu_block_idx = offloaded_block_idx * self.block_size_factor
251 - for i in range(self.block_size_factor):
328 + gpu_block_idx = offloaded_block_idx * self.config.block_size_factor
329 + for i in range(self.config.block_size_factor):
330 src_block_ids.append(block_ids[gpu_block_idx + i])
331 src_spec = GPULoadStoreSpec(
332 src_block_ids, group_sizes=(len(src_block_ids),)
333 )
334
335 reqs_to_store[req_id] = (src_spec, dst_spec)
258 - self._reqs_being_stored[req_id] |= block_hashes_to_store
336 + self._reqs_being_stored[req_id] |= keys_to_store
337
338 logger.debug(
339 "Request %s offloading %s blocks starting from block #%d",
340 req_id,
263 - len(block_hashes_to_store),
341 + len(keys_to_store),
342 start_block_idx,
343 )
344
@@ -279,10 +357,10 @@ class OffloadingConnectorScheduler:
357 # NOTE (orozery): we should move this logic to update_connector_output
358 # once KVConnectorOutput allows us to report completed transfers
359 for req_id in scheduler_output.preempted_req_ids or ():
282 - block_hashes = self._reqs_being_stored.get(req_id)
283 - if block_hashes:
284 - self.manager.complete_store(block_hashes)
285 - block_hashes.clear()
360 + keys = self._reqs_being_stored.get(req_id)
361 + if keys:
362 + self.manager.complete_store(keys)
363 + keys.clear()
364
365 return meta
366
@@ -295,16 +373,16 @@ class OffloadingConnectorScheduler:
373 connectors output.
374 """
375 for req_id in connector_output.finished_sending or []:
298 - block_hashes = self._reqs_being_stored.pop(req_id, None)
299 - if block_hashes:
300 - self.manager.complete_store(block_hashes)
376 + keys = self._reqs_being_stored.pop(req_id, None)
377 + if keys:
378 + self.manager.complete_store(keys)
379
380 for req_id in connector_output.finished_recving or []:
303 - block_hashes = self._reqs_being_loaded.pop(req_id, None)
304 - if block_hashes:
381 + keys = self._reqs_being_loaded.pop(req_id, None)
382 + if keys:
383 if self._blocks_being_loaded:
306 - self._blocks_being_loaded.difference_update(block_hashes)
307 - self.manager.complete_load(block_hashes)
384 + self._blocks_being_loaded.difference_update(keys)
385 + self.manager.complete_load(keys)
386
387 def request_finished(
388 self,
@@ -322,12 +400,10 @@ class OffloadingConnectorScheduler:
400 returned by the engine.
401 """
402 req_id = request.request_id
325 - self._requests.pop(req_id, None)
326 - self._request_block_ids.pop(req_id, None)
403
404 # TODO(orozery): possibly kickoff offload for last block
405 # which may have been deferred due to async scheduling
330 - self._next_stored_block_idx.pop(req_id, None)
406 + self._req_status.pop(req_id, None)
407
408 request_being_stored = req_id in self._reqs_being_stored
409 return request_being_stored, None
@@ -339,11 +415,12 @@ class OffloadingConnectorScheduler:
415 A list of KV cache events.
416 """
417 for event in self.manager.take_events():
418 + block_hashes = [get_offload_block_hash(key) for key in event.keys]
419 if event.removed:
343 - yield BlockRemoved(block_hashes=event.block_hashes, medium=event.medium)
420 + yield BlockRemoved(block_hashes=block_hashes, medium=event.medium)
421 else:
422 yield BlockStored(
346 - block_hashes=event.block_hashes,
423 + block_hashes=block_hashes,
424 parent_block_hash=None,
425 token_ids=[],
426 lora_id=None,
vllm/v1/kv_offload/abstract.py
+35 -18
@@ -30,8 +30,27 @@ 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
34
34 -from vllm.v1.core.kv_cache_utils import BlockHash
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.
37 +# Use the helper functions below to construct / decompose keys.
38 +OffloadKey = NewType("OffloadKey", bytes)
39 +
40 +
41 +def make_offload_key(block_hash: bytes, group_idx: int) -> OffloadKey:
42 + """Pack a block hash and group index into an `OffloadKey`."""
43 + return OffloadKey(block_hash + group_idx.to_bytes(4, "big", signed=False))
44 +
45 +
46 +def get_offload_block_hash(key: OffloadKey) -> bytes:
47 + """Extract the block hash from an `OffloadKey`."""
48 + return key[:-4]
49 +
50 +
51 +def get_offload_group_idx(key: OffloadKey) -> int:
52 + """Extract the group index from an `OffloadKey`."""
53 + return int.from_bytes(key[-4:], "big", signed=False)
54
55
56 class LoadStoreSpec(ABC):
@@ -52,14 +71,14 @@ class LoadStoreSpec(ABC):
71
72 @dataclass
73 class PrepareStoreOutput:
55 - block_hashes_to_store: list[BlockHash]
74 + keys_to_store: list[OffloadKey]
75 store_spec: LoadStoreSpec
57 - block_hashes_evicted: list[BlockHash]
76 + evicted_keys: list[OffloadKey]
77
78
79 @dataclass
80 class OffloadingEvent:
62 - block_hashes: list[BlockHash]
81 + keys: list[OffloadKey]
82 block_size: int
83 medium: str
84 # True if blocks are removed, False if stored
@@ -68,13 +87,13 @@ class OffloadingEvent:
87
88 class OffloadingManager(ABC):
89 @abstractmethod
71 - def lookup(self, block_hashes: Iterable[BlockHash]) -> int | None:
90 + def lookup(self, keys: Iterable[OffloadKey]) -> int | None:
91 """
92 Finds the length of the maximal series of blocks, starting from the
93 first one, that are all offloaded.
94
95 Args:
77 - block_hashes: the hashes identifying the blocks to lookup.
96 + keys: the keys identifying the blocks to lookup.
97
98 Returns:
99 An integer representing the maximal number of blocks that
@@ -85,7 +104,7 @@ class OffloadingManager(ABC):
104 pass
105
106 @abstractmethod
88 - def prepare_load(self, block_hashes: Iterable[BlockHash]) -> LoadStoreSpec:
107 + def prepare_load(self, keys: Iterable[OffloadKey]) -> LoadStoreSpec:
108 """
109 Prepare the given blocks to be read.
110 The given blocks will be protected from eviction until
@@ -93,7 +112,7 @@ class OffloadingManager(ABC):
112 It assumes all given blocks are offloaded.
113
114 Args:
96 - block_hashes: the hashes identifying the blocks.
115 + keys: the keys identifying the blocks.
116
117 Returns:
118 A LoadStoreSpec that can be used by a worker to locate and load
@@ -101,36 +120,34 @@ class OffloadingManager(ABC):
120 """
121 pass
122
104 - def touch(self, block_hashes: Iterable[BlockHash]):
123 + def touch(self, keys: Iterable[OffloadKey]):
124 """
125 Mark the given blocks as recently used.
126 This could in practice mean moving them to the end of an LRU list.
127
128 Args:
110 - block_hashes: the hashes identifying the blocks.
129 + keys: the keys identifying the blocks.
130 """
131 return
132
114 - def complete_load(self, block_hashes: Iterable[BlockHash]):
133 + def complete_load(self, keys: Iterable[OffloadKey]):
134 """
135 Marks previous blocks that were prepared to load as done loading.
136
137 Args:
119 - block_hashes: the hashes identifying the blocks.
138 + keys: the keys identifying the blocks.
139 """
140 return
141
142 @abstractmethod
124 - def prepare_store(
125 - self, block_hashes: Iterable[BlockHash]
126 - ) -> PrepareStoreOutput | None:
143 + def prepare_store(self, keys: Iterable[OffloadKey]) -> PrepareStoreOutput | None:
144 """
145 Prepare the given blocks to be offloaded.
146 The given blocks will be protected from eviction until
147 complete_store is called.
148
149 Args:
133 - block_hashes: the hashes identifying the blocks.
150 + keys: the keys identifying the blocks.
151
152 Returns:
153 A PrepareStoreOutput indicating which blocks need storing,
@@ -140,7 +157,7 @@ class OffloadingManager(ABC):
157 """
158 pass
159
143 - def complete_store(self, block_hashes: Iterable[BlockHash], success: bool = True):
160 + def complete_store(self, keys: Iterable[OffloadKey], success: bool = True):
161 """
162 Marks blocks which were previously prepared to be stored, as stored.
163 Following this call, the blocks become loadable.
@@ -148,7 +165,7 @@ class OffloadingManager(ABC):
165 removed.
166
167 Args:
151 - block_hashes: the hashes identifying the blocks.
168 + keys: the keys identifying the blocks.
169 success: whether the blocks were stored successfully.
170 """
171 return
vllm/v1/kv_offload/cpu/manager.py
+50 -58
@@ -3,11 +3,11 @@
3 from collections.abc import Iterable
4 from typing import Literal
5
6 -from vllm.v1.core.kv_cache_utils import BlockHash
6 from vllm.v1.kv_offload.abstract import (
7 LoadStoreSpec,
8 OffloadingEvent,
9 OffloadingManager,
10 + OffloadKey,
11 PrepareStoreOutput,
12 )
13 from vllm.v1.kv_offload.cpu.policies.abstract import BlockStatus, CachePolicy
@@ -57,11 +57,9 @@ class CPUOffloadingManager(OffloadingManager):
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
60 + def _allocate_blocks(self, keys: list[OffloadKey]) -> list[BlockStatus]:
61 + num_fresh = min(len(keys), self._num_blocks - self._num_allocated_blocks)
62 + num_reused = len(keys) - num_fresh
63 assert len(self._free_list) >= num_reused
64
65 # allocate fresh blocks
@@ -80,122 +78,116 @@ class CPUOffloadingManager(OffloadingManager):
78
79 def _get_load_store_spec(
80 self,
83 - block_hashes: Iterable[BlockHash],
81 + keys: Iterable[OffloadKey],
82 blocks: Iterable[BlockStatus],
83 ) -> CPULoadStoreSpec:
84 return CPULoadStoreSpec([block.block_id for block in blocks])
85
86 # --- OffloadingManager interface ---
87
90 - def lookup(self, block_hashes: Iterable[BlockHash]) -> int | None:
88 + def lookup(self, keys: Iterable[OffloadKey]) -> int | None:
89 hit_count = 0
92 - for block_hash in block_hashes:
93 - block = self._policy.get(block_hash)
90 + for key in keys:
91 + block = self._policy.get(key)
92 if block is None or not block.is_ready:
93 break
94 hit_count += 1
95 return hit_count
96
99 - def prepare_load(self, block_hashes: Iterable[BlockHash]) -> LoadStoreSpec:
97 + def prepare_load(self, keys: Iterable[OffloadKey]) -> LoadStoreSpec:
98 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"
99 + for key in keys:
100 + block = self._policy.get(key)
101 + assert block is not None, f"Block {key!r} not found in cache"
102 + assert block.is_ready, f"Block {key!r} is not ready for reading"
103 block.ref_cnt += 1
104 blocks.append(block)
107 - return self._get_load_store_spec(block_hashes, blocks)
105 + return self._get_load_store_spec(keys, blocks)
106
109 - def touch(self, block_hashes: Iterable[BlockHash]) -> None:
110 - self._policy.touch(block_hashes)
107 + def touch(self, keys: Iterable[OffloadKey]) -> None:
108 + self._policy.touch(keys)
109
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"
110 + def complete_load(self, keys: Iterable[OffloadKey]) -> None:
111 + for key in keys:
112 + block = self._policy.get(key)
113 + assert block is not None, f"Block {key!r} not found"
114 + assert block.ref_cnt > 0, f"Block {key!r} ref_cnt is already 0"
115 block.ref_cnt -= 1
116
119 - def prepare_store(
120 - self, block_hashes: Iterable[BlockHash]
121 - ) -> PrepareStoreOutput | None:
122 - block_hashes_list = list(block_hashes)
117 + def prepare_store(self, keys: Iterable[OffloadKey]) -> PrepareStoreOutput | None:
118 + keys_list = list(keys)
119
120 # 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 - ]
121 + keys_to_store = [k for k in keys_list if self._policy.get(k) is None]
122
129 - if not block_hashes_to_store:
123 + if not keys_to_store:
124 return PrepareStoreOutput(
131 - block_hashes_to_store=[],
125 + keys_to_store=[],
126 store_spec=self._get_load_store_spec([], []),
133 - block_hashes_evicted=[],
127 + evicted_keys=[],
128 )
129
136 - num_blocks_to_evict = len(block_hashes_to_store) - self._get_num_free_blocks()
130 + num_blocks_to_evict = len(keys_to_store) - self._get_num_free_blocks()
131
138 - to_evict: list[BlockHash] = []
132 + to_evict: list[OffloadKey] = []
133 if num_blocks_to_evict > 0:
134 # Blocks from the original input are excluded from eviction candidates:
135 # a block that was already stored must remain in the cache after this call.
142 - protected = set(block_hashes_list)
136 + protected = set(keys_list)
137 evicted = self._policy.evict(num_blocks_to_evict, protected)
138 if evicted is None:
139 return None
146 - for block_hash, block in evicted:
140 + for key, block in evicted:
141 self._free_block(block)
148 - to_evict.append(block_hash)
142 + to_evict.append(key)
143
144 if to_evict and self.events is not None:
145 self.events.append(
146 OffloadingEvent(
153 - block_hashes=to_evict,
147 + keys=to_evict,
148 block_size=self.block_size,
149 medium=self.medium,
150 removed=True,
151 )
152 )
153
160 - blocks = self._allocate_blocks(block_hashes_to_store)
161 - assert len(blocks) == len(block_hashes_to_store), (
154 + blocks = self._allocate_blocks(keys_to_store)
155 + assert len(blocks) == len(keys_to_store), (
156 "Block pool did not allocate the expected number of blocks"
157 )
158
165 - for block_hash, block in zip(block_hashes_to_store, blocks):
166 - self._policy.insert(block_hash, block)
159 + for key, block in zip(keys_to_store, blocks):
160 + self._policy.insert(key, block)
161
162 # build store specs for allocated blocks
169 - store_spec = self._get_load_store_spec(block_hashes_to_store, blocks)
163 + store_spec = self._get_load_store_spec(keys_to_store, blocks)
164
165 return PrepareStoreOutput(
172 - block_hashes_to_store=block_hashes_to_store,
166 + keys_to_store=keys_to_store,
167 store_spec=store_spec,
174 - block_hashes_evicted=to_evict,
168 + evicted_keys=to_evict,
169 )
170
177 - def complete_store(
178 - self, block_hashes: Iterable[BlockHash], success: bool = True
179 - ) -> None:
180 - stored_block_hashes: list[BlockHash] = []
171 + def complete_store(self, keys: Iterable[OffloadKey], success: bool = True) -> None:
172 + stored_keys: list[OffloadKey] = []
173
174 if success:
183 - for block_hash in block_hashes:
184 - block = self._policy.get(block_hash)
175 + for key in keys:
176 + block = self._policy.get(key)
177 if block is not None and not block.is_ready:
178 block.ref_cnt = 0
187 - stored_block_hashes.append(block_hash)
179 + stored_keys.append(key)
180 else:
189 - for block_hash in block_hashes:
190 - block = self._policy.get(block_hash)
181 + for key in keys:
182 + block = self._policy.get(key)
183 if block is not None and not block.is_ready:
192 - self._policy.remove(block_hash)
184 + self._policy.remove(key)
185 self._free_block(block)
186
195 - if stored_block_hashes and self.events is not None:
187 + if stored_keys and self.events is not None:
188 self.events.append(
189 OffloadingEvent(
198 - block_hashes=stored_block_hashes,
190 + keys=stored_keys,
191 block_size=self.block_size,
192 medium=self.medium,
193 removed=False,
vllm/v1/kv_offload/cpu/policies/abstract.py
+8 -8
@@ -4,7 +4,7 @@ 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
7 +from vllm.v1.kv_offload.abstract import OffloadKey
8
9
10 class BlockStatus(ctypes.Structure):
@@ -45,29 +45,29 @@ class CachePolicy(ABC):
45 def __init__(self, cache_capacity: int) -> None: ...
46
47 @abstractmethod
48 - def get(self, block_hash: BlockHash) -> BlockStatus | None:
48 + def get(self, key: OffloadKey) -> 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:
52 + def insert(self, key: OffloadKey, 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:
56 + def remove(self, key: OffloadKey) -> 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:
60 + def touch(self, keys: Iterable[OffloadKey]) -> 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:
65 + self, n: int, protected: set[OffloadKey]
66 + ) -> list[tuple[OffloadKey, 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,
70 + Returns a list of (key, 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
vllm/v1/kv_offload/cpu/policies/arc.py
+53 -53
@@ -3,7 +3,7 @@
3 from collections import OrderedDict
4 from collections.abc import Iterable
5
6 -from vllm.v1.core.kv_cache_utils import BlockHash
6 +from vllm.v1.kv_offload.abstract import OffloadKey
7 from vllm.v1.kv_offload.cpu.policies.abstract import BlockStatus, CachePolicy
8
9
@@ -23,7 +23,7 @@ class ARCCachePolicy(CachePolicy):
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):
26 + For each key (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.
@@ -48,88 +48,88 @@ class ARCCachePolicy(CachePolicy):
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)
51 + self.t1: OrderedDict[OffloadKey, BlockStatus] = OrderedDict()
52 + self.t2: OrderedDict[OffloadKey, BlockStatus] = OrderedDict()
53 + # key -> None (only care about presence)
54 + self.b1: OrderedDict[OffloadKey, None] = OrderedDict()
55 + self.b2: OrderedDict[OffloadKey, None] = OrderedDict()
56 +
57 + def get(self, key: OffloadKey) -> BlockStatus | None:
58 + return self.t1.get(key) or self.t2.get(key)
59 +
60 + def insert(self, key: OffloadKey, block: BlockStatus) -> None:
61 + self.t1[key] = block
62 + self.b1.pop(key, None)
63 + self.b2.pop(key, None)
64 +
65 + def remove(self, key: OffloadKey) -> None:
66 + if self.t1.pop(key, None) is None:
67 + self.t2.pop(key, None)
68 +
69 + def touch(self, keys: Iterable[OffloadKey]) -> None:
70 + for key in reversed(list(keys)):
71 + if key in self.t1:
72 + block = self.t1.pop(key)
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
76 + self.t1[key] = block
77 else:
78 - self.t2[block_hash] = block
78 + self.t2[key] = block
79
80 - elif block_hash in self.t2:
81 - self.t2.move_to_end(block_hash)
80 + elif key in self.t2:
81 + self.t2.move_to_end(key)
82
83 - elif block_hash in self.b1:
83 + elif key 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)
89 + self.b1.move_to_end(key)
90
91 - elif block_hash in self.b2:
91 + elif key 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)
95 + self.b2.move_to_end(key)
96
97 def evict(
98 - self, n: int, protected: set[BlockHash]
99 - ) -> list[tuple[BlockHash, BlockStatus]] | None:
98 + self, n: int, protected: set[OffloadKey]
99 + ) -> list[tuple[OffloadKey, 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()
106 + tuple[OffloadKey, BlockStatus, bool]
107 + ] = [] # (key, block, from_t1)
108 + already_selected: set[OffloadKey] = set()
109 virtual_t1_size = len(self.t1)
110
111 for _ in range(n):
112 - candidate: tuple[BlockHash, BlockStatus, bool] | None = None
112 + candidate: tuple[OffloadKey, BlockStatus, bool] | None = None
113
114 if virtual_t1_size >= int(self.target_t1_size):
115 - for block_hash, block in self.t1.items():
115 + for key, 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
118 + and key not in protected
119 + and key not in already_selected
120 ):
121 - candidate = (block_hash, block, True)
121 + candidate = (key, block, True)
122 virtual_t1_size -= 1
123 break
124
125 if candidate is None:
126 - for block_hash, block in self.t2.items():
126 + for key, 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
129 + and key not in protected
130 + and key not in already_selected
131 ):
132 - candidate = (block_hash, block, False)
132 + candidate = (key, block, False)
133 break
134 if candidate is None:
135 return None
@@ -138,15 +138,15 @@ class ARCCachePolicy(CachePolicy):
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:
141 + result: list[tuple[OffloadKey, BlockStatus]] = []
142 + for key, block, from_t1 in candidates:
143 if from_t1:
144 - del self.t1[block_hash]
145 - self.b1[block_hash] = None
144 + del self.t1[key]
145 + self.b1[key] = None
146 else:
147 - del self.t2[block_hash]
148 - self.b2[block_hash] = None
149 - result.append((block_hash, block))
147 + del self.t2[key]
148 + self.b2[key] = None
149 + result.append((key, block))
150
151 # Trim ghost lists to cache_capacity.
152 for ghost in (self.b1, self.b2):
vllm/v1/kv_offload/cpu/policies/lru.py
+20 -20
@@ -3,7 +3,7 @@
3 from collections import OrderedDict
4 from collections.abc import Iterable
5
6 -from vllm.v1.core.kv_cache_utils import BlockHash
6 +from vllm.v1.kv_offload.abstract import OffloadKey
7 from vllm.v1.kv_offload.cpu.policies.abstract import BlockStatus, CachePolicy
8
9
@@ -12,35 +12,35 @@ class LRUCachePolicy(CachePolicy):
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()
15 + self.blocks: OrderedDict[OffloadKey, BlockStatus] = OrderedDict()
16
17 - def get(self, block_hash: BlockHash) -> BlockStatus | None:
18 - return self.blocks.get(block_hash)
17 + def get(self, key: OffloadKey) -> BlockStatus | None:
18 + return self.blocks.get(key)
19
20 - def insert(self, block_hash: BlockHash, block: BlockStatus) -> None:
21 - self.blocks[block_hash] = block
20 + def insert(self, key: OffloadKey, block: BlockStatus) -> None:
21 + self.blocks[key] = block
22
23 - def remove(self, block_hash: BlockHash) -> None:
24 - del self.blocks[block_hash]
23 + def remove(self, key: OffloadKey) -> None:
24 + del self.blocks[key]
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)
26 + def touch(self, keys: Iterable[OffloadKey]) -> None:
27 + for key in reversed(list(keys)):
28 + if key in self.blocks:
29 + self.blocks.move_to_end(key)
30
31 def evict(
32 - self, n: int, protected: set[BlockHash]
33 - ) -> list[tuple[BlockHash, BlockStatus]] | None:
32 + self, n: int, protected: set[OffloadKey]
33 + ) -> list[tuple[OffloadKey, 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))
36 + candidates: list[tuple[OffloadKey, BlockStatus]] = []
37 + for key, block in self.blocks.items():
38 + if block.ref_cnt == 0 and key not in protected:
39 + candidates.append((key, 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]
44 + for key, _ in candidates:
45 + del self.blocks[key]
46 return candidates
vllm/v1/kv_offload/reuse_manager.py
+25 -29
@@ -10,11 +10,11 @@ FilterReusedOffloadingManager — OffloadingManager decorator that skips
10 from collections import OrderedDict
11 from collections.abc import Iterable
12
13 -from vllm.v1.core.kv_cache_utils import BlockHash
13 from vllm.v1.kv_offload.abstract import (
14 LoadStoreSpec,
15 OffloadingEvent,
16 OffloadingManager,
17 + OffloadKey,
18 PrepareStoreOutput,
19 )
20
@@ -26,8 +26,8 @@ class FilterReusedOffloadingManager(OffloadingManager):
26 All methods are delegated to the *backing* manager. Two methods are
27 intercepted:
28
29 - * ``lookup`` — records each visited block hash in an internal LRU counter.
30 - * ``prepare_store`` — filters out block hashes that have not yet
29 + * ``lookup`` — records each visited key in an internal LRU counter.
30 + * ``prepare_store`` — filters out keys that have not yet
31 crossed the threshold *before* calling the backing
32 ``prepare_store``.
33
@@ -59,61 +59,57 @@ class FilterReusedOffloadingManager(OffloadingManager):
59 self.store_threshold = store_threshold
60 self.max_tracker_size = max_tracker_size
61 # Ordered so we can evict the LRU entry in O(1).
62 - self.counts: OrderedDict[BlockHash, int] = OrderedDict()
62 + self.counts: OrderedDict[OffloadKey, int] = OrderedDict()
63
64 # ------------------------------------------------------------------
65 # Intercepted methods
66 # ------------------------------------------------------------------
67
68 - def lookup(self, block_hashes: Iterable[BlockHash]) -> int | None:
69 - """Record each hash, then delegate lookup to backing manager."""
70 - block_hashes = list(block_hashes)
71 - for block_hash in block_hashes:
72 - if block_hash in self.counts:
73 - self.counts.move_to_end(block_hash)
74 - self.counts[block_hash] += 1
68 + def lookup(self, keys: Iterable[OffloadKey]) -> int | None:
69 + """Record each key, then delegate lookup to backing manager."""
70 + keys = list(keys)
71 + for key in keys:
72 + if key in self.counts:
73 + self.counts.move_to_end(key)
74 + self.counts[key] += 1
75 else:
76 if len(self.counts) >= self.max_tracker_size:
77 self.counts.popitem(last=False) # evict LRU
78 - self.counts[block_hash] = 1
79 - return self._backing.lookup(block_hashes)
78 + self.counts[key] = 1
79 + return self._backing.lookup(keys)
80
81 - def prepare_store(
82 - self, block_hashes: Iterable[BlockHash]
83 - ) -> PrepareStoreOutput | None:
81 + def prepare_store(self, keys: Iterable[OffloadKey]) -> PrepareStoreOutput | None:
82 """Filter out blocks below threshold, then delegate to backing.
83
84 Filtering is evaluated *before* calling the backing manager's
85 ``prepare_store`` so that blocks that would be skipped do not
86 consume any CPU offload capacity.
87 """
90 - block_hashes = list(block_hashes)
88 + keys = list(keys)
89 eligible = [
92 - bh for bh in block_hashes if self.counts.get(bh, 0) >= self.store_threshold
90 + key for key in keys if self.counts.get(key, 0) >= self.store_threshold
91 ]
92
95 - # Delegate to the backing manager with only the eligible hashes.
93 # Passing an empty list is intentional and safe — CPUOffloadingManager
94 # handles it correctly, returning a PrepareStoreOutput with empty lists.
95 + # Delegate to the backing manager with only the eligible keys.
96 return self._backing.prepare_store(eligible)
97
98 # ------------------------------------------------------------------
99 # Delegated methods
100 # ------------------------------------------------------------------
101
104 - def prepare_load(self, block_hashes: Iterable[BlockHash]) -> LoadStoreSpec:
105 - return self._backing.prepare_load(block_hashes)
102 + def prepare_load(self, keys: Iterable[OffloadKey]) -> LoadStoreSpec:
103 + return self._backing.prepare_load(keys)
104
107 - def touch(self, block_hashes: Iterable[BlockHash]) -> None:
108 - return self._backing.touch(block_hashes)
105 + def touch(self, keys: Iterable[OffloadKey]) -> None:
106 + return self._backing.touch(keys)
107
110 - def complete_load(self, block_hashes: Iterable[BlockHash]) -> None:
111 - return self._backing.complete_load(block_hashes)
108 + def complete_load(self, keys: Iterable[OffloadKey]) -> None:
109 + return self._backing.complete_load(keys)
110
113 - def complete_store(
114 - self, block_hashes: Iterable[BlockHash], success: bool = True
115 - ) -> None:
116 - return self._backing.complete_store(block_hashes, success)
111 + def complete_store(self, keys: Iterable[OffloadKey], success: bool = True) -> None:
112 + return self._backing.complete_store(keys, success)
113
114 def take_events(self) -> Iterable[OffloadingEvent]:
115 return self._backing.take_events()