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

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 4, 11:15 UTC

[Refactor][kv_offload] KV Offloading maintainability improvements (#40538)

Signed-off-by: Martin Hickey <martin.hickey@ie.ibm.com> Co-authored-by: Or Ozeri <oro@il.ibm.com>

Martin Hickey committed Apr 30, 2026 at 03:55 UTC c42981d034713da814913ebd2e53269346f3ecea
24 files changed +446 -440
tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py
+1 -1
@@ -15,7 +15,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.scheduler import (
15 OffloadingConnectorScheduler,
16 )
17 from vllm.v1.core.kv_cache_utils import BlockHash
18 -from vllm.v1.kv_offload.abstract import (
18 +from vllm.v1.kv_offload.base import (
19 OffloadingEvent,
20 OffloadingManager,
21 ReqContext,
tests/v1/kv_connector/unit/offloading_connector/test_worker.py
+1 -1
@@ -20,7 +20,7 @@ from vllm.v1.kv_cache_interface import (
20 MLAAttentionSpec,
21 UniformTypeKVCacheSpecs,
22 )
23 -from vllm.v1.kv_offload.spec import (
23 +from vllm.v1.kv_offload.base import (
24 CanonicalKVCacheRef,
25 CanonicalKVCaches,
26 OffloadingSpec,
tests/v1/kv_connector/unit/offloading_connector/utils.py
+3 -3
@@ -37,15 +37,15 @@ from vllm.v1.kv_cache_interface import (
37 KVCacheConfig,
38 KVCacheGroupSpec,
39 )
40 -from vllm.v1.kv_offload.abstract import (
40 +from vllm.v1.kv_offload.base import (
41 + GPULoadStoreSpec,
42 LoadStoreSpec,
43 OffloadingManager,
44 + OffloadingSpec,
45 OffloadKey,
46 PrepareStoreOutput,
47 make_offload_key,
48 )
47 -from vllm.v1.kv_offload.mediums import GPULoadStoreSpec
48 -from vllm.v1.kv_offload.spec import OffloadingSpec
49 from vllm.v1.kv_offload.worker.worker import (
50 OffloadingHandler,
51 TransferResult,
tests/v1/kv_connector/unit/test_offloading_connector.py renamed
tests/v1/kv_offload/cpu/test_gpu_worker.py renamed
+5 -4
@@ -9,14 +9,15 @@ import torch
9
10 from vllm.platforms import current_platform
11 from vllm.utils.torch_utils import set_random_seed
12 -from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion
13 -from vllm.v1.kv_offload.mediums import CPULoadStoreSpec, GPULoadStoreSpec
14 -from vllm.v1.kv_offload.spec import (
12 +from vllm.v1.kv_offload.base import (
13 CanonicalKVCacheRef,
14 CanonicalKVCaches,
15 CanonicalKVCacheTensor,
16 + GPULoadStoreSpec,
17 )
19 -from vllm.v1.kv_offload.worker.cpu_gpu import CpuGpuOffloadingHandlers
18 +from vllm.v1.kv_offload.cpu.common import CPULoadStoreSpec
19 +from vllm.v1.kv_offload.cpu.gpu_worker import CpuGpuOffloadingHandlers
20 +from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion
21
22 NUM_GPU_BLOCKS = [64]
23 NUM_CPU_BLOCKS = [256]
tests/v1/kv_offload/cpu/test_manager.py renamed
+2 -2
@@ -6,7 +6,7 @@ from dataclasses import dataclass
6 import numpy as np
7 import pytest
8
9 -from vllm.v1.kv_offload.abstract import (
9 +from vllm.v1.kv_offload.base import (
10 LoadStoreSpec,
11 OffloadingEvent,
12 OffloadKey,
@@ -14,9 +14,9 @@ from vllm.v1.kv_offload.abstract import (
14 ReqContext,
15 make_offload_key,
16 )
17 +from vllm.v1.kv_offload.cpu.common import CPULoadStoreSpec
18 from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager
19 from vllm.v1.kv_offload.cpu.policies.arc import ARCCachePolicy
19 -from vllm.v1.kv_offload.mediums import CPULoadStoreSpec
20 from vllm.v1.kv_offload.reuse_manager import FilterReusedOffloadingManager
21
22
tests/v1/kv_offload/cpu/test_shared_offload_region.py renamed
tests/v1/kv_offload/test_worker.py
+1 -1
@@ -1,6 +1,6 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -from vllm.v1.kv_offload.abstract import LoadStoreSpec
3 +from vllm.v1.kv_offload.base import LoadStoreSpec
4 from vllm.v1.kv_offload.worker.worker import (
5 OffloadingHandler,
6 OffloadingWorker,
vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py
+3 -3
@@ -18,15 +18,15 @@ from vllm.logger import init_logger
18 from vllm.utils.math_utils import cdiv
19 from vllm.v1.core.kv_cache_manager import KVCacheBlocks
20 from vllm.v1.core.sched.output import SchedulerOutput
21 -from vllm.v1.kv_offload.abstract import (
21 +from vllm.v1.kv_offload.base import (
22 + GPULoadStoreSpec,
23 OffloadingManager,
24 + OffloadingSpec,
25 OffloadKey,
26 ReqContext,
27 get_offload_block_hash,
28 make_offload_key,
29 )
28 -from vllm.v1.kv_offload.mediums import GPULoadStoreSpec
29 -from vllm.v1.kv_offload.spec import OffloadingSpec
30 from vllm.v1.outputs import KVConnectorOutput
31 from vllm.v1.request import Request
32
vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py
+1 -1
@@ -25,7 +25,7 @@ from vllm.v1.kv_cache_interface import (
25 MambaSpec,
26 UniformTypeKVCacheSpecs,
27 )
28 -from vllm.v1.kv_offload.spec import (
28 +from vllm.v1.kv_offload.base import (
29 CanonicalKVCacheRef,
30 CanonicalKVCaches,
31 CanonicalKVCacheTensor,
vllm/v1/kv_offload/abstract.py deleted
-197
@@ -1,197 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -"""
4 -OffloadingManager class for managing KV data offloading in vLLM v1
5 -
6 -This class runs in the scheduler, tracks which blocks are offloaded
7 -and their address.
8 -
9 -The class provides the following primitives:
10 - lookup() - check whether a single block is offloaded and ready.
11 - prepare_load() - prepare given blocks to be read.
12 - The given blocks will be protected from eviction.
13 - This function returns a LoadSpec which encapsulates
14 - information required for performing the load.
15 - touch() - marks the give blocks as recently used. Can be used
16 - to track block's LRU. This function is separated from the
17 - prepare_load function to allow setting block recency even
18 - for blocks which do not need reading from the cache, such as
19 - blocks that are cached by the GPU prefix cache.
20 - complete_load() - mark blocks which were previously prepared to be
21 - loaded as done loading. This is to re-allow their eviction.
22 - prepare_store() - prepare the given blocks to be written.
23 - Returns a StoreSpec encapsulating offloading information,
24 - as well as a list of blocks that were evicted as a result.
25 - complete_store() - marks a previous store as completed.
26 - Following this call, the given blocks will become loadable.
27 -"""
28 -
29 -from abc import ABC, abstractmethod
30 -from collections.abc import Iterable, Sequence
31 -from dataclasses import dataclass
32 -from typing import Any, NewType
33 -
34 -# `OffloadKey` identifies an offloaded block. It combines a block hash with
35 -# its KV cache group index, encoded as raw bytes to avoid tuple GC overhead.
36 -# Use the helper functions below to construct / decompose keys.
37 -OffloadKey = NewType("OffloadKey", bytes)
38 -
39 -
40 -def make_offload_key(block_hash: bytes, group_idx: int) -> OffloadKey:
41 - """Pack a block hash and group index into an `OffloadKey`."""
42 - return OffloadKey(block_hash + group_idx.to_bytes(4, "big", signed=False))
43 -
44 -
45 -def get_offload_block_hash(key: OffloadKey) -> bytes:
46 - """Extract the block hash from an `OffloadKey`."""
47 - return key[:-4]
48 -
49 -
50 -def get_offload_group_idx(key: OffloadKey) -> int:
51 - """Extract the group index from an `OffloadKey`."""
52 - return int.from_bytes(key[-4:], "big", signed=False)
53 -
54 -
55 -@dataclass
56 -class ReqContext:
57 - kv_transfer_params: dict[str, Any] | None = None
58 -
59 -
60 -class LoadStoreSpec(ABC):
61 - """
62 - Abstract metadata that encapsulates information allowing a worker
63 - to load, and optionally also to store, blocks of KV data.
64 - """
65 -
66 - @staticmethod
67 - @abstractmethod
68 - def medium() -> str:
69 - """
70 - Returns a string representation of the medium type
71 - this store/load targets.
72 - """
73 - pass
74 -
75 -
76 -@dataclass
77 -class PrepareStoreOutput:
78 - keys_to_store: list[OffloadKey]
79 - store_spec: LoadStoreSpec
80 - evicted_keys: list[OffloadKey]
81 -
82 -
83 -@dataclass
84 -class OffloadingEvent:
85 - keys: list[OffloadKey]
86 - medium: str
87 - # True if blocks are removed, False if stored
88 - removed: bool
89 -
90 -
91 -class OffloadingManager(ABC):
92 - @abstractmethod
93 - def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None:
94 - """
95 - Checks whether a single block is offloaded and ready to be read.
96 -
97 - Args:
98 - key: the key identifying the block to lookup.
99 - req_context: per-request context (e.g. kv_transfer_params).
100 -
101 - Returns:
102 - True if the block is offloaded and ready, False if not,
103 - or None if the lookup should be retried later.
104 - Returning None will delay the request handling by the vLLM
105 - scheduler.
106 - """
107 - pass
108 -
109 - @abstractmethod
110 - def prepare_load(
111 - self,
112 - keys: Sequence[OffloadKey],
113 - req_context: ReqContext,
114 - ) -> LoadStoreSpec:
115 - """
116 - Prepare the given blocks to be read.
117 - The given blocks will be protected from eviction until
118 - complete_load is called.
119 - It assumes all given blocks are offloaded.
120 -
121 - Args:
122 - keys: the keys identifying the blocks.
123 - req_context: per-request context (e.g. kv_transfer_params).
124 -
125 - Returns:
126 - A LoadStoreSpec that can be used by a worker to locate and load
127 - the actual offloaded KV data.
128 - """
129 - pass
130 -
131 - def touch(self, keys: Sequence[OffloadKey]):
132 - """
133 - Mark the given blocks as recently used.
134 - This could in practice mean moving them to the end of an LRU list.
135 -
136 - Args:
137 - keys: the keys identifying the blocks.
138 - """
139 - return
140 -
141 - def complete_load(self, keys: Iterable[OffloadKey]):
142 - """
143 - Marks previous blocks that were prepared to load as done loading.
144 -
145 - Args:
146 - keys: the keys identifying the blocks.
147 - """
148 - return
149 -
150 - @abstractmethod
151 - def prepare_store(
152 - self,
153 - keys: Sequence[OffloadKey],
154 - req_context: ReqContext,
155 - ) -> PrepareStoreOutput | None:
156 - """
157 - Prepare the given blocks to be offloaded.
158 - The given blocks will be protected from eviction until
159 - complete_store is called.
160 -
161 - Args:
162 - keys: the keys identifying the blocks.
163 - req_context: per-request context (e.g. kv_transfer_params).
164 -
165 - Returns:
166 - A PrepareStoreOutput indicating which blocks need storing,
167 - where to store them (LoadStoreSpec), and list of blocks that
168 - were evicted as a result.
169 - None is returned if the blocks cannot be stored.
170 - """
171 - pass
172 -
173 - def complete_store(self, keys: Iterable[OffloadKey], success: bool = True):
174 - """
175 - Marks blocks which were previously prepared to be stored, as stored.
176 - Following this call, the blocks become loadable.
177 - If if_success is False, blocks that were not marked as stored will be
178 - removed.
179 -
180 - Args:
181 - keys: the keys identifying the blocks.
182 - success: whether the blocks were stored successfully.
183 - """
184 - return
185 -
186 - def take_events(self) -> Iterable[OffloadingEvent]:
187 - """
188 - Take the offloading events from the manager.
189 -
190 - Yields:
191 - New OffloadingEvents collected since the last call.
192 - """
193 - return ()
194 -
195 - def shutdown(self) -> None:
196 - """Shutdown the manager and release any resources."""
197 - return
vllm/v1/kv_offload/base.py new
+390
@@ -0,0 +1,390 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +Core abstractions for KV cache offloading in vLLM v1.
5 +"""
6 +
7 +from __future__ import annotations
8 +
9 +from abc import ABC, abstractmethod
10 +from collections.abc import Iterable, Iterator, Sequence
11 +from dataclasses import dataclass
12 +from typing import TYPE_CHECKING, Any, NewType
13 +
14 +import numpy as np
15 +import torch
16 +
17 +from vllm.logger import init_logger
18 +
19 +if TYPE_CHECKING:
20 + from vllm.config import VllmConfig
21 + from vllm.v1.kv_cache_interface import KVCacheConfig
22 + from vllm.v1.kv_offload.worker.worker import OffloadingHandler
23 +
24 +# `OffloadKey` identifies an offloaded block. It combines a block hash with
25 +# its KV cache group index, encoded as raw bytes to avoid tuple GC overhead.
26 +# Use the helper functions below to construct / decompose keys.
27 +OffloadKey = NewType("OffloadKey", bytes)
28 +
29 +logger = init_logger(__name__)
30 +
31 +
32 +def make_offload_key(block_hash: bytes, group_idx: int) -> OffloadKey:
33 + """Pack a block hash and group index into an `OffloadKey`."""
34 + return OffloadKey(block_hash + group_idx.to_bytes(4, "big", signed=False))
35 +
36 +
37 +def get_offload_block_hash(key: OffloadKey) -> bytes:
38 + """Extract the block hash from an `OffloadKey`."""
39 + return key[:-4]
40 +
41 +
42 +def get_offload_group_idx(key: OffloadKey) -> int:
43 + """Extract the group index from an `OffloadKey`."""
44 + return int.from_bytes(key[-4:], "big", signed=False)
45 +
46 +
47 +@dataclass
48 +class ReqContext:
49 + kv_transfer_params: dict[str, Any] | None = None
50 +
51 +
52 +class LoadStoreSpec(ABC):
53 + """
54 + Abstract metadata that encapsulates information allowing a worker
55 + to load, and optionally also to store, blocks of KV data.
56 + """
57 +
58 + @staticmethod
59 + @abstractmethod
60 + def medium() -> str:
61 + """
62 + Returns a string representation of the medium type
63 + this store/load targets.
64 + """
65 + pass
66 +
67 +
68 +@dataclass
69 +class PrepareStoreOutput:
70 + keys_to_store: list[OffloadKey]
71 + store_spec: LoadStoreSpec
72 + evicted_keys: list[OffloadKey]
73 +
74 +
75 +@dataclass
76 +class OffloadingEvent:
77 + keys: list[OffloadKey]
78 + medium: str
79 + # True if blocks are removed, False if stored
80 + removed: bool
81 +
82 +
83 +"""
84 +OffloadingManager class for managing KV data offloading in vLLM v1
85 +
86 +This class runs in the scheduler, tracks which blocks are offloaded
87 +and their address.
88 +
89 +The class provides the following primitives:
90 + lookup() - check whether a single block is offloaded and ready.
91 + prepare_load() - prepare given blocks to be read.
92 + The given blocks will be protected from eviction.
93 + This function returns a LoadSpec which encapsulates
94 + information required for performing the load.
95 + touch() - marks the give blocks as recently used. Can be used
96 + to track block's LRU. This function is separated from the
97 + prepare_load function to allow setting block recency even
98 + for blocks which do not need reading from the cache, such as
99 + blocks that are cached by the GPU prefix cache.
100 + complete_load() - mark blocks which were previously prepared to be
101 + loaded as done loading. This is to re-allow their eviction.
102 + prepare_store() - prepare the given blocks to be written.
103 + Returns a StoreSpec encapsulating offloading information,
104 + as well as a list of blocks that were evicted as a result.
105 + complete_store() - marks a previous store as completed.
106 + Following this call, the given blocks will become loadable.
107 +"""
108 +
109 +
110 +class OffloadingManager(ABC):
111 + @abstractmethod
112 + def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None:
113 + """
114 + Checks whether a single block is offloaded and ready to be read.
115 +
116 + Args:
117 + key: the key identifying the block to lookup.
118 + req_context: per-request context (e.g. kv_transfer_params).
119 +
120 + Returns:
121 + True if the block is offloaded and ready, False if not,
122 + or None if the lookup should be retried later.
123 + Returning None will delay the request handling by the vLLM
124 + scheduler.
125 + """
126 + pass
127 +
128 + @abstractmethod
129 + def prepare_load(
130 + self,
131 + keys: Sequence[OffloadKey],
132 + req_context: ReqContext,
133 + ) -> LoadStoreSpec:
134 + """
135 + Prepare the given blocks to be read.
136 + The given blocks will be protected from eviction until
137 + complete_load is called.
138 + It assumes all given blocks are offloaded.
139 +
140 + Args:
141 + keys: the keys identifying the blocks.
142 + req_context: per-request context (e.g. kv_transfer_params).
143 +
144 + Returns:
145 + A LoadStoreSpec that can be used by a worker to locate and load
146 + the actual offloaded KV data.
147 + """
148 + pass
149 +
150 + def touch(self, keys: Sequence[OffloadKey]):
151 + """
152 + Mark the given blocks as recently used.
153 + This could in practice mean moving them to the end of an LRU list.
154 +
155 + Args:
156 + keys: the keys identifying the blocks.
157 + """
158 + return
159 +
160 + def complete_load(self, keys: Iterable[OffloadKey]):
161 + """
162 + Marks previous blocks that were prepared to load as done loading.
163 +
164 + Args:
165 + keys: the keys identifying the blocks.
166 + """
167 + return
168 +
169 + @abstractmethod
170 + def prepare_store(
171 + self,
172 + keys: Sequence[OffloadKey],
173 + req_context: ReqContext,
174 + ) -> PrepareStoreOutput | None:
175 + """
176 + Prepare the given blocks to be offloaded.
177 + The given blocks will be protected from eviction until
178 + complete_store is called.
179 +
180 + Args:
181 + keys: the keys identifying the blocks.
182 + req_context: per-request context (e.g. kv_transfer_params).
183 +
184 + Returns:
185 + A PrepareStoreOutput indicating which blocks need storing,
186 + where to store them (LoadStoreSpec), and list of blocks that
187 + were evicted as a result.
188 + None is returned if the blocks cannot be stored.
189 + """
190 + pass
191 +
192 + def complete_store(self, keys: Iterable[OffloadKey], success: bool = True):
193 + """
194 + Marks blocks which were previously prepared to be stored, as stored.
195 + Following this call, the blocks become loadable.
196 + If if_success is False, blocks that were not marked as stored will be
197 + removed.
198 +
199 + Args:
200 + keys: the keys identifying the blocks.
201 + success: whether the blocks were stored successfully.
202 + """
203 + return
204 +
205 + def take_events(self) -> Iterable[OffloadingEvent]:
206 + """
207 + Take the offloading events from the manager.
208 +
209 + Yields:
210 + New OffloadingEvents collected since the last call.
211 + """
212 + return ()
213 +
214 + def shutdown(self) -> None:
215 + """Shutdown the manager and release any resources."""
216 + return
217 +
218 +
219 +class BlockIDsLoadStoreSpec(LoadStoreSpec, ABC):
220 + """
221 + Spec for loading/storing KV blocks from given block numbers.
222 + """
223 +
224 + def __init__(self, block_ids: list[int]):
225 + self.block_ids = np.array(block_ids, dtype=np.int64)
226 +
227 + def __repr__(self) -> str:
228 + return repr(self.block_ids)
229 +
230 +
231 +class GPULoadStoreSpec(BlockIDsLoadStoreSpec):
232 + """
233 + Spec for loading/storing a KV block to GPU memory.
234 +
235 + If there are multiple KV groups, the blocks are expected to be
236 + ordered by the group index.
237 + In that case, group_sizes[i] determines the number of blocks
238 + per the i-th KV group, and thus sum(group_sizes) == len(block_ids).
239 + group_sizes=None indicates a single KV group.
240 +
241 + If block_indices is given, each group (determined by group_sizes) of block IDs
242 + will correspond to logically contiguous blocks, e.g. blocks 5-10 of a some request.
243 + block_indices[i] will represent the block index of the first block in group #i.
244 + Thus, len(block_indices) == len(group_sizes) = number of KV cache groups.
245 + This information is required in order to support off/loading from offloaded blocks
246 + which are larger than GPU blocks.
247 + In such cases, the first GPU block per each group may be unaligned to the offloaded
248 + block size, and so knowing block_indices[i] allows the worker to correctly
249 + skip part of the first matching offloaded block.
250 + """
251 +
252 + def __init__(
253 + self,
254 + block_ids: list[int],
255 + group_sizes: Sequence[int],
256 + block_indices: Sequence[int],
257 + ):
258 + super().__init__(block_ids)
259 + assert sum(group_sizes) == len(block_ids)
260 + assert len(block_indices) == len(group_sizes)
261 + self.group_sizes: Sequence[int] = group_sizes
262 + self.block_indices: Sequence[int] = block_indices
263 +
264 + @staticmethod
265 + def medium() -> str:
266 + return "GPU"
267 +
268 +
269 +@dataclass
270 +class CanonicalKVCacheTensor:
271 + """
272 + A canonicalized KV cache tensor whose first dimension is num_blocks.
273 +
274 + For attention backends where the raw tensor has num_blocks at a
275 + non-leading physical dimension (e.g. FlashAttention's
276 + (2, num_blocks, ...) layout), the tensor is split so that each
277 + resulting CanonicalKVCacheTensor starts with (num_blocks, ...).
278 + """
279 +
280 + # The KV cache tensor with shape (num_blocks, ...)
281 + tensor: torch.Tensor
282 + # The (possibly padded) page size per block in bytes
283 + page_size_bytes: int
284 +
285 +
286 +@dataclass
287 +class CanonicalKVCacheRef:
288 + """
289 + Per-layer (or group of layers) reference to a specific (by index)
290 + CanonicalKVCacheTensor and records the un-padded page size used by that layer.
291 + """
292 +
293 + # Index into the list of CanonicalKVCacheTensor objects
294 + tensor_idx: int
295 + # The un-padded page size per block in bytes
296 + page_size_bytes: int
297 +
298 +
299 +@dataclass
300 +class CanonicalKVCaches:
301 + """
302 + Canonicalized block-level representation of the KV caches.
303 +
304 + Composed of:
305 + - Unique list of KV cache data tensors,
306 + each with shape (num_blocks, page_size_in_bytes) and int8 dtype.
307 + - Per-group data references of the tensors.
308 + i.e. how each KV cache group maps to the tensors.
309 + """
310 +
311 + # Ordered list of unique block tensors, each with shape
312 + # (num_blocks, ...).
313 + tensors: list[CanonicalKVCacheTensor]
314 + # Per-KV-cache-group list of data references that map each layer
315 + # in the group to the appropriate entry in the tensors list.
316 + group_data_refs: list[list[CanonicalKVCacheRef]]
317 +
318 +
319 +class OffloadingSpec(ABC):
320 + """Spec for an offloading connector"""
321 +
322 + def __init__(self, vllm_config: VllmConfig, kv_cache_config: KVCacheConfig):
323 + logger.warning(
324 + "Initializing OffloadingSpec. This API is experimental and "
325 + "subject to change in the future as we iterate the design."
326 + )
327 + self.vllm_config = vllm_config
328 + self.kv_cache_config = kv_cache_config
329 +
330 + kv_transfer_config = vllm_config.kv_transfer_config
331 + assert kv_transfer_config is not None
332 + self.extra_config = kv_transfer_config.kv_connector_extra_config
333 +
334 + # block size used by vLLM for hashing request tokens for the sake
335 + # of enabling prefix caching
336 + self.hash_block_size = vllm_config.cache_config.block_size
337 + # gpu block size per group
338 + self.gpu_block_size: tuple[int, ...] = tuple(
339 + kv_cache_group.kv_cache_spec.block_size
340 + for kv_cache_group in kv_cache_config.kv_cache_groups
341 + )
342 +
343 + for block_size in self.gpu_block_size:
344 + assert block_size % self.hash_block_size == 0, (
345 + f"gpu_block_size={block_size} not divisible by "
346 + f"hash_block_size={self.hash_block_size}. "
347 + f"Hybrid models (e.g. Mamba+Attention) need "
348 + f"--enable-prefix-caching to align block sizes."
349 + )
350 +
351 + # offloaded_block_size / gpu_block_size
352 + self.block_size_factor: int = 1
353 +
354 + offloaded_block_size = self.extra_config.get("block_size")
355 + if offloaded_block_size is not None:
356 + offloaded_block_size_int = int(offloaded_block_size)
357 + gpu_block_sizes = set(self.gpu_block_size)
358 + assert len(gpu_block_sizes) == 1, (
359 + "If 'block_size' is specified in kv_connector_extra_config, "
360 + "there must be at least one KV cache group, "
361 + "and all groups must have the same block size."
362 + )
363 + gpu_block_size = gpu_block_sizes.pop()
364 +
365 + assert offloaded_block_size_int % gpu_block_size == 0
366 + self.block_size_factor = offloaded_block_size_int // gpu_block_size
367 +
368 + @abstractmethod
369 + def get_manager(self) -> OffloadingManager:
370 + """
371 + Get an OffloadingManager that will be used
372 + by the scheduler-side offloading connector to track
373 + offloaded blocks and manage evictions.
374 + """
375 + pass
376 +
377 + @abstractmethod
378 + def get_handlers(
379 + self, kv_caches: CanonicalKVCaches
380 + ) -> Iterator[tuple[type[LoadStoreSpec], type[LoadStoreSpec], OffloadingHandler]]:
381 + """
382 + Get offloading handlers along with their respective src and dst types.
383 +
384 + Args:
385 + kv_caches: Canonicalized KV caches.
386 +
387 + Yields:
388 + Tuples of (src_type, dst_type, offloading_handler).
389 + """
390 + pass
vllm/v1/kv_offload/cpu/common.py new
+13
@@ -0,0 +1,13 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from vllm.v1.kv_offload.base import BlockIDsLoadStoreSpec
4 +
5 +
6 +class CPULoadStoreSpec(BlockIDsLoadStoreSpec):
7 + """
8 + Spec for loading/storing a KV block to CPU memory.
9 + """
10 +
11 + @staticmethod
12 + def medium() -> str:
13 + return "CPU"
vllm/v1/kv_offload/cpu/gpu_worker.py renamed
+6 -2
@@ -11,9 +11,13 @@ from vllm import _custom_ops as ops
11 from vllm.logger import init_logger
12 from vllm.utils.math_utils import cdiv
13 from vllm.utils.platform_utils import is_pin_memory_available
14 +from vllm.v1.kv_offload.base import (
15 + BlockIDsLoadStoreSpec,
16 + CanonicalKVCacheRef,
17 + CanonicalKVCaches,
18 + GPULoadStoreSpec,
19 +)
20 from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion
15 -from vllm.v1.kv_offload.mediums import BlockIDsLoadStoreSpec, GPULoadStoreSpec
16 -from vllm.v1.kv_offload.spec import CanonicalKVCacheRef, CanonicalKVCaches
21 from vllm.v1.kv_offload.worker.worker import (
22 OffloadingHandler,
23 TransferResult,
vllm/v1/kv_offload/cpu/manager.py
+3 -3
@@ -3,7 +3,7 @@
3 from collections.abc import Iterable, Sequence
4 from typing import Literal
5
6 -from vllm.v1.kv_offload.abstract import (
6 +from vllm.v1.kv_offload.base import (
7 LoadStoreSpec,
8 OffloadingEvent,
9 OffloadingManager,
@@ -11,10 +11,10 @@ from vllm.v1.kv_offload.abstract import (
11 PrepareStoreOutput,
12 ReqContext,
13 )
14 -from vllm.v1.kv_offload.cpu.policies.abstract import BlockStatus, CachePolicy
14 +from vllm.v1.kv_offload.cpu.common import CPULoadStoreSpec
15 from vllm.v1.kv_offload.cpu.policies.arc import ARCCachePolicy
16 +from vllm.v1.kv_offload.cpu.policies.base import BlockStatus, CachePolicy
17 from vllm.v1.kv_offload.cpu.policies.lru import LRUCachePolicy
17 -from vllm.v1.kv_offload.mediums import CPULoadStoreSpec
18
19 _CACHE_POLICIES: dict[str, type[CachePolicy]] = {
20 "lru": LRUCachePolicy,
vllm/v1/kv_offload/cpu/policies/arc.py
+2 -2
@@ -3,8 +3,8 @@
3 from collections import OrderedDict
4 from collections.abc import Iterable
5
6 -from vllm.v1.kv_offload.abstract import OffloadKey
7 -from vllm.v1.kv_offload.cpu.policies.abstract import BlockStatus, CachePolicy
6 +from vllm.v1.kv_offload.base import OffloadKey
7 +from vllm.v1.kv_offload.cpu.policies.base import BlockStatus, CachePolicy
8
9
10 class ARCCachePolicy(CachePolicy):
vllm/v1/kv_offload/cpu/policies/base.py renamed
+1 -1
@@ -4,7 +4,7 @@ import ctypes
4 from abc import ABC, abstractmethod
5 from collections.abc import Iterable
6
7 -from vllm.v1.kv_offload.abstract import OffloadKey
7 +from vllm.v1.kv_offload.base import OffloadKey
8
9
10 class BlockStatus(ctypes.Structure):
vllm/v1/kv_offload/cpu/policies/lru.py
+2 -2
@@ -3,8 +3,8 @@
3 from collections import OrderedDict
4 from collections.abc import Iterable
5
6 -from vllm.v1.kv_offload.abstract import OffloadKey
7 -from vllm.v1.kv_offload.cpu.policies.abstract import BlockStatus, CachePolicy
6 +from vllm.v1.kv_offload.base import OffloadKey
7 +from vllm.v1.kv_offload.cpu.policies.base import BlockStatus, CachePolicy
8
9
10 class LRUCachePolicy(CachePolicy):
vllm/v1/kv_offload/cpu/spec.py
+9 -4
@@ -5,12 +5,17 @@ from collections.abc import Iterator
5 from vllm.config import VllmConfig
6 from vllm.platforms import current_platform
7 from vllm.v1.kv_cache_interface import KVCacheConfig
8 -from vllm.v1.kv_offload.abstract import LoadStoreSpec, OffloadingManager
8 +from vllm.v1.kv_offload.base import (
9 + CanonicalKVCaches,
10 + GPULoadStoreSpec,
11 + LoadStoreSpec,
12 + OffloadingManager,
13 + OffloadingSpec,
14 +)
15 +from vllm.v1.kv_offload.cpu.common import CPULoadStoreSpec
16 +from vllm.v1.kv_offload.cpu.gpu_worker import CpuGpuOffloadingHandlers
17 from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager
10 -from vllm.v1.kv_offload.mediums import CPULoadStoreSpec, GPULoadStoreSpec
18 from vllm.v1.kv_offload.reuse_manager import FilterReusedOffloadingManager
12 -from vllm.v1.kv_offload.spec import CanonicalKVCaches, OffloadingSpec
13 -from vllm.v1.kv_offload.worker.cpu_gpu import CpuGpuOffloadingHandlers
19 from vllm.v1.kv_offload.worker.worker import OffloadingHandler
20
21
vllm/v1/kv_offload/factory.py
+1 -1
@@ -5,7 +5,7 @@ from collections.abc import Callable
5 from typing import TYPE_CHECKING
6
7 from vllm.logger import init_logger
8 -from vllm.v1.kv_offload.spec import OffloadingSpec
8 +from vllm.v1.kv_offload.base import OffloadingSpec
9
10 if TYPE_CHECKING:
11 from vllm.config import VllmConfig
vllm/v1/kv_offload/mediums.py deleted
-68
@@ -1,68 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -from abc import ABC
4 -from collections.abc import Sequence
5 -
6 -import numpy as np
7 -
8 -from vllm.v1.kv_offload.abstract import LoadStoreSpec
9 -
10 -
11 -class BlockIDsLoadStoreSpec(LoadStoreSpec, ABC):
12 - """
13 - Spec for loading/storing KV blocks from given block numbers.
14 - """
15 -
16 - def __init__(self, block_ids: list[int]):
17 - self.block_ids = np.array(block_ids, dtype=np.int64)
18 -
19 - def __repr__(self) -> str:
20 - return repr(self.block_ids)
21 -
22 -
23 -class GPULoadStoreSpec(BlockIDsLoadStoreSpec):
24 - """
25 - Spec for loading/storing a KV block to GPU memory.
26 -
27 - If there are multiple KV groups, the blocks are expected to be
28 - ordered by the group index.
29 - In that case, group_sizes[i] determines the number of blocks
30 - per the i-th KV group, and thus sum(group_sizes) == len(block_ids).
31 - group_sizes=None indicates a single KV group.
32 -
33 - If block_indices is given, each group (determined by group_sizes) of block IDs
34 - will correspond to logically contiguous blocks, e.g. blocks 5-10 of a some request.
35 - block_indices[i] will represent the block index of the first block in group #i.
36 - Thus, len(block_indices) == len(group_sizes) = number of KV cache groups.
37 - This information is required in order to support off/loading from offloaded blocks
38 - which are larger than GPU blocks.
39 - In such cases, the first GPU block per each group may be unaligned to the offloaded
40 - block size, and so knowing block_indices[i] allows the worker to correctly
41 - skip part of the first matching offloaded block.
42 - """
43 -
44 - def __init__(
45 - self,
46 - block_ids: list[int],
47 - group_sizes: Sequence[int],
48 - block_indices: Sequence[int],
49 - ):
50 - super().__init__(block_ids)
51 - assert sum(group_sizes) == len(block_ids)
52 - assert len(block_indices) == len(group_sizes)
53 - self.group_sizes: Sequence[int] = group_sizes
54 - self.block_indices: Sequence[int] = block_indices
55 -
56 - @staticmethod
57 - def medium() -> str:
58 - return "GPU"
59 -
60 -
61 -class CPULoadStoreSpec(BlockIDsLoadStoreSpec):
62 - """
63 - Spec for loading/storing a KV block to CPU memory.
64 - """
65 -
66 - @staticmethod
67 - def medium() -> str:
68 - return "CPU"
vllm/v1/kv_offload/reuse_manager.py
+1 -1
@@ -10,7 +10,7 @@ FilterReusedOffloadingManager — OffloadingManager decorator that skips
10 from collections import OrderedDict
11 from collections.abc import Iterable, Sequence
12
13 -from vllm.v1.kv_offload.abstract import (
13 +from vllm.v1.kv_offload.base import (
14 LoadStoreSpec,
15 OffloadingEvent,
16 OffloadingManager,
vllm/v1/kv_offload/spec.py deleted
-142
@@ -1,142 +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 collections.abc import Iterator
5 -from dataclasses import dataclass
6 -from typing import TYPE_CHECKING
7 -
8 -import torch
9 -
10 -from vllm.logger import init_logger
11 -from vllm.v1.kv_offload.abstract import LoadStoreSpec, OffloadingManager
12 -from vllm.v1.kv_offload.worker.worker import OffloadingHandler
13 -
14 -if TYPE_CHECKING:
15 - from vllm.config import VllmConfig
16 - from vllm.v1.kv_cache_interface import KVCacheConfig
17 -
18 -logger = init_logger(__name__)
19 -
20 -
21 -@dataclass
22 -class CanonicalKVCacheTensor:
23 - """
24 - A canonicalized KV cache tensor whose first dimension is num_blocks.
25 -
26 - For attention backends where the raw tensor has num_blocks at a
27 - non-leading physical dimension (e.g. FlashAttention's
28 - (2, num_blocks, ...) layout), the tensor is split so that each
29 - resulting CanonicalKVCacheTensor starts with (num_blocks, ...).
30 - """
31 -
32 - # The KV cache tensor with shape (num_blocks, ...)
33 - tensor: torch.Tensor
34 - # The (possibly padded) page size per block in bytes
35 - page_size_bytes: int
36 -
37 -
38 -@dataclass
39 -class CanonicalKVCacheRef:
40 - """
41 - Per-layer (or group of layers) reference to a specific (by index)
42 - CanonicalKVCacheTensor and records the un-padded page size used by that layer.
43 - """
44 -
45 - # Index into the list of CanonicalKVCacheTensor objects
46 - tensor_idx: int
47 - # The un-padded page size per block in bytes
48 - page_size_bytes: int
49 -
50 -
51 -@dataclass
52 -class CanonicalKVCaches:
53 - """
54 - Canonicalized block-level representation of the KV caches.
55 -
56 - Composed of:
57 - - Unique list of KV cache data tensors,
58 - each with shape (num_blocks, page_size_in_bytes) and int8 dtype.
59 - - Per-group data references of the tensors.
60 - i.e. how each KV cache group maps to the tensors.
61 - """
62 -
63 - # Ordered list of unique block tensors, each with shape
64 - # (num_blocks, ...).
65 - tensors: list[CanonicalKVCacheTensor]
66 - # Per-KV-cache-group list of data references that map each layer
67 - # in the group to the appropriate entry in the tensors list.
68 - group_data_refs: list[list[CanonicalKVCacheRef]]
69 -
70 -
71 -class OffloadingSpec(ABC):
72 - """Spec for an offloading connector"""
73 -
74 - def __init__(self, vllm_config: "VllmConfig", kv_cache_config: "KVCacheConfig"):
75 - logger.warning(
76 - "Initializing OffloadingSpec. This API is experimental and "
77 - "subject to change in the future as we iterate the design."
78 - )
79 - self.vllm_config = vllm_config
80 - self.kv_cache_config = kv_cache_config
81 -
82 - kv_transfer_config = vllm_config.kv_transfer_config
83 - assert kv_transfer_config is not None
84 - self.extra_config = kv_transfer_config.kv_connector_extra_config
85 -
86 - # block size used by vLLM for hashing request tokens for the sake
87 - # of enabling prefix caching
88 - self.hash_block_size = vllm_config.cache_config.block_size
89 - # gpu block size per group
90 - self.gpu_block_size: tuple[int, ...] = tuple(
91 - kv_cache_group.kv_cache_spec.block_size
92 - for kv_cache_group in kv_cache_config.kv_cache_groups
93 - )
94 -
95 - for block_size in self.gpu_block_size:
96 - assert block_size % self.hash_block_size == 0, (
97 - f"gpu_block_size={block_size} not divisible by "
98 - f"hash_block_size={self.hash_block_size}. "
99 - f"Hybrid models (e.g. Mamba+Attention) need "
100 - f"--enable-prefix-caching to align block sizes."
101 - )
102 -
103 - # offloaded_block_size / gpu_block_size
104 - self.block_size_factor: int = 1
105 -
106 - offloaded_block_size = self.extra_config.get("block_size")
107 - if offloaded_block_size is not None:
108 - offloaded_block_size_int = int(offloaded_block_size)
109 - gpu_block_sizes = set(self.gpu_block_size)
110 - assert len(gpu_block_sizes) == 1, (
111 - "If 'block_size' is specified in kv_connector_extra_config, "
112 - "there must be at least one KV cache group, "
113 - "and all groups must have the same block size."
114 - )
115 - gpu_block_size = gpu_block_sizes.pop()
116 -
117 - assert offloaded_block_size_int % gpu_block_size == 0
118 - self.block_size_factor = offloaded_block_size_int // gpu_block_size
119 -
120 - @abstractmethod
121 - def get_manager(self) -> OffloadingManager:
122 - """
123 - Get an OffloadingManager that will be used
124 - by the scheduler-side offloading connector to track
125 - offloaded blocks and manage evictions.
126 - """
127 - pass
128 -
129 - @abstractmethod
130 - def get_handlers(
131 - self, kv_caches: CanonicalKVCaches
132 - ) -> Iterator[tuple[type[LoadStoreSpec], type[LoadStoreSpec], OffloadingHandler]]:
133 - """
134 - Get offloading handlers along with their respective src and dst types.
135 -
136 - Args:
137 - kv_caches: Canonicalized KV caches.
138 -
139 - Yields:
140 - Tuples of (src_type, dst_type, offloading_handler).
141 - """
142 - pass
vllm/v1/kv_offload/worker/worker.py
+1 -1
@@ -4,7 +4,7 @@ from abc import ABC, abstractmethod
4 from dataclasses import dataclass
5
6 from vllm.logger import init_logger
7 -from vllm.v1.kv_offload.abstract import LoadStoreSpec
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]