@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 93da882e73

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 2, 06:00 UTC

[kv_offload] Add `@override` decorators to subclass method implementations (#44177)

Signed-off-by: Ronen Schaffer <ronen.schaffer@ibm.com>

Ronen Schaffer committed Jun 2, 2026 at 11:07 UTC 93da882e73787fb7c4d5e122ea5c48ff3ad1e0c3

10 files changed +68

vllm/v1/kv_offload/base.py

@@ -12,6 +12,7 @@ from typing import TYPE_CHECKING, Any, NewType

13 import numpy as np

14 import torch

15 +from typing_extensions import override

17 from vllm.logger import init_logger

18 from vllm.v1.core.kv_cache_utils import resolve_kv_cache_block_sizes

@@ -311,6 +312,7 @@ class GPULoadStoreSpec(BlockIDsLoadStoreSpec):

312 self.block_indices: Sequence[int] = block_indices

313

314 @staticmethod

315 + @override

316 def medium() -> str:

317 return "GPU"

318

vllm/v1/kv_offload/cpu/common.py

@@ -1,5 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

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

3 +from typing_extensions import override

4 +

5 from vllm.v1.kv_offload.base import BlockIDsLoadStoreSpec

@@ -9,5 +11,6 @@ class CPULoadStoreSpec(BlockIDsLoadStoreSpec):

11 """

13 @staticmethod

14 + @override

15 def medium() -> str:

16 return "CPU"

vllm/v1/kv_offload/cpu/gpu_worker.py

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

7 import numpy as np

8 import torch

9 +from typing_extensions import override

11 from vllm import _custom_ops as ops

12 from vllm.logger import init_logger

@@ -179,6 +180,7 @@ class SingleDirectionOffloadingHandler(OffloadingHandler):

180 # list of CUDA events available for re-use

181 self._event_pool: list[torch.Event] = []

182

183 + @override

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

185 src_spec, dst_spec = transfer_spec

186 assert isinstance(src_spec, BlockIDsLoadStoreSpec)

@@ -348,6 +350,7 @@ class SingleDirectionOffloadingHandler(OffloadingHandler):

350 # success

351 return True

352

353 + @override

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

355 results: list[TransferResult] = []

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

@@ -370,12 +373,14 @@ class SingleDirectionOffloadingHandler(OffloadingHandler):

373 del self._transfer_events[transfer.job_id]

374 return results

375

376 + @override

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

378 for job_id in job_ids:

379 event = self._transfer_events.get(job_id)

380 if event is not None:

381 event.synchronize()

382

383 + @override

384 def shutdown(self) -> None:

385 while self._transfers:

386 transfer = self._transfers.popleft()

vllm/v1/kv_offload/cpu/manager.py

+11

@@ -4,6 +4,8 @@ from collections import OrderedDict

4 from collections.abc import Collection, Iterable

5 from typing import Literal

7 +from typing_extensions import override

8 +

9 from vllm.v1.kv_offload.base import (

10 LoadStoreSpec,

11 OffloadingEvent,

@@ -95,9 +97,11 @@ class CPUOffloadingManager(OffloadingManager):

98 # --- OffloadingManager interface ---

100 + @override

101 def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext:

102 return RequestOffloadingContext()

103

104 + @override

105 def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None:

106 if self.counts is not None:

107 if key in self.counts:

@@ -114,6 +118,7 @@ class CPUOffloadingManager(OffloadingManager):

118 return None # write in-flight; caller should retry

119 return True

120

121 + @override

122 def prepare_load(

123 self,

124 keys: Collection[OffloadKey],

@@ -128,9 +133,11 @@ class CPUOffloadingManager(OffloadingManager):

133 blocks.append(block)

134 return self._get_load_store_spec(keys, blocks)

135

136 + @override

137 def touch(self, keys: Collection[OffloadKey], req_context: ReqContext) -> None:

138 self._policy.touch(keys)

139

140 + @override

141 def complete_load(

142 self, keys: Collection[OffloadKey], req_context: ReqContext

143 ) -> None:

@@ -140,6 +147,7 @@ class CPUOffloadingManager(OffloadingManager):

147 assert block.ref_cnt > 0, f"Block {key!r} ref_cnt is already 0"

148 block.ref_cnt -= 1

149

150 + @override

151 def prepare_store(

152 self,

153 keys: Collection[OffloadKey],

@@ -197,6 +205,7 @@ class CPUOffloadingManager(OffloadingManager):

205 evicted_keys=to_evict,

206 )

207

208 + @override

209 def complete_store(

210 self,

211 keys: Collection[OffloadKey],

@@ -227,6 +236,7 @@ class CPUOffloadingManager(OffloadingManager):

236 )

237 )

238

239 + @override

240 def reset_cache(self) -> None:

241 # Clear ALL blocks unconditionally. The scheduler's _stale_job_threshold

242 # guarantees that complete_load / complete_store are never called for

@@ -238,6 +248,7 @@ class CPUOffloadingManager(OffloadingManager):

248 self._free_list.clear()

249 self._num_allocated_blocks = 0

250

251 + @override

252 def take_events(self) -> Iterable[OffloadingEvent]:

253 if self.events is not None:

254 yield from self.events

vllm/v1/kv_offload/cpu/policies/arc.py

@@ -3,6 +3,8 @@

3 from collections import OrderedDict

4 from collections.abc import Iterable

6 +from typing_extensions import override

7 +

8 from vllm.v1.kv_offload.base import OffloadKey

9 from vllm.v1.kv_offload.cpu.policies.base import BlockStatus, CachePolicy

@@ -54,18 +56,22 @@ class ARCCachePolicy(CachePolicy):

56 self.b1: OrderedDict[OffloadKey, None] = OrderedDict()

57 self.b2: OrderedDict[OffloadKey, None] = OrderedDict()

59 + @override

60 def get(self, key: OffloadKey) -> BlockStatus | None:

61 return self.t1.get(key) or self.t2.get(key)

63 + @override

64 def insert(self, key: OffloadKey, block: BlockStatus) -> None:

65 self.t1[key] = block

66 self.b1.pop(key, None)

67 self.b2.pop(key, None)

69 + @override

70 def remove(self, key: OffloadKey) -> None:

71 if self.t1.pop(key, None) is None:

72 self.t2.pop(key, None)

74 + @override

75 def touch(self, keys: Iterable[OffloadKey]) -> None:

76 for key in reversed(list(keys)):

77 if key in self.t1:

@@ -94,6 +100,7 @@ class ARCCachePolicy(CachePolicy):

100 # move to MRU position (end) to keep it fresh in the ghost list

101 self.b2.move_to_end(key)

102

103 + @override

104 def clear(self) -> None:

105 self.t1.clear()

106 self.t2.clear()

@@ -101,6 +108,7 @@ class ARCCachePolicy(CachePolicy):

108 self.b2.clear()

109 self.target_t1_size = 0.0

110

111 + @override

112 def evict(

113 self, n: int, protected: set[OffloadKey]

114 ) -> list[tuple[OffloadKey, BlockStatus]] | None:

vllm/v1/kv_offload/cpu/policies/lru.py

@@ -3,6 +3,8 @@

3 from collections import OrderedDict

4 from collections.abc import Iterable

6 +from typing_extensions import override

7 +

8 from vllm.v1.kv_offload.base import OffloadKey

9 from vllm.v1.kv_offload.cpu.policies.base import BlockStatus, CachePolicy

@@ -14,23 +16,29 @@ class LRUCachePolicy(CachePolicy):

16 # cache_capacity unused by LRU but accepted for a uniform constructor

17 self.blocks: OrderedDict[OffloadKey, BlockStatus] = OrderedDict()

19 + @override

20 def get(self, key: OffloadKey) -> BlockStatus | None:

21 return self.blocks.get(key)

23 + @override

24 def insert(self, key: OffloadKey, block: BlockStatus) -> None:

25 self.blocks[key] = block

27 + @override

28 def remove(self, key: OffloadKey) -> None:

29 del self.blocks[key]

31 + @override

32 def touch(self, keys: Iterable[OffloadKey]) -> None:

33 for key in reversed(list(keys)):

34 if key in self.blocks:

35 self.blocks.move_to_end(key)

37 + @override

38 def clear(self) -> None:

39 self.blocks.clear()

41 + @override

42 def evict(

43 self, n: int, protected: set[OffloadKey]

44 ) -> list[tuple[OffloadKey, BlockStatus]] | None:

vllm/v1/kv_offload/cpu/spec.py

@@ -2,6 +2,8 @@

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

3 from collections.abc import Iterator

5 +from typing_extensions import override

6 +

7 from vllm.config import VllmConfig

8 from vllm.platforms import current_platform

9 from vllm.v1.kv_cache_interface import KVCacheConfig

@@ -57,6 +59,7 @@ class CPUOffloadingSpec(OffloadingSpec):

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

62 + @override

63 def get_manager(self) -> OffloadingManager:

64 if not self._manager:

65 kv_events_config = self.vllm_config.kv_events_config

@@ -88,6 +91,7 @@ class CPUOffloadingSpec(OffloadingSpec):

91 num_cpu_blocks=self.num_blocks,

92 )

94 + @override

95 def get_handlers(

96 self, kv_caches: CanonicalKVCaches

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

vllm/v1/kv_offload/tiering/example/manager.py

@@ -13,6 +13,8 @@ import logging

13 from collections.abc import Iterable

14 from typing import TYPE_CHECKING

16 +from typing_extensions import override

17 +

18 from vllm.v1.kv_offload.base import OffloadKey, ReqContext, RequestOffloadingContext

19 from vllm.v1.kv_offload.tiering.base import (

20 JobMetadata,

@@ -64,6 +66,7 @@ class ExampleSecondaryTierManager(SecondaryTierManager):

66 # Completed jobs waiting to be retrieved by get_finished_jobs()

67 self.completed_jobs: list[JobResult] = []

69 + @override

70 def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None:

71 """

72 Check whether a block exists in this secondary tier.

@@ -77,6 +80,7 @@ class ExampleSecondaryTierManager(SecondaryTierManager):

80 """

81 return key in self.blocks

83 + @override

84 def submit_store(self, job_metadata: JobMetadata) -> None:

85 """

86 Submit a job to store blocks from primary tier to this tier.

@@ -96,6 +100,7 @@ class ExampleSecondaryTierManager(SecondaryTierManager):

100 self.blocks[key] = True

101 self.completed_jobs.append(JobResult(job_id=job_metadata.job_id, success=True))

102

103 + @override

104 def submit_load(self, job_metadata: JobMetadata) -> None:

105 """

106 Submit a job to load blocks from this tier to primary tier.

@@ -120,6 +125,7 @@ class ExampleSecondaryTierManager(SecondaryTierManager):

125

126 self.completed_jobs.append(JobResult(job_id=job_metadata.job_id, success=True))

127

128 + @override

129 def get_finished_jobs(self) -> Iterable[JobResult]:

130 """

131 Poll for finished jobs.

@@ -132,6 +138,7 @@ class ExampleSecondaryTierManager(SecondaryTierManager):

138 self.completed_jobs = []

139 return result

140

141 + @override

142 def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext:

143 return RequestOffloadingContext()

144

vllm/v1/kv_offload/tiering/fs/manager.py

@@ -21,6 +21,8 @@ import os

21 from collections.abc import Iterable

22 from typing import TYPE_CHECKING

24 +from typing_extensions import override

25 +

26 from vllm.logger import init_logger

27 from vllm.v1.kv_offload.base import OffloadKey, ReqContext

28 from vllm.v1.kv_offload.file_mapper import FileMapper

@@ -101,14 +103,17 @@ class FileSystemTierManager(SecondaryTierManager):

103 thread_name_prefix="vllm_kv_py_fs",

104 )

