@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / c59a0eca42
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 14, 05:30 UTC

[KV offload][4/N] Offloading KV connector (#22595)

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

Or Ozeri committed Sep 19, 2025 at 22:07 UTC c59a0eca4204522ecaa28af1f1e38d50b1e2626f
6 files changed +1111 -1
tests/v1/kv_connector/unit/test_offloading_connector.py new
+505
@@ -0,0 +1,505 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import copy
4 +from collections.abc import Iterable, Iterator
5 +from dataclasses import dataclass
6 +from typing import Any
7 +from unittest.mock import MagicMock
8 +
9 +import pytest
10 +import torch
11 +
12 +from vllm import SamplingParams
13 +from vllm.config import KVTransferConfig, VllmConfig
14 +from vllm.distributed.kv_events import BlockRemoved, BlockStored
15 +from vllm.distributed.kv_transfer.kv_connector.v1 import KVConnectorRole
16 +from vllm.distributed.kv_transfer.kv_connector.v1.offloading_connector import (
17 + OffloadingConnector, OffloadingConnectorMetadata)
18 +from vllm.forward_context import ForwardContext
19 +from vllm.utils import sha256
20 +from vllm.v1.core.kv_cache_utils import (BlockHash, get_request_block_hasher,
21 + init_none_hash)
22 +from vllm.v1.core.sched.scheduler import Scheduler
23 +from vllm.v1.kv_offload.abstract import (LoadStoreSpec, OffloadingEvent,
24 + OffloadingManager, PrepareStoreOutput)
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 (OffloadingHandler,
28 + TransferResult, TransferSpec)
29 +from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT, KVConnectorOutput
30 +from vllm.v1.request import Request
31 +
32 +from .utils import (EOS_TOKEN_ID, create_model_runner_output, create_scheduler,
33 + create_vllm_config)
34 +
35 +
36 +class MockLoadStoreSpec(LoadStoreSpec):
37 +
38 + def __init__(self, block_hashes: Iterable[BlockHash]):
39 + self.block_hashes: list[BlockHash] = list(block_hashes)
40 +
41 + @staticmethod
42 + def medium() -> str:
43 + return "Mock"
44 +
45 + def __repr__(self) -> str:
46 + return repr(self.block_hashes)
47 +
48 +
49 +class MockOffloadingHandler(OffloadingHandler):
50 +
51 + def __init__(self):
52 + self.completed_transfers: list[TransferResult] = []
53 + self.completed_specs: list[TransferSpec] = []
54 +
55 + def get_finished(self) -> list[TransferResult]:
56 + finished = self.completed_transfers
57 + self.completed_transfers = []
58 + return finished
59 +
60 + def transfer_async(self, job_id: int, spec: TransferSpec) -> bool:
61 + self.completed_specs.append(spec)
62 + self.completed_transfers.append((job_id, True))
63 + return True
64 +
65 +
66 +class MockOffloadingSpec(OffloadingSpec):
67 +
68 + def __init__(self, vllm_config: VllmConfig):
69 + super().__init__(vllm_config)
70 +
71 + self.manager = MagicMock(spec=OffloadingManager)
72 + self.manager.lookup.return_value = 0
73 + self.manager.prepare_load = lambda block_hashes: (MockLoadStoreSpec(
74 + block_hashes))
75 + self.handler = MockOffloadingHandler()
76 +
77 + def get_manager(self) -> OffloadingManager:
78 + return self.manager
79 +
80 + def get_handlers(
81 + self, _
82 + ) -> Iterator[tuple[type[LoadStoreSpec], type[LoadStoreSpec],
83 + OffloadingHandler]]:
84 +
85 + yield GPULoadStoreSpec, MockLoadStoreSpec, self.handler
86 + yield MockLoadStoreSpec, GPULoadStoreSpec, self.handler
87 +
88 + def get_completed_transfers(self) -> list[TransferSpec]:
89 + specs = self.handler.completed_specs
90 + self.handler.completed_specs = []
91 + return specs
92 +
93 +
94 +@dataclass
95 +class TransferSummary:
96 + gpu_block_indices: list[int]
97 + offload_addresses: list[Any]
98 +
99 +
100 +class RequestRunner:
101 +
102 + def __init__(self, offloaded_block_size: int, gpu_block_size: int,
103 + num_gpu_blocks: int):
104 + self.offloaded_block_size: int = offloaded_block_size
105 + self.gpu_block_size: int = gpu_block_size
106 + self.num_gpu_blocks: int = num_gpu_blocks
107 +
108 + self.req_id: int = -1
109 +
110 + vllm_config = create_vllm_config(block_size=gpu_block_size,
111 + max_num_batched_tokens=1000)
112 + vllm_config.kv_transfer_config = KVTransferConfig(
113 + kv_connector="OffloadingConnector",
114 + kv_role="kv_both",
115 + kv_connector_extra_config={
116 + "spec_name": "MockOffloadingSpec",
117 + "spec_module_path":
118 + "tests.v1.kv_connector.unit.test_offloading_connector",
119 + "block_size": offloaded_block_size,
120 + })
121 +
122 + self.scheduler: Scheduler = create_scheduler(vllm_config,
123 + num_blocks=num_gpu_blocks)
124 + self.worker_connector = OffloadingConnector(vllm_config,
125 + KVConnectorRole.WORKER)
126 +
127 + # register worker kv_caches to enable OffloadingWorker creations
128 + self.worker_connector.register_kv_caches(
129 + kv_caches={"a": torch.empty(0)})
130 +
131 + # extract connector of scheduler
132 + scheduler_connector = self.scheduler.connector
133 + assert scheduler_connector is not None
134 + assert isinstance(scheduler_connector, OffloadingConnector)
135 + self.scheduler_connector: OffloadingConnector = scheduler_connector
136 +
137 + # extract mocked OffloadingManager of scheduler connector
138 + connector_scheduler = scheduler_connector.connector_scheduler
139 + assert connector_scheduler is not None
140 + manager = connector_scheduler.manager
141 + assert isinstance(manager, MagicMock)
142 + self.manager: MagicMock = manager
143 +
144 + assert connector_scheduler.gpu_block_size == gpu_block_size
145 + assert connector_scheduler.offloaded_block_size == offloaded_block_size
146 +
147 + # extract OffloadingSpec of worker_connector
148 + connector_worker = self.worker_connector.connector_worker
149 + assert connector_worker is not None
150 + offloading_spec = connector_worker.spec
151 + assert isinstance(offloading_spec, MockOffloadingSpec)
152 + self.offloading_spec: MockOffloadingSpec = offloading_spec
153 +
154 + # mapping (offloading address) -> gpu_block_index
155 + self.offloaded: dict[Any, int] = {}
156 +
157 + self.pending_loads_count: int = 0
158 + self.pending_stores_count: int = 0
159 +
160 + self.completed_loads: list[TransferSummary] = []
161 + self.completed_stores: list[TransferSummary] = []
162 +
163 + # maps {block_id: block_offset}
164 + self.gpu_block_index: dict[int, int] = {}
165 +
166 + init_none_hash(sha256)
167 + self._block_hasher = get_request_block_hasher(gpu_block_size, sha256)
168 +
169 + self._dummy_ctx: ForwardContext = ForwardContext(no_compile_layers={},
170 + attn_metadata={},
171 + virtual_engine=0)
172 +
173 + def new_request(self, token_ids: list[int]):
174 + assert not self.scheduler.requests
175 + self.req_id += 1
176 +
177 + req = Request(
178 + request_id=str(self.req_id),
179 + prompt_token_ids=token_ids,
180 + sampling_params=SamplingParams(max_tokens=1000),
181 + pooling_params=None,
182 + eos_token_id=EOS_TOKEN_ID,
183 + block_hasher=self._block_hasher,
184 + )
185 +
186 + self.scheduler.add_request(req)
187 +
188 + def _wait_for_transfers(self):
189 + block_size_factor = self.offloaded_block_size // self.gpu_block_size
190 +
191 + while self.pending_loads_count or self.pending_stores_count:
192 + for transfer_spec in (
193 + self.offloading_spec.get_completed_transfers()):
194 + src_spec, dst_spec = transfer_spec
195 +
196 + if isinstance(src_spec, GPULoadStoreSpec):
197 + store = True
198 + gpu_spec = src_spec
199 + offload_spec = dst_spec
200 + else:
201 + store = False
202 + gpu_spec = dst_spec
203 + offload_spec = src_spec
204 +
205 + assert isinstance(offload_spec, MockLoadStoreSpec)
206 + assert isinstance(gpu_spec, GPULoadStoreSpec)
207 +
208 + gpu_block_indices: list[int] = []
209 + for block_id in gpu_spec.block_ids:
210 + gpu_block_indices.append(
211 + self.gpu_block_index[block_id.item()])
212 +
213 + # list of (block_hash, sub_block_offset)
214 + offload_addresses: list[Any] = []
215 + for block_hash in offload_spec.block_hashes:
216 + for sub_block_idx in range(block_size_factor):
217 + offload_addresses.append((block_hash, sub_block_idx))
218 +
219 + if store:
220 + assert len(gpu_block_indices) == len(offload_addresses)
221 +
222 + self.completed_stores.append(
223 + TransferSummary(gpu_block_indices, offload_addresses))
224 + self.pending_stores_count -= 1
225 + else:
226 + remainder_sub_block_count = (len(offload_addresses) -
227 + len(gpu_block_indices))
228 + assert remainder_sub_block_count >= 0
229 + assert remainder_sub_block_count < block_size_factor
230 + offload_addresses = offload_addresses[
231 + remainder_sub_block_count:]
232 +
233 + self.completed_loads.append(
234 + TransferSummary(gpu_block_indices, offload_addresses))
235 + self.pending_loads_count -= 1
236 +
237 + def _update_gpu_block_idx(self):
238 + for blocks in (self.scheduler.kv_cache_manager.coordinator.
239 + single_type_managers[0].req_to_blocks.values()):
240 + for block_idx, block in enumerate(blocks):
241 + self.gpu_block_index[block.block_id] = block_idx
242 +
243 + def _run(self, decoded_tokens: list[int]):
244 + """
245 + Runs multiple engine (scheduler + worker) steps.
246 + Assumes a single request is running.
247 +
248 + Args:
249 + decoded_tokens: the tokens to yield at each step.
250 + """
251 +
252 + tokens_iter = iter(decoded_tokens)
253 + token_id = next(tokens_iter, None)
254 + while token_id is not None:
255 + assert self.scheduler.requests
256 +
257 + scheduler_output = self.scheduler.schedule()
258 + self._update_gpu_block_idx()
259 +
260 + kv_connector_metadata = scheduler_output.kv_connector_metadata
261 + assert kv_connector_metadata is not None
262 + assert isinstance(kv_connector_metadata,
263 + OffloadingConnectorMetadata)
264 +
265 + self.pending_loads_count += len(kv_connector_metadata.reqs_to_load)
266 + self.pending_stores_count += len(
267 + kv_connector_metadata.reqs_to_store)
268 +
269 + self.worker_connector.bind_connector_metadata(
270 + kv_connector_metadata)
271 + self.worker_connector.start_load_kv(self._dummy_ctx)
272 +
273 + if scheduler_output.total_num_scheduled_tokens > 0:
274 + self.worker_connector.wait_for_save()
275 +
276 + finished_sending, finished_recving = (
277 + self.worker_connector.get_finished(
278 + scheduler_output.finished_req_ids))
279 +
280 + self.worker_connector.clear_connector_metadata()
281 +
282 + model_runner_output = create_model_runner_output(
283 + reqs=self.scheduler.running,
284 + finished_sending=list(finished_sending),
285 + finished_recving=list(finished_recving),
286 + token_id=token_id)
287 +
288 + if self.scheduler.running:
289 + token_id = next(tokens_iter, None)
290 +
291 + self.scheduler.update_from_output(scheduler_output,
292 + model_runner_output)
293 +
294 + self._wait_for_transfers()
295 +
296 + # run one more step to update finished stored
297 + if EOS_TOKEN_ID in decoded_tokens:
298 + assert not self.scheduler.running
299 +
300 + while self.scheduler.requests:
301 + scheduler_output = self.scheduler.schedule()
302 +
303 + finished_sending, finished_recving = (
304 + self.worker_connector.get_finished(
305 + scheduler_output.finished_req_ids))
306 +
307 + assert not finished_recving
308 +
309 + model_runner_output = copy.deepcopy(EMPTY_MODEL_RUNNER_OUTPUT)
310 + model_runner_output.kv_connector_output = KVConnectorOutput(
311 + finished_sending=finished_sending)
312 +
313 + self.scheduler.update_from_output(scheduler_output,
314 + model_runner_output)
315 +
316 + def run(
317 + self,
318 + decoded_tokens: list[int],
319 + expected_stored_gpu_block_indexes: tuple[int, ...] = (),
320 + expected_loaded_gpu_block_indexes: tuple[int, ...] = (),
321 + ):
322 + """
323 + Runs multiple engine (scheduler + worker) steps.
324 + Assumes a single request is running.
325 +
326 + Args:
327 + decoded_tokens: the tokens to yield at each step.
328 + expected_stored_gpu_block_indexes: GPU block indexes
329 + that are expected to be written during the run.
330 + expected_loaded_gpu_block_indexes: GPU block indexes
331 + that are expected to be loaded during the run.
332 + """
333 +
334 + self.manager.reset_mock()
335 + self._run(decoded_tokens)
336 +
337 + loaded_gpu_block_indexes: set[int] = set()
338 + for transfer in self.completed_loads:
339 + for gpu_block_idx, offloaded_address in zip(
340 + transfer.gpu_block_indices, transfer.offload_addresses):
341 + loaded_gpu_block_indexes.add(gpu_block_idx)
342 + assert gpu_block_idx == self.offloaded[offloaded_address]
343 +
344 + assert (
345 + set(expected_loaded_gpu_block_indexes) == loaded_gpu_block_indexes)
346 + self.completed_loads.clear()
347 +
348 + stored_gpu_block_indexes: set[int] = set()
349 + for transfer in self.completed_stores:
350 + for gpu_block_idx, offloaded_address in zip(
351 + transfer.gpu_block_indices, transfer.offload_addresses):
352 + stored_gpu_block_indexes.add(gpu_block_idx)
353 + self.offloaded[offloaded_address] = gpu_block_idx
354 +
355 + assert (
356 + set(expected_stored_gpu_block_indexes) == stored_gpu_block_indexes)
357 + self.completed_stores.clear()
358 +
359 +
360 +@pytest.fixture
361 +def request_runner():
362 + runners = []
363 +
364 + def runner_factory(offloaded_block_size, gpu_block_size, num_gpu_blocks):
365 + runner = RequestRunner(offloaded_block_size=offloaded_block_size,
366 + gpu_block_size=gpu_block_size,
367 + num_gpu_blocks=num_gpu_blocks)
368 + runners.append(runner)
369 + return runner
370 +
371 + yield runner_factory # pass factory to the test
372 +
373 +
374 +def generate_store_output(block_hashes: Iterable[BlockHash]):
375 + block_hashes = list(block_hashes)
376 + return PrepareStoreOutput(
377 + block_hashes_to_store=list(block_hashes),
378 + store_spec=MockLoadStoreSpec(block_hashes),
379 + block_hashes_evicted=[],
380 + )
381 +
382 +
383 +def test_offloading_connector(request_runner):
384 + offloaded_block_size = 12
385 + gpu_block_size = 4
386 + num_gpu_blocks = 100
387 + block_size_factor = offloaded_block_size // gpu_block_size
388 +
389 + runner = request_runner(offloaded_block_size=offloaded_block_size,
390 + gpu_block_size=gpu_block_size,
391 + num_gpu_blocks=num_gpu_blocks)
392 +
393 + # 3 blocks, store just the middle block (skip first and last)
394 + # blocks = [0, 1, 2], [3, 4, 5], [6, 7, 8]
395 + runner.new_request(token_ids=[0] * offloaded_block_size * 3)
396 + runner.manager.prepare_store.side_effect = \
397 + lambda block_hashes: generate_store_output(list(block_hashes)[1:2])
398 + runner.run(decoded_tokens=[0], expected_stored_gpu_block_indexes=(3, 4, 5))
399 +
400 + # add block missing 1 token -> no offload
401 + runner.run(decoded_tokens=[0] * (offloaded_block_size - 1))
402 + runner.manager.prepare_store.assert_not_called()
403 +
404 + # +1 token -> single block, fail prepare_store
405 + runner.manager.prepare_store.side_effect = \
406 + lambda block_hashes: None
407 + runner.run(decoded_tokens=[0])
408 + runner.manager.prepare_store.assert_called()
409 +
410 + # 1 more block, now set block_hashes_to_store = []
411 + runner.manager.prepare_store.side_effect = \
412 + lambda block_hashes: generate_store_output([])
413 + runner.run(decoded_tokens=[0] * offloaded_block_size)
414 +
415 + # 1 more block, now check touch was called with all 6 blocks
416 + runner.manager.prepare_store.side_effect = \
417 + lambda block_hashes: generate_store_output(block_hashes)
418 + runner.run(decoded_tokens=[0] * offloaded_block_size,
419 + expected_stored_gpu_block_indexes=(15, 16, 17))
420 + runner.manager.touch.assert_called()
421 + block_hashes1 = list(runner.manager.touch.call_args.args[0])
422 + assert len(block_hashes1) == 6
423 +
424 + # terminate request
425 + runner.run(decoded_tokens=[EOS_TOKEN_ID])
426 +
427 + # create a new request differing only on the last token
428 + runner.new_request(token_ids=[0] * (offloaded_block_size * 6 - 1) + [1])
429 + runner.run(decoded_tokens=[0],
430 + expected_stored_gpu_block_indexes=tuple(
431 + range(6 * block_size_factor)))
432 + runner.manager.touch.assert_called()
433 + block_hashes2 = list(runner.manager.touch.call_args.args[0])
434 + assert len(block_hashes2) == 6
435 +
436 + # verify hashes are the same, except for the last block
437 + assert block_hashes1[:5] == block_hashes2[:5]
438 + assert block_hashes1[5] != block_hashes2[5]
439 +
440 + # terminate request
441 + runner.run(decoded_tokens=[EOS_TOKEN_ID])
442 +
443 + # full_block_tokens - num_computed_tokens < offloaded_block_size
444 + runner.new_request(token_ids=[0] * gpu_block_size + [1] *
445 + (offloaded_block_size - gpu_block_size))
446 + runner.manager.prepare_store.side_effect = \
447 + lambda block_hashes: generate_store_output([])
448 + runner.run(decoded_tokens=[EOS_TOKEN_ID])
449 + runner.manager.lookup.assert_not_called()
450 +
451 + # single block lookup with no hits
452 + runner.new_request(token_ids=[1] * offloaded_block_size)
453 + runner.manager.prepare_store.side_effect = \
454 + lambda block_hashes: generate_store_output([])
455 + runner.run(decoded_tokens=[EOS_TOKEN_ID])
456 + runner.manager.lookup.assert_called()
457 + assert len(list(runner.manager.lookup.call_args.args[0])) == 1
458 +
459 + # single block lookup with a hit
460 + runner.scheduler.reset_prefix_cache()
461 + runner.new_request(token_ids=[0] * offloaded_block_size)
462 + runner.manager.prepare_store.side_effect = \
463 + lambda block_hashes: generate_store_output([])
464 + runner.manager.lookup.return_value = 1
465 + runner.run(decoded_tokens=[EOS_TOKEN_ID],
466 + expected_loaded_gpu_block_indexes=(0, 1, 2))
467 +
468 + # single block lookup with a hit in a middle block
469 + runner.new_request(token_ids=[0] * offloaded_block_size * 2 +
470 + [1] * offloaded_block_size)
471 + runner.manager.prepare_store.side_effect = \
472 + lambda block_hashes: generate_store_output([])
473 + runner.manager.lookup.return_value = 1
474 + runner.run(decoded_tokens=[EOS_TOKEN_ID],
475 + expected_loaded_gpu_block_indexes=(3, 4, 5))
476 +
477 + # test take_events
478 + def to_hashes(int_hashes: list[int]) -> list[BlockHash]:
479 + return [BlockHash(str(i).encode()) for i in int_hashes]
480 +
481 + def take_events() -> Iterable[OffloadingEvent]:
482 + yield OffloadingEvent(block_hashes=to_hashes([1, 2, 3]),
483 + block_size=16,
484 + medium="A",
485 + removed=False)
486 + yield OffloadingEvent(block_hashes=to_hashes([4, 5, 6]),
487 + block_size=32,
488 + medium="B",
489 + removed=True)
490 +
491 + runner.manager.take_events.side_effect = take_events
492 + events = list(runner.scheduler_connector.take_events())
493 + assert len(events) == 2
494 + event = events[0]
495 + assert isinstance(event, BlockStored)
496 + assert event.block_hashes == to_hashes([1, 2, 3])
497 + assert event.block_size == 16
498 + assert event.medium == "A"
499 + assert event.token_ids == []
500 + assert event.parent_block_hash is None
501 + assert event.lora_id is None
502 + event = events[1]
503 + assert isinstance(event, BlockRemoved)
504 + assert event.block_hashes == to_hashes([4, 5, 6])
505 + assert event.medium == "B"
tests/v1/kv_connector/unit/utils.py
+2 -1
@@ -176,6 +176,7 @@ def create_model_runner_output(
176 finished_sending: Optional[list[str]] = None,
177 finished_recving: Optional[list[str]] = None,
178 use_eos: bool = False,
179 + token_id: int = 0,
180 ) -> ModelRunnerOutput:
181 """Make dummy model runner output for testing."""
182
@@ -184,7 +185,7 @@ def create_model_runner_output(
185 req_id_to_index = {req_id: idx for idx, req_id in enumerate(req_ids)}
186
187 # Make sampled tokens.
187 - sampled_token = EOS_TOKEN_ID if use_eos else 0
188 + sampled_token = EOS_TOKEN_ID if use_eos else token_id
189 sampled_token_ids = [[sampled_token] for _ in req_ids]
190
191 kv_connector_output = None if (
vllm/distributed/kv_transfer/kv_connector/factory.py
+5
@@ -106,3 +106,8 @@ KVConnectorFactory.register_connector(
106 "MultiConnector",
107 "vllm.distributed.kv_transfer.kv_connector.v1.multi_connector",
108 "MultiConnector")
109 +
110 +KVConnectorFactory.register_connector(
111 + "OffloadingConnector",
112 + "vllm.distributed.kv_transfer.kv_connector.v1.offloading_connector",
113 + "OffloadingConnector")
vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py new
+485
@@ -0,0 +1,485 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from collections import defaultdict
4 +from collections.abc import Iterable, Iterator
5 +from dataclasses import dataclass
6 +from itertools import islice
7 +from typing import Any, Optional
8 +
9 +import torch
10 +
11 +from vllm.attention import AttentionMetadata
12 +from vllm.config import VllmConfig
13 +from vllm.distributed.kv_events import BlockRemoved, BlockStored, KVCacheEvent
14 +from vllm.distributed.kv_transfer.kv_connector.v1 import (KVConnectorBase_V1,
15 + KVConnectorRole)
16 +from vllm.distributed.kv_transfer.kv_connector.v1.base import (
17 + KVConnectorMetadata)
18 +from vllm.forward_context import ForwardContext
19 +from vllm.logger import init_logger
20 +from vllm.v1.core.kv_cache_manager import KVCacheBlocks
21 +from vllm.v1.core.kv_cache_utils import BlockHash
22 +from vllm.v1.core.sched.output import SchedulerOutput
23 +from vllm.v1.kv_offload.abstract import OffloadingManager
24 +from vllm.v1.kv_offload.factory import OffloadingSpecFactory
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 OffloadingWorker, TransferSpec
28 +from vllm.v1.outputs import KVConnectorOutput
29 +from vllm.v1.request import Request
30 +
31 +ReqId = str
32 +
33 +logger = init_logger(__name__)
34 +
35 +
36 +@dataclass
37 +class OffloadingConnectorMetadata(KVConnectorMetadata):
38 + reqs_to_load: dict[ReqId, TransferSpec]
39 + reqs_to_store: dict[ReqId, TransferSpec]
40 +
41 +
42 +class OffloadingConnector(KVConnectorBase_V1):
43 +
44 + def __init__(self, vllm_config: VllmConfig, role: KVConnectorRole):
45 + super().__init__(vllm_config, role)
46 +
47 + spec = OffloadingSpecFactory.create_spec(vllm_config)
48 +
49 + self.connector_scheduler: Optional[OffloadingConnectorScheduler] = None
50 + self.connector_worker: Optional[OffloadingConnectorWorker] = None
51 + if role == KVConnectorRole.SCHEDULER:
52 + self.connector_scheduler = OffloadingConnectorScheduler(spec)
53 + elif role == KVConnectorRole.WORKER:
54 + self.connector_worker = OffloadingConnectorWorker(spec)
55 +
56 + def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
57 + assert self.connector_worker is not None
58 + self.connector_worker.register_kv_caches(kv_caches)
59 +
60 + def start_load_kv(self, forward_context: "ForwardContext",
61 + **kwargs) -> None:
62 + assert self.connector_worker is not None
63 + assert isinstance(self._connector_metadata,
64 + OffloadingConnectorMetadata)
65 + self.connector_worker.start_load_kv(self._connector_metadata)
66 +
67 + def wait_for_layer_load(self, layer_name: str) -> None:
68 + pass
69 +
70 + def save_kv_layer(self, layer_name: str, kv_layer: torch.Tensor,
71 + attn_metadata: "AttentionMetadata", **kwargs) -> None:
72 + pass
73 +
74 + def wait_for_save(self):
75 + assert self.connector_worker is not None
76 + assert isinstance(self._connector_metadata,
77 + OffloadingConnectorMetadata)
78 + self.connector_worker.start_store_kv(self._connector_metadata)
79 +
80 + def get_finished(self,
81 + finished_req_ids: set[str]) -> tuple[set[str], set[str]]:
82 + assert self.connector_worker is not None
83 + return self.connector_worker.get_finished(finished_req_ids)
84 +
85 + def get_num_new_matched_tokens(
86 + self, request: "Request",
87 + num_computed_tokens: int) -> tuple[int, bool]:
88 + assert self.connector_scheduler is not None
89 + return self.connector_scheduler.get_num_new_matched_tokens(
90 + request, num_computed_tokens)
91 +
92 + def update_state_after_alloc(self, request: "Request",
93 + blocks: "KVCacheBlocks",
94 + num_external_tokens: int):
95 + assert self.connector_scheduler is not None
96 + return self.connector_scheduler.update_state_after_alloc(
97 + request, blocks, num_external_tokens)
98 +
99 + def build_connector_meta(
100 + self, scheduler_output: SchedulerOutput) -> KVConnectorMetadata:
101 + assert self.connector_scheduler is not None
102 + return self.connector_scheduler.build_connector_meta(scheduler_output)
103 +
104 + def update_connector_output(self, connector_output: KVConnectorOutput):
105 + assert self.connector_scheduler is not None
106 + self.connector_scheduler.update_connector_output(connector_output)
107 +
108 + def request_finished(
109 + self,
110 + request: "Request",
111 + block_ids: list[int],
112 + ) -> tuple[bool, Optional[dict[str, Any]]]:
113 + assert self.connector_scheduler is not None
114 + return self.connector_scheduler.request_finished(request, block_ids)
115 +
116 + def take_events(self) -> Iterable[KVCacheEvent]:
117 + assert self.connector_scheduler is not None
118 + return self.connector_scheduler.take_events()
119 +
120 +
121 +class OffloadingConnectorScheduler:
122 + """Implementation of Scheduler side methods"""
123 +
124 + def __init__(self, spec: OffloadingSpec):
125 + self.gpu_block_size = spec.gpu_block_size
126 + self.offloaded_block_size = spec.offloaded_block_size
127 + self.block_size_factor = (self.offloaded_block_size //
128 + self.gpu_block_size)
129 + self.manager: OffloadingManager = spec.get_manager()
130 +
131 + self._requests: dict[ReqId, Request] = {}
132 + # list of GPU block IDs per request
133 + self._request_block_ids: dict[ReqId, list[int]] = {}
134 + # requests to load for the current scheduler step
135 + self._reqs_to_load: dict[ReqId, TransferSpec] = {}
136 + # request blocks are stored in order
137 + # index of next block (of size offloaded_block_size) to offload
138 + self._next_stored_block_idx: dict[ReqId, int] = {}
139 +
140 + # request ID -> set(block hashes being stored/load)
141 + self._reqs_being_stored = defaultdict[ReqId, set[BlockHash]](set)
142 + self._reqs_being_loaded = defaultdict[ReqId, set[BlockHash]](set)
143 +
144 + def _get_block_hashes(
145 + self,
146 + req: Request,
147 + start_idx: int = 0,
148 + end_idx: Optional[int] = None,
149 + ) -> Iterable[BlockHash]:
150 + return islice(
151 + req.block_hashes,
152 + self.block_size_factor * start_idx + self.block_size_factor - 1,
153 + self.block_size_factor * end_idx if end_idx else None,
154 + self.block_size_factor)
155 +
156 + def get_num_new_matched_tokens(
157 + self, request: Request,
158 + num_computed_tokens: int) -> tuple[int, bool]:
159 + """
160 + Get number of new tokens that can be loaded beyond the
161 + num_computed_tokens.
162 +
163 + Args:
164 + request (Request): the request object.
165 + num_computed_tokens (int): the number of locally
166 + computed tokens for this request
167 +
168 + Returns:
169 + A tuple with the following elements:
170 + - The number of tokens that can be loaded beyond what is
171 + already computed.
172 + - `True` if tokens will be loaded asynchronously
173 + (between scheduler steps).
174 + """
175 + num_blocks = request.num_tokens // self.offloaded_block_size
176 +
177 + assert (len(request.block_hashes) //
178 + self.block_size_factor == num_blocks)
179 + block_hashes = self._get_block_hashes(request)
180 +
181 + self.manager.touch(block_hashes)
182 +
183 + full_block_tokens = self.offloaded_block_size * num_blocks
184 + if full_block_tokens - num_computed_tokens < self.offloaded_block_size:
185 + # we can load less than a block, skip
186 + return 0, False
187 +
188 + start_block_idx = num_computed_tokens // self.offloaded_block_size
189 + hits = self.manager.lookup(
190 + self._get_block_hashes(request, start_idx=start_block_idx))
191 + if hits == 0:
192 + return 0, False
193 +
194 + num_hit_tokens = (self.offloaded_block_size *
195 + (start_block_idx + hits) - num_computed_tokens)
196 + logger.debug(
197 + "Request %s hit %s offloaded tokens after %s GPU hit tokens",
198 + request.request_id,
199 + num_hit_tokens,
200 + num_computed_tokens,
201 + )
202 + if num_hit_tokens < self.offloaded_block_size:
203 + return 0, False
204 +
205 + return num_hit_tokens, True
206 +
207 + def update_state_after_alloc(self, request: Request, blocks: KVCacheBlocks,
208 + num_external_tokens: int):
209 + self._requests[request.request_id] = request
210 + # the block ids are updated in _get_reqs_to_store
211 + self._request_block_ids[request.request_id] = []
212 +
213 + if num_external_tokens == 0:
214 + return
215 +
216 + block_groups = blocks.get_block_ids()
217 + block_ids = block_groups[0]
218 +
219 + num_computed_gpu_blocks = sum(block.block_hash is not None
220 + for block in blocks.blocks[0])
221 + num_computed_tokens = num_computed_gpu_blocks * self.gpu_block_size
222 + full_block_tokens = num_computed_tokens + num_external_tokens
223 + assert full_block_tokens % self.offloaded_block_size == 0
224 +
225 + num_pending_gpu_blocks = len(block_ids) - num_computed_gpu_blocks
226 + assert (num_external_tokens == num_pending_gpu_blocks *
227 + self.gpu_block_size)
228 +
229 + start_block_idx = num_computed_tokens // self.offloaded_block_size
230 + num_blocks = full_block_tokens // self.offloaded_block_size
231 +
232 + assert (len(request.block_hashes) // self.block_size_factor
233 + >= num_blocks)
234 + block_hashes = self._get_block_hashes(request,
235 + start_idx=start_block_idx,
236 + end_idx=num_blocks)
237 +
238 + src_spec = self.manager.prepare_load(block_hashes)
239 + dst_spec = GPULoadStoreSpec(block_ids[num_computed_gpu_blocks:])
240 +
241 + block_hashes = self._get_block_hashes(request,
242 + start_idx=start_block_idx,
243 + end_idx=num_blocks)
244 +
245 + self._reqs_to_load[request.request_id] = (src_spec, dst_spec)
246 + self._reqs_being_loaded[request.request_id].update(block_hashes)
247 + self._next_stored_block_idx[request.request_id] = num_blocks
248 +
249 + def _get_reqs_to_store(self, scheduler_output: SchedulerOutput):
250 + reqs_to_store: dict[ReqId, TransferSpec] = {}
251 + # iterate over both new and cached requests
252 + for req_id, new_block_id_groups, preempted in yield_req_data(
253 + scheduler_output):
254 +
255 + if preempted:
256 + self._request_block_ids[req_id] = []
257 +
258 + if new_block_id_groups:
259 + new_block_ids = new_block_id_groups[0]
260 + self._request_block_ids[req_id] += new_block_ids
261 +
262 + block_ids = self._request_block_ids[req_id]
263 +
264 + req = self._requests[req_id]
265 + new_tokens = scheduler_output.num_scheduled_tokens[req_id]
266 + total_tokens = req.num_computed_tokens + new_tokens
267 + num_blocks = total_tokens // self.offloaded_block_size
268 + start_block_idx = self._next_stored_block_idx.get(req_id, 0)
269 + num_new_blocks = num_blocks - start_block_idx
270 +
271 + if num_new_blocks <= 0:
272 + continue
273 +
274 + num_gpu_blocks = num_blocks * self.block_size_factor
275 + assert len(req.block_hashes) >= num_gpu_blocks
276 +
277 + new_block_hashes = self._get_block_hashes(
278 + req, start_idx=start_block_idx, end_idx=num_blocks)
279 + store_output = self.manager.prepare_store(new_block_hashes)
280 + if store_output is None:
281 + logger.warning("Cannot store %s blocks", num_new_blocks)
282 + break
283 +
284 + self._next_stored_block_idx[req_id] = num_blocks
285 +
286 + if not store_output.block_hashes_to_store:
287 + continue
288 + block_hashes_to_store = set(store_output.block_hashes_to_store)
289 +
290 + block_hashes = self._get_block_hashes(req, end_idx=num_blocks)
291 + self.manager.touch(block_hashes)
292 +
293 + new_block_hashes = self._get_block_hashes(
294 + req, start_idx=start_block_idx, end_idx=num_blocks)
295 + dst_spec = store_output.store_spec
296 + src_block_ids: list[int] = []
297 + for idx, blk_hash in enumerate(new_block_hashes):
298 + if blk_hash not in block_hashes_to_store:
299 + continue
300 + offloaded_block_idx = start_block_idx + idx
301 + gpu_block_idx = offloaded_block_idx * self.block_size_factor
302 + for i in range(self.block_size_factor):
303 + src_block_ids.append(block_ids[gpu_block_idx + i])
304 + src_spec = GPULoadStoreSpec(src_block_ids)
305 +
306 + reqs_to_store[req_id] = (src_spec, dst_spec)
307 + self._reqs_being_stored[req_id] |= block_hashes_to_store
308 +
309 + logger.debug(
310 + "Request %s offloading %s blocks starting from block #%d",
311 + req_id,
312 + len(block_hashes_to_store),
313 + start_block_idx,
314 + )
315 +
316 + return reqs_to_store
317 +
318 + def build_connector_meta(
319 + self, scheduler_output: SchedulerOutput) -> KVConnectorMetadata:
320 + meta = OffloadingConnectorMetadata(
321 + reqs_to_load=self._reqs_to_load,
322 + reqs_to_store=self._get_reqs_to_store(scheduler_output))
323 + self._reqs_to_load = {}
324 + return meta
325 +
326 + def update_connector_output(self, connector_output: KVConnectorOutput):
327 + """
328 + Update KVConnector state from worker-side connectors output.
329 +
330 + Args:
331 + connector_output (KVConnectorOutput): the worker-side
332 + connectors output.
333 + """
334 + for req_id in connector_output.finished_sending or []:
335 + block_hashes = self._reqs_being_stored.pop(req_id, None)
336 + if block_hashes:
337 + self.manager.complete_store(block_hashes)
338 +
339 + for req_id in connector_output.finished_recving or []:
340 + block_hashes = self._reqs_being_loaded.pop(req_id, None)
341 + if block_hashes:
342 + self.manager.complete_load(block_hashes)
343 +
344 + def request_finished(
345 + self,
346 + request: Request,
347 + block_ids: list[int],
348 + ) -> tuple[bool, Optional[dict[str, Any]]]:
349 + """
350 + Called when a request has finished, before its blocks are freed.
351 +
352 + Returns:
353 + True if the request is being saved/sent asynchronously and blocks
354 + should not be freed until the request_id is returned from
355 + get_finished().
356 + Optional KVTransferParams to be included in the request outputs
357 + returned by the engine.
358 + """
359 + req_id = request.request_id
360 + self._requests.pop(req_id, None)
361 + self._request_block_ids.pop(req_id, None)
362 + self._next_stored_block_idx.pop(req_id, None)
363 +
364 + request_being_stored = req_id in self._reqs_being_stored
365 + return request_being_stored, None
366 +
367 + def take_events(self) -> Iterable[KVCacheEvent]:
368 + """Take the KV cache events from the connector.
369 +
370 + Returns:
371 + A list of KV cache events.
372 + """
373 + for event in self.manager.take_events():
374 + if event.removed:
375 + yield BlockRemoved(block_hashes=event.block_hashes,
376 + medium=event.medium)
377 + else:
378 + yield BlockStored(block_hashes=event.block_hashes,
379 + parent_block_hash=None,
380 + token_ids=[],
381 + lora_id=None,
382 + block_size=event.block_size,
383 + medium=event.medium)
384 +
385 +
386 +class OffloadingConnectorWorker:
387 + """Implementation of Worker side methods"""
388 +
389 + def __init__(self, spec: OffloadingSpec):
390 + self.spec = spec
391 + self.worker = OffloadingWorker()
392 +
393 + self._job_counter = 0
394 +
395 + # req_id -> (job_id, store)
396 + self._jobs: dict[int, tuple[ReqId, bool]] = {}
397 + # req_id -> active job IDs
398 + self._load_job: dict[ReqId, int] = {}
399 + # req_id -> set(active job IDs)
400 + self._store_jobs = defaultdict[ReqId, set[int]](set)
401 +
402 + self._finished_reqs_waiting_for_store: set[ReqId] = set()
403 +
404 + def _generate_job_id(self) -> int:
405 + job_id = self._job_counter
406 + self._job_counter = job_id + 1
407 + return job_id
408 +
409 + def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
410 + for src_cls, dst_cls, handler in (self.spec.get_handlers(kv_caches)):
411 + self.worker.register_handler(src_cls, dst_cls, handler)
412 +
413 + def start_load_kv(self, metadata: OffloadingConnectorMetadata):
414 + for req_id, transfer_spec in metadata.reqs_to_load.items():
415 + job_id = self._generate_job_id()
416 + self._jobs[job_id] = (req_id, False)
417 + assert req_id not in self._load_job
418 + self._load_job[req_id] = job_id
419 + assert self.worker.transfer_async(job_id, transfer_spec)
420 +
421 + def start_store_kv(self, metadata: OffloadingConnectorMetadata):
422 + for req_id, transfer_spec in metadata.reqs_to_store.items():
423 + job_id = self._generate_job_id()
424 + self._jobs[job_id] = (req_id, True)
425 + self._store_jobs[req_id].add(job_id)
426 + assert self.worker.transfer_async(job_id, transfer_spec)
427 +
428 + def get_finished(self,
429 + finished_req_ids: set[str]) -> tuple[set[str], set[str]]:
430 + """
431 + Notifies worker-side connector ids of requests that have
432 + finished generating tokens.
433 + Returns a list of request IDs that finished loading or storing.
434 +
435 + Returns:
436 + ids of requests that have finished asynchronous transfer
437 + tuple of (sending/saving ids, recving/loading ids).
438 + """
439 + finished_sending = set()
440 + finished_recving = set()
441 + for job_id, success in self.worker.get_finished():
442 + # we currently do not support job failures
443 + assert success
444 + req_id, store = self._jobs.pop(job_id)
445 + if store:
446 + req_jobs = self._store_jobs[req_id]
447 + req_jobs.remove(job_id)
448 + if req_jobs:
449 + continue
450 +
451 + if req_id in self._finished_reqs_waiting_for_store:
452 + self._finished_reqs_waiting_for_store.remove(req_id)
453 + finished_sending.add(req_id)
454 + del self._store_jobs[req_id]
455 + else:
456 + req_job = self._load_job[req_id]
457 + assert job_id == req_job
458 + del self._load_job[req_id]
459 + finished_recving.add(req_id)
460 +
461 + for req_id in finished_req_ids:
462 + pending_req_jobs = self._store_jobs.get(req_id)
463 + if pending_req_jobs:
464 + self._finished_reqs_waiting_for_store.add(req_id)
465 + elif pending_req_jobs is not None:
466 + finished_sending.add(req_id)
467 + del self._store_jobs[req_id]
468 +
469 + return finished_sending, finished_recving
470 +
471 +
472 +def yield_req_data(
473 + scheduler_output) -> Iterator[tuple[str, tuple[list[int], ...], bool]]:
474 + """
475 + Yields:
476 + (req_id, new_block_id_groups, preempted)
477 + """
478 + # new requests
479 + for req_data in scheduler_output.scheduled_new_reqs:
480 + yield req_data.req_id, req_data.block_ids, False
481 +
482 + # cached requests
483 + cached_reqs = scheduler_output.scheduled_cached_reqs
484 + yield from zip(cached_reqs.req_ids, cached_reqs.new_block_ids,
485 + cached_reqs.resumed_from_preemption)
vllm/v1/kv_offload/factory.py new
+53
@@ -0,0 +1,53 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import importlib
4 +from typing import TYPE_CHECKING, Callable
5 +
6 +from vllm.logger import init_logger
7 +from vllm.v1.kv_offload.spec import OffloadingSpec
8 +
9 +if TYPE_CHECKING:
10 + from vllm.config import VllmConfig
11 +
12 +logger = init_logger(__name__)
13 +
14 +
15 +class OffloadingSpecFactory:
16 + _registry: dict[str, Callable[[], type[OffloadingSpec]]] = {}
17 +
18 + @classmethod
19 + def register_spec(cls, name: str, module_path: str,
20 + class_name: str) -> None:
21 + """Register a spec with a lazy-loading module and class name."""
22 + if name in cls._registry:
23 + raise ValueError(f"Connector '{name}' is already registered.")
24 +
25 + def loader() -> type[OffloadingSpec]:
26 + module = importlib.import_module(module_path)
27 + return getattr(module, class_name)
28 +
29 + cls._registry[name] = loader
30 +
31 + @classmethod
32 + def create_spec(
33 + cls,
34 + config: "VllmConfig",
35 + ) -> OffloadingSpec:
36 + kv_transfer_config = config.kv_transfer_config
37 + assert kv_transfer_config is not None
38 + extra_config = kv_transfer_config.kv_connector_extra_config
39 + spec_name = extra_config.get("spec_name", "CPUOffloadingSpec")
40 + if spec_name in cls._registry:
41 + spec_cls = cls._registry[spec_name]()
42 + else:
43 + spec_module_path = extra_config.get("spec_module_path")
44 + if spec_module_path is None:
45 + raise ValueError(f"Unsupported spec type: {spec_name}")
46 + spec_module = importlib.import_module(spec_module_path)
47 + spec_cls = getattr(spec_module, spec_name)
48 + assert issubclass(spec_cls, OffloadingSpec)
49 + logger.info("Creating offloading spec with name: %s", spec_name)
50 + return spec_cls(config)
51 +
52 +
53 +# Register various specs here.
vllm/v1/kv_offload/spec.py new
+61
@@ -0,0 +1,61 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from abc import ABC, abstractmethod
4 +from collections.abc import Iterator
5 +from typing import TYPE_CHECKING
6 +
7 +import torch
8 +
9 +from vllm.logger import init_logger
10 +from vllm.v1.kv_offload.abstract import LoadStoreSpec, OffloadingManager
11 +from vllm.v1.kv_offload.worker.worker import OffloadingHandler
12 +
13 +if TYPE_CHECKING:
14 + from vllm.config import VllmConfig
15 +
16 +logger = init_logger(__name__)
17 +
18 +
19 +class OffloadingSpec(ABC):
20 + """Spec for an offloading connector"""
21 +
22 + def __init__(self, vllm_config: "VllmConfig"):
23 + logger.warning(
24 + "Initializing OffloadingSpec. This API is experimental and "
25 + "subject to change in the future as we iterate the design.")
26 + self.vllm_config = vllm_config
27 +
28 + kv_transfer_config = vllm_config.kv_transfer_config
29 + assert kv_transfer_config is not None
30 + self.extra_config = kv_transfer_config.kv_connector_extra_config
31 +
32 + self.gpu_block_size = vllm_config.cache_config.block_size
33 + self.offloaded_block_size = int(
34 + self.extra_config.get("block_size", self.gpu_block_size))
35 +
36 + assert self.offloaded_block_size % self.gpu_block_size == 0
37 +
38 + @abstractmethod
39 + def get_manager(self) -> OffloadingManager:
40 + """
41 + Get an OffloadingManager that will be used
42 + by the scheduler-side offloading connector to track
43 + offloaded blocks and manage evictions.
44 + """
45 + pass
46 +
47 + @abstractmethod
48 + def get_handlers(
49 + self, kv_caches: dict[str, torch.Tensor]
50 + ) -> Iterator[tuple[type[LoadStoreSpec], type[LoadStoreSpec],
51 + OffloadingHandler]]:
52 + """
53 + Get offloading handlers along with their respective src and dst types.
54 +
55 + Args:
56 + kv_caches: A dictionary of layer_name -> gpu_kv_cache tensor.
57 +
58 + Yields:
59 + Tuples of (src_type, dst_type, offloading_handler).
60 + """
61 + pass