@tduncan / Nvidia-Optimization-Stack-VLLM / commits / f237e16b41

[KV Offload] Replace OffloadingHandler with OffloadingWorker (#45053)

Signed-off-by: Martin Hickey <martin.hickey@ie.ibm.com>

Martin Hickey committed Jun 24, 2026 at 12:44 UTC f237e16b41bb444b3c9994260a36f9c2388bd019

13 files changed +176 -465

tests/v1/kv_connector/unit/offloading_connector/test_worker.py

+3 -3

@@ -98,7 +98,7 @@ def _make_worker(kv_cache_config: KVCacheConfig):

98 spec = MagicMock(spec=OffloadingSpec)

99 spec.kv_cache_config = kv_cache_config

100 spec.vllm_config = MagicMock()

101 - spec.get_handlers.return_value = iter([])

101 + spec.get_worker.return_value = MagicMock()

102

103 worker = OffloadingConnectorWorker(spec=spec)

104 worker.worker = MagicMock()

@@ -279,7 +279,7 @@ def test_register_kv_caches(backend):

279 worker, spec = _make_worker(kv_cache_config)

280 worker.register_kv_caches(kv_caches)

281

282 - canonical = spec.get_handlers.call_args[0][0]

282 + canonical = spec.get_worker.call_args[0][0]

283 assert isinstance(canonical, CanonicalKVCaches)

284

285 # -- Expected block tensors ----------------------------------------------

@@ -422,7 +422,7 @@ def test_register_kv_caches_uniform_type(backend):

422 worker, spec = _make_worker(kv_cache_config)

423 worker.register_kv_caches(kv_caches)

424

425 - canonical = spec.get_handlers.call_args[0][0]

425 + canonical = spec.get_worker.call_args[0][0]

426 assert isinstance(canonical, CanonicalKVCaches)

427

428 for block_tensor in canonical.tensors:

tests/v1/kv_connector/unit/offloading_connector/utils.py

+24 -24

@@ -1,6 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 -from collections.abc import Callable, Iterable, Iterator

3 +from collections.abc import Callable, Iterable

4 from dataclasses import dataclass

5 from typing import Any

6 from unittest.mock import MagicMock

@@ -43,19 +43,17 @@ from vllm.v1.kv_cache_interface import (

43 KVCacheGroupSpec,

44 )

45 from vllm.v1.kv_offload.base import (

46 + CanonicalKVCaches,

47 GPULoadStoreSpec,

48 LoadStoreSpec,

49 OffloadingManager,

50 OffloadingSpec,

51 + OffloadingWorker,

52 OffloadKey,

53 PrepareStoreOutput,

54 RequestOffloadingContext,

53 - make_offload_key,

54 -)

55 -from vllm.v1.kv_offload.worker.worker import (

56 - OffloadingHandler,

55 TransferResult,

58 - TransferSpec,

56 + make_offload_key,

57 )

58 from vllm.v1.request import Request

59 from vllm.v1.structured_output import StructuredOutputManager

@@ -81,9 +79,9 @@ class MockLoadStoreSpec(LoadStoreSpec):

79 return repr(self.offload_keys)

84 -class MockOffloadingHandler(OffloadingHandler):

82 +class MockOffloadingWorker(OffloadingWorker):

83 def __init__(self):

86 - self.transfer_specs: dict[int, TransferSpec] = {}

84 + self.transfer_specs: dict[int, tuple[LoadStoreSpec, LoadStoreSpec]] = {}

85 self.completed_transfers: list[TransferResult] = []

86 self.waiting_jobs: set[int] = set()

87 self.completed_jobs: list[int] = []

@@ -94,8 +92,17 @@ class MockOffloadingHandler(OffloadingHandler):

92 self.completed_transfers = []

93 return finished

97 - def transfer_async(self, job_id: int, spec: TransferSpec) -> bool:

98 - self.transfer_specs[job_id] = spec

95 + def submit_store(

96 + self, job_id: int, src_spec: LoadStoreSpec, dst_spec: LoadStoreSpec

97 + ) -> bool: # type: ignore[override]

98 + self.transfer_specs[job_id] = (src_spec, dst_spec)

99 + self.waiting_jobs.add(job_id)

100 + return True

101 +

102 + def submit_load(

103 + self, job_id: int, src_spec: LoadStoreSpec, dst_spec: LoadStoreSpec

104 + ) -> bool: # type: ignore[override]

105 + self.transfer_specs[job_id] = (src_spec, dst_spec)

106 self.waiting_jobs.add(job_id)

107 return True

108

@@ -109,7 +116,6 @@ class MockOffloadingHandler(OffloadingHandler):

116 success=True,

117 transfer_size=None,

118 transfer_time=None,

112 - transfer_type=None,

119 )

120 self.completed_transfers.append(result)

121

@@ -127,21 +133,18 @@ class MockOffloadingSpec(OffloadingSpec):

133 self.manager.prepare_load = lambda keys, req_context: MockLoadStoreSpec(keys)

134 self.manager.lookup.return_value = False

135 self.manager.on_new_request.return_value = RequestOffloadingContext()

130 - self.handler = MockOffloadingHandler()

136 + self.handler = MockOffloadingWorker()

137

138 def get_manager(self) -> OffloadingManager:

139 return self.manager

140

135 - def get_handlers(

136 - self, _

137 - ) -> Iterator[tuple[type[LoadStoreSpec], type[LoadStoreSpec], OffloadingHandler]]:

138 - yield GPULoadStoreSpec, MockLoadStoreSpec, self.handler

139 - yield MockLoadStoreSpec, GPULoadStoreSpec, self.handler

141 + def get_worker(self, _: CanonicalKVCaches) -> OffloadingWorker:

142 + return self.handler

143

144 def complete_transfers(self):

145 self.handler.complete_jobs(self.handler.waiting_jobs.copy())

146

144 - def get_completed_transfers(self) -> list[TransferSpec]:

147 + def get_completed_transfers(self) -> list[tuple[LoadStoreSpec, LoadStoreSpec]]:

148 specs = [

149 self.handler.transfer_specs[job_id]

150 for job_id in self.handler.completed_jobs

@@ -149,7 +152,7 @@ class MockOffloadingSpec(OffloadingSpec):

152 self.handler.completed_jobs.clear()

153 return specs

154

152 - def get_flushed_transfers(self):

155 + def get_flushed_transfers(self) -> list[tuple[LoadStoreSpec, LoadStoreSpec]]:

156 specs = [

157 self.handler.transfer_specs[job_id] for job_id in self.handler.flushed_jobs

158 ]

@@ -362,8 +365,7 @@ class RequestRunner:

365 self.scheduler.add_request(req)

366

367 def _parse_transfers(self):

365 - for transfer_spec in self.offloading_spec.get_flushed_transfers():

366 - src_spec, dst_spec = transfer_spec

368 + for src_spec, dst_spec in self.offloading_spec.get_flushed_transfers():

369 if isinstance(src_spec, GPULoadStoreSpec):

370 # store flush

371 for block_id in src_spec.block_ids:

@@ -375,9 +377,7 @@ class RequestRunner:

377

378 block_size_factor = self.block_size_factor

379

378 - for transfer_spec in self.offloading_spec.get_completed_transfers():

379 - src_spec, dst_spec = transfer_spec

380 -

380 + for src_spec, dst_spec in self.offloading_spec.get_completed_transfers():

381 if isinstance(src_spec, GPULoadStoreSpec):

382 store = True

383 gpu_spec = src_spec

tests/v1/kv_offload/cpu/test_gpu_worker.py

+20 -21

@@ -17,7 +17,7 @@ from vllm.v1.kv_offload.base import (

17 GPULoadStoreSpec,

18 )

19 from vllm.v1.kv_offload.cpu.common import CPULoadStoreSpec

20 -from vllm.v1.kv_offload.cpu.gpu_worker import CpuGpuOffloadingHandlers

20 +from vllm.v1.kv_offload.cpu.gpu_worker import CPUOffloadingWorker

21 from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion

23 NUM_GPU_BLOCKS = [64]

@@ -103,7 +103,7 @@ def test_transfer(

103 cpu_page_size=cpu_page_size,

104 )

105

106 - handlers = CpuGpuOffloadingHandlers(

106 + worker = CPUOffloadingWorker(

107 kv_caches=kv_caches,

108 block_size_factor=block_size_factor,

109 num_cpu_blocks=num_cpu_blocks,

@@ -130,7 +130,7 @@ def test_transfer(

130

131 # set transfer direction

132 if gpu_to_cpu:

133 - handler = handlers.gpu_to_cpu_handler

133 + handler = worker._store_handler

134 src_spec = GPULoadStoreSpec(

135 gpu_blocks, group_sizes=(len(gpu_blocks),), block_indices=(blocks_to_skip,)

136 )

@@ -138,7 +138,7 @@ def test_transfer(

138 dst_to_src = dict(zip(cpu_blocks_expanded, gpu_blocks))

139 num_dst_sub_blocks = num_gpu_blocks

140 else:

141 - handler = handlers.cpu_to_gpu_handler

141 + handler = worker._load_handler

142 src_spec = CPULoadStoreSpec(cpu_blocks)

143 dst_spec = GPULoadStoreSpec(

144 gpu_blocks, group_sizes=(len(gpu_blocks),), block_indices=(blocks_to_skip,)

@@ -156,23 +156,21 @@ def test_transfer(

156 orig_src_tensors = [x.clone() for x in handler.src_tensors]

157 orig_dst_tensors = [x.clone() for x in handler.dst_tensors]

158

159 - # call transfer function

159 + # call transfer function via public API

160 start_time = time.time()

161 - assert handler.transfer_async(1, (src_spec, dst_spec))

161 + if gpu_to_cpu:

162 + assert worker.submit_store(1, src_spec, dst_spec)

163 + else:

164 + assert worker.submit_load(1, src_spec, dst_spec)

165 assert {x.job_id for x in handler._transfers} == {1}

166

167 # wait for transfer to complete

168 end_time = time.time() + 10

169 while time.time() < end_time:

167 - finished = handler.get_finished()

170 + finished = worker.get_finished()

171 if finished:

172 assert finished[0].job_id == 1

173 assert finished[0].success

171 - assert (

172 - finished[0].transfer_type == ("GPU", "CPU")

173 - if gpu_to_cpu

174 - else ("CPU", "GPU")

175 - )

174 assert finished[0].transfer_size == (

175 len(gpu_blocks)

176 * sum([x.page_size_bytes for x in handler.kv_cache_groups_data_refs[0]])

@@ -208,8 +206,7 @@ def test_transfer(

206 del orig_tensor, tensor, src_tensor, dst_tensor, orig_dst_tensor

207 del src_view, dst_view, orig_dst_view, expected

208

211 - handlers.cpu_to_gpu_handler.shutdown()

212 - handlers.gpu_to_cpu_handler.shutdown()

209 + worker.shutdown()

210 if mmap_region:

211 mmap_region.cleanup()

212

@@ -276,7 +273,7 @@ def test_transfer_multi_group(

273 tensors=kv_cache_tensors, group_data_refs=kv_cache_groups_data_refs

274 )

275

279 - handlers = CpuGpuOffloadingHandlers(

276 + worker = CPUOffloadingWorker(

277 kv_caches=canonical_kv_caches,

278 block_size_factor=block_size_factor,

279 num_cpu_blocks=num_cpu_blocks,

@@ -338,7 +335,7 @@ def test_transfer_multi_group(

335 block_indices: list[int] = [0, 0, sub_blocks_to_skip]

336

337 if gpu_to_cpu:

341 - handler = handlers.gpu_to_cpu_handler

338 + handler = worker._store_handler

339 src_spec = GPULoadStoreSpec(

340 gpu_blocks, group_sizes=group_sizes, block_indices=block_indices

341 )

@@ -352,7 +349,7 @@ def test_transfer_multi_group(

349 ]

350 num_dst_sub_blocks = num_cpu_blocks * block_size_factor

351 else:

355 - handler = handlers.cpu_to_gpu_handler

352 + handler = worker._load_handler

353 src_spec = CPULoadStoreSpec(cpu_blocks)

354 dst_spec = GPULoadStoreSpec(

355 gpu_blocks, group_sizes=group_sizes, block_indices=block_indices

@@ -375,12 +372,15 @@ def test_transfer_multi_group(

372 orig_src_tensors = [x.clone() for x in handler.src_tensors]

373 orig_dst_tensors = [x.clone() for x in handler.dst_tensors]

374

378 - assert handler.transfer_async(1, (src_spec, dst_spec))

375 + if gpu_to_cpu:

376 + assert worker.submit_store(1, src_spec, dst_spec)

377 + else:

378 + assert worker.submit_load(1, src_spec, dst_spec)

379 assert {x.job_id for x in handler._transfers} == {1}

380

381 end_time = time.time() + 10

382 while time.time() < end_time:

383 - finished = handler.get_finished()

383 + finished = worker.get_finished()

384 if finished:

385 assert finished[0].job_id == 1

386 assert finished[0].success

@@ -418,5 +418,4 @@ def test_transfer_multi_group(

418 dst_view[dst_sub_block].cpu(), expected.cpu()

419 )

420

421 - handlers.cpu_to_gpu_handler.shutdown()

422 - handlers.gpu_to_cpu_handler.shutdown()

421 + worker.shutdown()

tests/v1/kv_offload/test_worker.py deleted

-165

@@ -1,165 +0,0 @@

1 -# SPDX-License-Identifier: Apache-2.0

2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 -from vllm.v1.kv_offload.base import LoadStoreSpec

4 -from vllm.v1.kv_offload.worker.worker import (

5 - OffloadingHandler,

6 - OffloadingWorker,

7 - TransferResult,

8 - TransferSpec,

9 -)

10 -

11 -

12 -class LoadStoreSpec1(LoadStoreSpec):

13 - def __init__(

14 - self,

15 - submit_success: bool = True,

16 - async_success: bool = True,

17 - exception: bool = False,

18 - ):

19 - self.finished = False

20 - self.submit_success = submit_success

21 - self.async_success = async_success

22 - self.exception = exception

23 -

24 - @staticmethod

25 - def medium() -> str:

26 - return "1"

27 -

28 - def __repr__(self):

29 - return f"{self.medium()}: {id(self)}"

30 -

31 -

32 -class LoadStoreSpec2(LoadStoreSpec):

33 - @staticmethod

34 - def medium() -> str:

35 - return "2"

36 -

37 - def __repr__(self):

38 - return f"{self.medium()}: {id(self)}"

39 -

40 -

41 -class OffloadingHandler1To2(OffloadingHandler):

42 - def __init__(self):

43 - self.transfers: dict[int, LoadStoreSpec1] = {}

44 -

45 - def transfer_async(self, job_id: int, spec: TransferSpec) -> bool:

46 - src, dst = spec

47 - assert isinstance(src, LoadStoreSpec1)

48 - assert isinstance(dst, LoadStoreSpec2)

49 -

50 - if src.exception:

51 - raise Exception("An expected exception. Don't worry!")

52 - if not src.submit_success:

53 - return False

54 -

55 - self.transfers[job_id] = src

56 - return True

57 -

58 - def get_finished(self) -> list[TransferResult]:

59 - finished = []

60 - for job_id, spec in list(self.transfers.items()):

61 - if spec.finished:

62 - finished.append((job_id, spec.async_success))

63 - del self.transfers[job_id]

64 - return finished

65 -

66 - def wait(self, job_ids: set[int]) -> None:

67 - for job_id in job_ids:

68 - spec = self.transfers.get(job_id)

69 - if spec:

70 - assert spec.finished

71 -

72 -

73 -class OffloadingHandler2To1(OffloadingHandler):

74 - def __init__(self):

75 - self.transfers: dict[int, LoadStoreSpec1] = {}

76 -

77 - def transfer_async(self, job_id: int, spec: TransferSpec) -> bool:

78 - src, dst = spec

79 - assert isinstance(src, LoadStoreSpec2)

80 - assert isinstance(dst, LoadStoreSpec1)

81 -

82 - self.transfers[job_id] = dst

83 - return True

84 -

85 - def get_finished(self) -> list[TransferResult]:

86 - finished = []

87 - for job_id, spec in list(self.transfers.items()):

88 - if spec.finished:

89 - finished.append((job_id, spec.async_success))

90 - del self.transfers[job_id]

91 - return finished

92 -

93 - def wait(self, job_ids: set[int]) -> None:

94 - for job_id in job_ids:

95 - spec = self.transfers.get(job_id)

96 - if spec:

97 - assert spec.finished

98 -

99 -

100 -def test_offloading_worker():

101 - """

102 - Tests OffloadingWorker with 2 handlers.

103 - One handler performs 1->2 transfers, and the other handles 2->1.

104 - """

105 - worker = OffloadingWorker()

106 - handler1to2 = OffloadingHandler1To2()

107 - handler2to1 = OffloadingHandler2To1()

108 - worker.register_handler(LoadStoreSpec1, LoadStoreSpec2, handler1to2)

109 - worker.register_handler(LoadStoreSpec2, LoadStoreSpec1, handler2to1)

110 -

111 - # 1st transfer 1->2 (exception)

112 - src1 = LoadStoreSpec1(exception=True)

113 - dst1 = LoadStoreSpec2()

114 - assert not worker.transfer_async(1, (src1, dst1))

115 -

116 - # 2ed transfer 1->2 (failure to submit)

117 - src2 = LoadStoreSpec1(submit_success=False)

118 - dst2 = LoadStoreSpec2()

119 - assert not worker.transfer_async(2, (src2, dst2))

120 -

121 - # 3rd transfer 1->2 (failure)

122 - src3 = LoadStoreSpec1(async_success=False)

123 - dst3 = LoadStoreSpec2()

124 - assert worker.transfer_async(3, (src3, dst3))

125 -

126 - # 4th transfer 1->2 (success)

127 - src4 = LoadStoreSpec1()

128 - dst4 = LoadStoreSpec2()

129 - worker.transfer_async(4, (src4, dst4))

130 - assert set(handler1to2.transfers.keys()) == {3, 4}

131 -

132 - # 5th transfer 2->1

133 - src5 = LoadStoreSpec2()

134 - dst5 = LoadStoreSpec1()

135 - worker.transfer_async(5, (src5, dst5))

136 - assert set(handler2to1.transfers.keys()) == {5}

137 -

138 - # no transfer completed yet

139 - assert worker.get_finished() == []

140 -

141 - # complete 3rd, 4th

142 - src3.finished = True

143 - src4.finished = True

144 -

145 - # 6th transfer 1->2

146 - src6 = LoadStoreSpec1()

147 - dst6 = LoadStoreSpec2()

148 - worker.transfer_async(6, (src6, dst6))

149 -

150 - # 7th transfer 2->1

151 - src7 = LoadStoreSpec2()

152 - dst7 = LoadStoreSpec1()

153 - worker.transfer_async(7, (src7, dst7))

154 -

155 - # 6th and 7th transfers started

156 - assert 6 in handler1to2.transfers

157 - assert 7 in handler2to1.transfers

158 -

159 - # verify result of 3rd and 4th transfers

160 - assert sorted(worker.get_finished()) == [(3, False), (4, True)]

161 -

162 - # complete 6th and 7th transfers

163 - src6.finished = True

164 - dst7.finished = True

165 - assert sorted(worker.get_finished()) == [(6, True), (7, True)]

vllm/distributed/kv_transfer/kv_connector/v1/offloading/common.py

+3 -2

@@ -6,7 +6,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.base import (

6 KVConnectorMetadata,

7 KVConnectorWorkerMetadata,

8 )

9 -from vllm.v1.kv_offload.worker.worker import TransferSpec

9 +from vllm.v1.kv_offload.base import LoadStoreSpec

11 ReqId = str

@@ -60,7 +60,8 @@ class TransferJob:

60 """

62 req_id: ReqId

63 - transfer_spec: TransferSpec

63 + src_spec: LoadStoreSpec

64 + dst_spec: LoadStoreSpec

67 @dataclass

vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py

+3 -2

@@ -762,7 +762,8 @@ class OffloadingConnectorScheduler:

762 load_job_id = self._generate_job_id()

763 self._current_batch_load_jobs[load_job_id] = TransferJob(

764 req_id=request.request_id,

765 - transfer_spec=(src_spec, dst_spec),

765 + src_spec=src_spec,

766 + dst_spec=dst_spec,

767 )

768 # a load can only be issued when no other jobs are pending.

769 assert not req_status.transfer_jobs

@@ -998,7 +999,7 @@ class OffloadingConnectorScheduler:

999 )

1000

1001 store_jobs[job_id] = TransferJob(

1001 - req_id=req_id, transfer_spec=(src_spec, dst_spec)

1002 + req_id=req_id, src_spec=src_spec, dst_spec=dst_spec

1003 )

1004

1005 logger.debug(

vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py

+26 -18

@@ -21,11 +21,10 @@ from vllm.v1.kv_offload.base import (

21 CanonicalKVCacheRef,

22 CanonicalKVCaches,

23 CanonicalKVCacheTensor,

24 + GPULoadStoreSpec,

25 + LoadStoreSpec,

26 OffloadingSpec,

25 -)

26 -from vllm.v1.kv_offload.worker.worker import (

27 OffloadingWorker,

28 - TransferSpec,

28 )

30 logger = init_logger(__name__)

@@ -36,16 +35,17 @@ class OffloadingConnectorWorker:

36 def __init__(self, spec: OffloadingSpec):

37 self.spec = spec

39 - self.worker = OffloadingWorker()

38 + self.worker: OffloadingWorker | None = None

40 # job_id -> req_id for in-flight loads.

41 self._load_jobs: dict[int, ReqId] = {}

43 - self._unsubmitted_store_jobs: list[tuple[int, TransferSpec]] = []

42 + self._unsubmitted_store_jobs: list[

43 + tuple[int, GPULoadStoreSpec, LoadStoreSpec]

44 + ] = []

45 self._connector_worker_meta = OffloadingWorkerMetadata()

46 - def _register_handlers(self, kv_caches: CanonicalKVCaches):

47 - for src_cls, dst_cls, handler in self.spec.get_handlers(kv_caches):

48 - self.worker.register_handler(src_cls, dst_cls, handler)

47 + def _init_worker(self, kv_caches: CanonicalKVCaches) -> None:

48 + self.worker = self.spec.get_worker(kv_caches)

50 def register_kv_caches(

51 self, kv_caches: dict[str, torch.Tensor | list[torch.Tensor]]

@@ -138,7 +138,7 @@ class OffloadingConnectorWorker:

138 (block_stride, 1),

139 storage_offset=0,

140 )

141 - self._register_handlers(

141 + self._init_worker(

142 CanonicalKVCaches(

143 [CanonicalKVCacheTensor(packed_tensor, block_stride)],

144 [

@@ -204,7 +204,7 @@ class OffloadingConnectorWorker:

204 group_data_refs=group_data_refs,

205 )

206

207 - self._register_handlers(canonical_kv_caches)

207 + self._init_worker(canonical_kv_caches)

208

209 def register_cross_layers_kv_cache(

210 self, kv_cache: torch.Tensor, attn_backend: type[AttentionBackend]

@@ -251,11 +251,12 @@ class OffloadingConnectorWorker:

251 tensors=[kv_cache_tensor], group_data_refs=[[kv_cache_data_ref]]

252 )

253

254 - self._register_handlers(canonical_kv_caches)

254 + self._init_worker(canonical_kv_caches)

255

256 def handle_preemptions(self, kv_connector_metadata: OffloadingConnectorMetadata):

257 - for job_id, transfer_spec in self._unsubmitted_store_jobs:

258 - success = self.worker.transfer_async(job_id, transfer_spec)

257 + assert self.worker is not None

258 + for job_id, src_spec, dst_spec in self._unsubmitted_store_jobs:

259 + success = self.worker.submit_store(job_id, src_spec, dst_spec)

260 assert success

261 self._unsubmitted_store_jobs.clear()

262

@@ -263,14 +264,16 @@ class OffloadingConnectorWorker:

264 self.worker.wait(kv_connector_metadata.jobs_to_flush)

265

266 def start_kv_transfers(self, metadata: OffloadingConnectorMetadata):

266 - for job_id, transfer_spec in self._unsubmitted_store_jobs:

267 - success = self.worker.transfer_async(job_id, transfer_spec)

267 + assert self.worker is not None

268 + for job_id, src_spec, dst_spec in self._unsubmitted_store_jobs:

269 + success = self.worker.submit_store(job_id, src_spec, dst_spec)

270 assert success

271 self._unsubmitted_store_jobs.clear()

272

273 for job_id, entry in metadata.load_jobs.items():

274 self._load_jobs[job_id] = entry.req_id

273 - success = self.worker.transfer_async(job_id, entry.transfer_spec)

275 + assert isinstance(entry.dst_spec, GPULoadStoreSpec)

276 + success = self.worker.submit_load(job_id, entry.src_spec, entry.dst_spec)

277 assert success

278

279 def prepare_store_kv(self, metadata: OffloadingConnectorMetadata):

@@ -278,7 +281,10 @@ class OffloadingConnectorWorker:

281 # NOTE(orozery): defer the store to the beginning of the next

282 # engine step, so that offloading starts AFTER transfers related

283 # to token sampling, thereby avoiding delays to token generation.

281 - self._unsubmitted_store_jobs.append((job_id, entry.transfer_spec))

284 + assert isinstance(entry.src_spec, GPULoadStoreSpec)

285 + self._unsubmitted_store_jobs.append(

286 + (job_id, entry.src_spec, entry.dst_spec)

287 + )

288

289 def get_finished(self, finished_req_ids: set[str]) -> tuple[set[str], set[str]]:

290 """

@@ -290,6 +296,7 @@ class OffloadingConnectorWorker:

296 finished_recving so the base scheduler can resume requests

297 blocked on remote KV (and free aborted-during-load reqs).

298 """

299 + assert self.worker is not None

300 finished_recving: set[str] = set()

301 for transfer_result in self.worker.get_finished():

302 # we currently do not support job failures

@@ -328,4 +335,5 @@ class OffloadingConnectorWorker:

335 self._unsubmitted_store_jobs.clear()

336 self._load_jobs.clear()

337 self._connector_worker_meta = OffloadingWorkerMetadata()

331 - self.worker.shutdown()

338 + if self.worker is not None:

339 + self.worker.shutdown()

vllm/v1/kv_offload/base.py

+40 -8

@@ -5,7 +5,7 @@ Core abstractions for KV cache offloading in vLLM v1.

5 """

7 from abc import ABC, abstractmethod

8 -from collections.abc import Collection, Iterable, Iterator, Sequence

8 +from collections.abc import Collection, Iterable, Sequence

9 from dataclasses import dataclass

10 from enum import Enum

11 from typing import TYPE_CHECKING, Any, NewType

@@ -23,7 +23,6 @@ if TYPE_CHECKING:

23 OffloadingConnectorStats,

24 )

25 from vllm.v1.kv_cache_interface import KVCacheConfig

26 - from vllm.v1.kv_offload.worker.worker import OffloadingHandler

27 # `OffloadKey` identifies an offloaded block. It combines a block hash with

28 # its KV cache group index, encoded as raw bytes to avoid tuple GC overhead.

@@ -432,6 +431,41 @@ class CanonicalKVCaches:

431 group_data_refs: list[list[CanonicalKVCacheRef]]

432

433

434 +@dataclass

435 +class TransferResult:

436 + job_id: int

437 + success: bool

438 + transfer_size: int | None = None

439 + transfer_time: float | None = None

440 +

441 +

442 +class OffloadingWorker(ABC):

443 + """Runs in the worker process. Performs async KV transfers for ONE

444 + offloaded medium (e.g. CPU). Direction is explicit via submit_store /

445 + submit_load, so there is no (src_medium, dst_medium) routing."""

446 +

447 + @abstractmethod

448 + def submit_store(

449 + self, job_id: int, src_spec: GPULoadStoreSpec, dst_spec: LoadStoreSpec

450 + ) -> bool:

451 + """Async GPU -> offloaded medium."""

452 +

453 + @abstractmethod

454 + def submit_load(

455 + self, job_id: int, src_spec: LoadStoreSpec, dst_spec: GPULoadStoreSpec

456 + ) -> bool:

457 + """Async offloaded medium -> GPU."""

458 +

459 + @abstractmethod

460 + def get_finished(self) -> list[TransferResult]: ...

461 +

462 + @abstractmethod

463 + def wait(self, job_ids: set[int]) -> None: ...

464 +

465 + def shutdown(self) -> None:

466 + return

467 +

468 +

469 class OffloadingSpec(ABC):

470 """Spec for an offloading connector"""

471

@@ -524,16 +558,14 @@ class OffloadingSpec(ABC):

558 pass

559

560 @abstractmethod

527 - def get_handlers(

528 - self, kv_caches: CanonicalKVCaches

529 - ) -> Iterator[tuple[type[LoadStoreSpec], type[LoadStoreSpec], "OffloadingHandler"]]:

561 + def get_worker(self, kv_caches: CanonicalKVCaches) -> OffloadingWorker:

562 """

531 - Get offloading handlers along with their respective src and dst types.

563 + Get an OffloadingWorker that handles async KV transfers for this spec.

564

565 Args:

566 kv_caches: Canonicalized KV caches.

567

536 - Yields:

537 - Tuples of (src_type, dst_type, offloading_handler).

568 + Returns:

569 + An OffloadingWorker instance for this medium.

570 """

571 pass

vllm/v1/kv_offload/cpu/gpu_worker.py

+41 -21

@@ -7,7 +7,6 @@ from dataclasses import dataclass

8 import numpy as np

9 import torch

10 -from typing_extensions import override

11 from vllm import _custom_ops as ops

12 from vllm.logger import init_logger

@@ -20,17 +19,15 @@ from vllm.v1.kv_offload.base import (

19 CanonicalKVCacheRef,

20 CanonicalKVCaches,

21 GPULoadStoreSpec,

22 + LoadStoreSpec,

23 + OffloadingWorker,

24 + TransferResult,

25 )

26 from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion

27 from vllm.v1.kv_offload.cpu.swap_blocks_triton import (

28 THRESHOLD_BYTES,

29 swap_blocks_batch,

30 )

29 -from vllm.v1.kv_offload.worker.worker import (

30 - OffloadingHandler,

31 - TransferResult,

32 - TransferSpec,

33 -)

32 logger = init_logger(__name__)

@@ -166,10 +163,9 @@ def _new_descriptor_buffers(

163 )

164

165

169 -class SingleDirectionOffloadingHandler(OffloadingHandler):

166 +class SingleDirectionOffloadingHandler:

167 """

171 - SingleDirectionOffloadingHandler handles transfers for a single direction,

172 - either CPU->GPU or GPU->CPU.

168 + Handles transfers for a single direction, either CPU->GPU or GPU->CPU.

169 Transfers are guaranteed to be executed in order of their submission.

170 Each transfer uses a unique CUDA stream, and its stream will start

171 executing only after the streams of previous transfers have finished.

@@ -228,7 +224,6 @@ class SingleDirectionOffloadingHandler(OffloadingHandler):

224 self.src_block_size_factor = 1 if self.gpu_to_cpu else block_size_factor

225 self.dst_block_size_factor = block_size_factor if self.gpu_to_cpu else 1

226

231 - self.transfer_type = ("GPU", "CPU") if self.gpu_to_cpu else ("CPU", "GPU")

227 # mmap_region to clean up on shutdown (gpu_to_cpu handler owns it)

228 self._mmap_region = mmap_region

229 # job_id -> event

@@ -242,9 +237,9 @@ class SingleDirectionOffloadingHandler(OffloadingHandler):

237 # list of pinned descriptor buffer sets available for re-use

238 self._buffer_pool: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = []

239

245 - @override

246 - def transfer_async(self, job_id: int, transfer_spec: TransferSpec) -> bool:

247 - src_spec, dst_spec = transfer_spec

240 + def transfer_async(

241 + self, job_id: int, src_spec: LoadStoreSpec, dst_spec: LoadStoreSpec

242 + ) -> bool:

243 assert isinstance(src_spec, BlockIDsLoadStoreSpec)

244 assert isinstance(dst_spec, BlockIDsLoadStoreSpec)

245

@@ -425,7 +420,6 @@ class SingleDirectionOffloadingHandler(OffloadingHandler):

420 # success

421 return True

422

428 - @override

423 def get_finished(self) -> list[TransferResult]:

424 results: list[TransferResult] = []

425 while self._transfers and self._transfers[0].end_event.query():

@@ -438,7 +432,6 @@ class SingleDirectionOffloadingHandler(OffloadingHandler):

432 success=True,

433 transfer_size=transfer.num_bytes,

434 transfer_time=transfer_time,

441 - transfer_type=self.transfer_type,

435 )

436

437 results.append(result)

@@ -451,14 +444,12 @@ class SingleDirectionOffloadingHandler(OffloadingHandler):

444 del self._transfer_events[transfer.job_id]

445 return results

446

454 - @override

447 def wait(self, job_ids: set[int]):

448 for job_id in job_ids:

449 event = self._transfer_events.get(job_id)

450 if event is not None:

451 event.synchronize()

452

461 - @override

453 def shutdown(self) -> None:

454 while self._transfers:

455 transfer = self._transfers.popleft()

@@ -474,7 +465,14 @@ class SingleDirectionOffloadingHandler(OffloadingHandler):

465 self._mmap_region = None

466

467

477 -class CpuGpuOffloadingHandlers:

468 +class CPUOffloadingWorker(OffloadingWorker):

469 + """OffloadingWorker for CPU offloading.

470 +

471 + Composes two SingleDirectionOffloadingHandler instances (one for each

472 + direction) and exposes them through the explicit submit_store /

473 + submit_load API.

474 + """

475 +

476 def __init__(

477 self,

478 kv_caches: CanonicalKVCaches,

@@ -484,7 +482,6 @@ class CpuGpuOffloadingHandlers:

482 ):

483 pin_memory = PIN_MEMORY

484 logger.info("Allocating %d CPU tensors...", len(kv_caches.tensors))

487 - self._mmap_region = mmap_region

485 if mmap_region is not None and pin_memory:

486 pin_mmap_region(mmap_region)

487

@@ -518,7 +515,7 @@ class CpuGpuOffloadingHandlers:

515 gpu_tensors.append(gpu_tensor)

516 cpu_tensors.append(cpu_tensor)

517

521 - self.gpu_to_cpu_handler = SingleDirectionOffloadingHandler(

518 + self._store_handler = SingleDirectionOffloadingHandler(

519 gpu_tensors=gpu_tensors,

520 cpu_tensors=cpu_tensors,

521 block_size_factor=block_size_factor,

@@ -527,10 +524,33 @@ class CpuGpuOffloadingHandlers:

524 mmap_region=mmap_region,

525 )

526

530 - self.cpu_to_gpu_handler = SingleDirectionOffloadingHandler(

527 + self._load_handler = SingleDirectionOffloadingHandler(

528 gpu_tensors=gpu_tensors,

529 cpu_tensors=cpu_tensors,

530 block_size_factor=block_size_factor,

531 kv_cache_groups_data_refs=kv_caches.group_data_refs,

532 gpu_to_cpu=False,

533 )

534 +

535 + def submit_store(

536 + self, job_id: int, src_spec: GPULoadStoreSpec, dst_spec: LoadStoreSpec

537 + ) -> bool:

538 + """Async GPU -> CPU."""

539 + return self._store_handler.transfer_async(job_id, src_spec, dst_spec)

540 +

541 + def submit_load(

542 + self, job_id: int, src_spec: LoadStoreSpec, dst_spec: GPULoadStoreSpec

543 + ) -> bool:

544 + """Async CPU -> GPU."""

545 + return self._load_handler.transfer_async(job_id, src_spec, dst_spec)

546 +

547 + def get_finished(self) -> list[TransferResult]:

548 + return self._store_handler.get_finished() + self._load_handler.get_finished()

549 +

550 + def wait(self, job_ids: set[int]) -> None:

551 + self._store_handler.wait(job_ids)

552 + self._load_handler.wait(job_ids)

553 +

554 + def shutdown(self) -> None:

555 + self._store_handler.shutdown()

556 + self._load_handler.shutdown()

vllm/v1/kv_offload/cpu/spec.py

+11 -20

@@ -1,6 +1,5 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 -from collections.abc import Iterator

3 from typing import Any

5 from typing_extensions import override

@@ -11,21 +10,16 @@ from vllm.utils.math_utils import round_up

10 from vllm.v1.kv_cache_interface import KVCacheConfig

11 from vllm.v1.kv_offload.base import (

12 CanonicalKVCaches,

14 - GPULoadStoreSpec,

15 - LoadStoreSpec,

13 OffloadingCounterMetadata,

14 OffloadingGaugeMetadata,

15 OffloadingManager,

16 OffloadingMetricMetadata,

17 OffloadingSpec,

18 + OffloadingWorker,

19 )

22 -from vllm.v1.kv_offload.cpu.common import (

23 - CPULoadStoreSpec,

24 - CPUOffloadingMetrics,

25 -)

26 -from vllm.v1.kv_offload.cpu.gpu_worker import CpuGpuOffloadingHandlers

20 +from vllm.v1.kv_offload.cpu.common import CPUOffloadingMetrics

21 +from vllm.v1.kv_offload.cpu.gpu_worker import CPUOffloadingWorker

22 from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager

28 -from vllm.v1.kv_offload.worker.worker import OffloadingHandler

25 class CPUOffloadingSpec(OffloadingSpec):

@@ -107,7 +101,7 @@ class CPUOffloadingSpec(OffloadingSpec):

101 self._manager: OffloadingManager | None = None

102

103 # worker-side

110 - self._handlers: CpuGpuOffloadingHandlers | None = None

104 + self._worker: CPUOffloadingWorker | None = None

105

106 self.eviction_policy: str = self.extra_config.get("eviction_policy", "lru")

107

@@ -131,25 +125,22 @@ class CPUOffloadingSpec(OffloadingSpec):

125 )

126 return self._manager

127

134 - def create_handlers(self, kv_caches: CanonicalKVCaches) -> CpuGpuOffloadingHandlers:

135 - return CpuGpuOffloadingHandlers(

128 + def create_worker(self, kv_caches: CanonicalKVCaches) -> CPUOffloadingWorker:

129 + return CPUOffloadingWorker(

130 kv_caches=kv_caches,

131 block_size_factor=self.block_size_factor,

132 num_cpu_blocks=self.num_blocks,

133 )

134

135 @override

142 - def get_handlers(

143 - self, kv_caches: CanonicalKVCaches

144 - ) -> Iterator[tuple[type[LoadStoreSpec], type[LoadStoreSpec], OffloadingHandler]]:

145 - if not self._handlers:

136 + def get_worker(self, kv_caches: CanonicalKVCaches) -> OffloadingWorker:

137 + if not self._worker:

138 if not (current_platform.is_cuda_alike() or current_platform.is_xpu()):

139 raise Exception(

140 "CPU Offloading is currently only supported on CUDA-alike "

141 "and XPU GPUs"

142 )

151 - self._handlers = self.create_handlers(kv_caches)

143 + self._worker = self.create_worker(kv_caches)

144

153 - assert self._handlers is not None

154 - yield GPULoadStoreSpec, CPULoadStoreSpec, self._handlers.gpu_to_cpu_handler

155 - yield CPULoadStoreSpec, GPULoadStoreSpec, self._handlers.cpu_to_gpu_handler

145 + assert self._worker is not None

146 + return self._worker

vllm/v1/kv_offload/tiering/spec.py

+5 -5

@@ -44,7 +44,7 @@ from vllm.v1.kv_offload.base import (

44 OffloadingManager,

45 OffloadingMetricMetadata,

46 )

47 -from vllm.v1.kv_offload.cpu.gpu_worker import CpuGpuOffloadingHandlers

47 +from vllm.v1.kv_offload.cpu.gpu_worker import CPUOffloadingWorker

48 from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion

49 from vllm.v1.kv_offload.cpu.spec import CPUOffloadingSpec

50 from vllm.v1.kv_offload.tiering.factory import SecondaryTierFactory

@@ -163,8 +163,8 @@ class TieringOffloadingSpec(CPUOffloadingSpec):

163 raise

164

165 # Create TieringOffloadingManager. GPU↔CPU transfers use the inherited

166 - # get_handlers(); secondary tier transfers are handled by the

167 - # secondary tier managers and need no additional handlers here.

166 + # get_worker(). Secondary tier transfers are handled by the

167 + # secondary tier managers and need no additional workers here.

168 tiering_manager = TieringOffloadingManager(

169 primary_tier=primary_tier,

170 secondary_tiers=secondary_tiers,

@@ -187,7 +187,7 @@ class TieringOffloadingSpec(CPUOffloadingSpec):

187 return self._manager

188

189 @override

190 - def create_handlers(self, kv_caches: CanonicalKVCaches) -> CpuGpuOffloadingHandlers:

190 + def create_worker(self, kv_caches: CanonicalKVCaches) -> CPUOffloadingWorker:

191 rank = torch.accelerator.current_device_index()

192 worker_mmap = SharedOffloadRegion(

193 instance_id=self.vllm_config.instance_id,

@@ -196,7 +196,7 @@ class TieringOffloadingSpec(CPUOffloadingSpec):

196 kv_bytes_per_block=self.kv_bytes_per_offloaded_block,

197 cpu_page_size=self.cpu_page_size_per_worker,

198 )

199 - return CpuGpuOffloadingHandlers(

199 + return CPUOffloadingWorker(

200 kv_caches=kv_caches,

201 block_size_factor=self.block_size_factor,

202 num_cpu_blocks=self.num_blocks,

vllm/v1/kv_offload/worker/__init__.py

vllm/v1/kv_offload/worker/worker.py deleted

-176

@@ -1,176 +0,0 @@

1 -# SPDX-License-Identifier: Apache-2.0

2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 -from abc import ABC, abstractmethod

4 -from dataclasses import dataclass

5 -

6 -from vllm.logger import init_logger

7 -from vllm.v1.kv_offload.base import LoadStoreSpec

8 -

9 -# a single transfer spec (src_blocks_spec, dst_blocks_spec)

10 -TransferSpec = tuple[LoadStoreSpec, LoadStoreSpec]

11 -# transfers are forwarded to workers by (src_medium, dst_medium)

12 -TransferType = tuple[str, str]

13 -

14 -logger = init_logger(__name__)

15 -

16 -

17 -@dataclass

18 -class TransferResult:

19 - job_id: int

20 - success: bool

21 - transfer_size: int | None = None # Size in bytes

22 - transfer_time: float | None = None

23 - transfer_type: TransferType | None = None

24 -

25 -

26 -class OffloadingHandler(ABC):

27 - """

28 - OffloadingHandler class for managing asynchronous KV data transfers

29 -

30 - This class runs in the worker.

31 - It kicks off async KV data transfer requests, and allows

32 - collecting back completion statuses.

33 -

34 - The class provides the following primitives:

35 - transfer_async() - kicks off a new transfer job

36 - get_finished() - returns a list of newly finished job IDs.

37 - """

38 -

39 - @abstractmethod

40 - def transfer_async(self, job_id: int, spec: TransferSpec) -> bool:

41 - """

42 - Initiates an asynchronous transfer of KV data.

43 -

44 - Args:

45 - job_id: a unique ID that will be used when notifying back on

46 - transfer completion.

47 - spec: the (src, dst) spec of the KV data transfer.

48 -

49 - Returns:

50 - True if transfer was submitted successfully.

51 - """

52 - pass

53 -

54 - @abstractmethod

55 - def get_finished(self) -> list[TransferResult]:

56 - """

57 - Get transfers finished since last call.

58 -

59 - Returns:

60 - A list of (job_id, success) of transfers.

61 - """

62 - pass

63 -

64 - @abstractmethod

65 - def wait(self, job_ids: set[int]) -> None:

66 - """

67 - Wait for jobs to finish (blocking).

68 - Args:

69 - job_ids: The set of job IDs to wait for.

70 - """

71 -

72 - def shutdown(self) -> None:

73 - """Shutdown the handler and release any resources."""

74 - return

75 -

76 -

77 -class OffloadingWorker:

78 - """

79 - OffloadingWorker class for managing asynchronous KV data transfers

80 - using multiple OffloadingHandlers

81 -

82 - This class runs in the worker.

83 - It kicks off async KV data transfer requests, by delegating

84 - to one of its registered OffloadingHandlers, based on the transfer type.

85 -

86 - The class provides the following primitives:

87 - register_handler() - registers a new handler to handle

88 - a specific transfer type

89 - transfer_async() - kicks off a new transfer job

90 - using one of the registered handlers.

91 - get_finished() - returns a list of newly finished job IDs

92 - from all handlers.

93 - """

94 -

95 - def __init__(self):

96 - self.handlers: set[OffloadingHandler] = set()

97 - self.transfer_type_to_handler: dict[TransferType, OffloadingHandler] = {}

98 -

99 - def register_handler(

100 - self,

101 - src_cls: type[LoadStoreSpec],

102 - dst_cls: type[LoadStoreSpec],

103 - handler: OffloadingHandler,

104 - ) -> None:

105 - """

106 - Registers a new handler.

107 -

108 - Args:

109 - src_cls: the source type of transfers handled by this handler.

110 - dst_cls: the destination type of transfers handled by this handler.

111 - handler: the handler that will handle transfers.

112 - """

113 - transfer_type = (src_cls.medium(), dst_cls.medium())

114 - assert transfer_type not in self.transfer_type_to_handler

115 - self.handlers.add(handler)

116 - self.transfer_type_to_handler[transfer_type] = handler

117 -

118 - def transfer_async(self, job_id: int, spec: TransferSpec) -> bool:

119 - """

120 - Initiates an asynchronous transfer of KV data.

121 -

122 - Args:

123 - job_id: a unique ID that will be used when notifying back on

124 - transfer completion.

125 - spec: the (src, dst) spec of the KV data transfer.

126 -

127 - Returns:

128 - True if transfer was submitted successfully.

129 - """

130 - src, dst = spec

131 - transfer_type = (src.medium(), dst.medium())

132 - handler = self.transfer_type_to_handler.get(transfer_type)

133 - assert handler is not None

134 - try:

135 - success = handler.transfer_async(job_id, spec)

136 - except Exception as e:

137 - logger.warning(

138 - "Exception in %r transfer %d: %r",

139 - transfer_type,

140 - job_id,

141 - e,

142 - exc_info=True,

143 - )

144 - return False

145 -

146 - if not success:

147 - logger.warning("Failed to submit %r transfer %d", transfer_type, job_id)

148 - else:

149 - logger.debug("Submitted %r transfer %d: %r", transfer_type, job_id, spec)

150 - return success

151 -

152 - def get_finished(self) -> list[TransferResult]:

153 - """

154 - Get transfers finished since last call.

155 -

156 - Returns:

157 - A list of TransferResults

158 - """

159 - finished = []

160 - for handler in self.handlers:

161 - finished.extend(handler.get_finished())

162 - return finished

163 -

164 - def wait(self, job_ids: set[int]) -> None:

165 - """

166 - Wait for jobs to finish (blocking).

167 -

168 - Args:

169 - job_ids: The set of job IDs to wait for.

170 - """

171 - for handler in self.handlers:

172 - handler.wait(job_ids)

173 -

174 - def shutdown(self) -> None:

175 - for handler in self.handlers:

176 - handler.shutdown()