105

106 + @override

107 def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext:

108 return RequestOffloadingContext()

109

110 + @override

111 def lookup(

112 self, key: OffloadKey, req_context: ReqContext | None = None

113 ) -> bool | None:

114 return os.path.exists(self.file_mapper.get_file_name(key))

115

116 + @override

117 def submit_store(self, job_metadata: JobMetadata) -> None:

118 tasks = (

119 functools.partial(

@@ -122,6 +127,7 @@ class FileSystemTierManager(SecondaryTierManager):

127 )

128 self._pool.enqueue_store(job_metadata.job_id, len(job_metadata.keys), tasks)

129

130 + @override

131 def submit_load(self, job_metadata: JobMetadata) -> None:

132 tasks = (

133 functools.partial(

@@ -135,6 +141,7 @@ class FileSystemTierManager(SecondaryTierManager):

141 )

142 self._pool.enqueue_load(job_metadata.job_id, len(job_metadata.keys), tasks)

143

144 + @override

145 def get_finished_jobs(self) -> Iterable[JobResult]:

146 """

147 Collect completed jobs from the finished-jobs queue.

@@ -144,6 +151,7 @@ class FileSystemTierManager(SecondaryTierManager):

151 for job_id, success in self._pool.get_finished()

152 )

153

154 + @override

155 def shutdown(self) -> None:

156 """

157 Release resources held by this tier.

vllm/v1/kv_offload/tiering/manager.py

+12

@@ -25,6 +25,7 @@ from collections.abc import Collection, Iterable, Sequence

25 from dataclasses import dataclass, field

27 import numpy as np

28 +from typing_extensions import override

30 from vllm.logger import init_logger

31 from vllm.v1.kv_offload.base import (

@@ -100,6 +101,7 @@ class CPUPrimaryTierOffloadingManager(CPUOffloadingManager):

101 """

102 return self._kv_memoryview

103

104 + @override

105 def shutdown(self) -> None:

106 super().shutdown()

107 self._kv_memoryview.release()

@@ -222,6 +224,7 @@ class TieringOffloadingManager(OffloadingManager):

224 job_metadata.keys, job_metadata.req_context

225 )

226

227 + @override

228 def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None:

229 """

230 Check whether a single block is offloaded and ready.

@@ -338,6 +341,7 @@ class TieringOffloadingManager(OffloadingManager):

341

342 self._pending_load_submissions.clear()

343

344 + @override

345 def prepare_load(

346 self, keys: Collection[OffloadKey], req_context: ReqContext

347 ) -> LoadStoreSpec:

@@ -362,6 +366,7 @@ class TieringOffloadingManager(OffloadingManager):

366

367 return self.primary_tier.prepare_load(keys, req_context)

368

369 + @override

370 def touch(self, keys: Collection[OffloadKey], req_context: ReqContext):

371 """

372 Mark blocks as recently used in all tiers.

@@ -374,6 +379,7 @@ class TieringOffloadingManager(OffloadingManager):

379 for tier in self.secondary_tiers:

380 tier.touch(keys, req_context)

381

382 + @override

383 def complete_load(self, keys: Collection[OffloadKey], req_context: ReqContext):

384 """

385 Mark blocks as done loading from primary tier to GPU.

@@ -387,6 +393,7 @@ class TieringOffloadingManager(OffloadingManager):

393 """

394 self.primary_tier.complete_load(keys, req_context)

395

396 + @override

397 def prepare_store(

398 self, keys: Collection[OffloadKey], req_context: ReqContext

399 ) -> PrepareStoreOutput | None:

@@ -470,6 +477,7 @@ class TieringOffloadingManager(OffloadingManager):

477 self._transfer_jobs[job_id] = job_metadata

478 tier.submit_store(job_metadata)

479

480 + @override

481 def complete_store(

482 self,

483 keys: Collection[OffloadKey],

@@ -528,6 +536,7 @@ class TieringOffloadingManager(OffloadingManager):

536 # Note: The async transfers are now in flight. Their completion is

537 # tracked via get_finished_jobs() / _maybe_process_finished_jobs().

538

539 + @override

540 def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext:

541 """

542 Query each secondary tier for its offload policy preference.

@@ -547,12 +556,14 @@ class TieringOffloadingManager(OffloadingManager):

556 )

557 return RequestOffloadingContext(policy=policy)

558

559 + @override

560 def on_request_finished(self, req_context: ReqContext) -> None:

561 self.primary_tier.on_request_finished(req_context)

562 for tier in self.secondary_tiers:

563 tier.on_request_finished(req_context)

564 self._request_level_tiers.pop(req_context.req_id, None)

565

566 + @override

567 def take_events(self) -> Iterable[OffloadingEvent]:

568 """

569 End-of-step hook: flush deferred work, yield events, reset per-step state.

@@ -585,6 +596,7 @@ class TieringOffloadingManager(OffloadingManager):

596

597 yield from self.primary_tier.take_events()

598

599 + @override

600 def shutdown(self) -> None:

601 """Shutdown all tiers and release resources."""

602 for tier in self.secondary_tiers: