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, 10:00 UTC
[KV Offload] Use `Collection` instead of `Sequence/Iterable` for OffloadingManager key parameters (#41361)
Signed-off-by: Ronen Schaffer <ronen.schaffer@ibm.com>
Ronen Schaffer committed
May 1, 2026 at 05:18 UTC
415a8798996c518bbc33f377c9652c2776e074a7
3 files changed
+22
-18
vllm/v1/kv_offload/base.py
+6
-6
@@ -7,7 +7,7 @@ Core abstractions for KV cache offloading in vLLM v1.
7
from __future__ import annotations
8
9
from abc import ABC, abstractmethod
10
-from collections.abc import Iterable, Iterator, Sequence
10
+from collections.abc import Collection, Iterable, Iterator, Sequence
11
from dataclasses import dataclass
12
from typing import TYPE_CHECKING, Any, NewType
13
@@ -128,7 +128,7 @@ class OffloadingManager(ABC):
128
@abstractmethod
129
def prepare_load(
130
self,
131
- keys: Sequence[OffloadKey],
131
+ keys: Collection[OffloadKey],
132
req_context: ReqContext,
133
) -> LoadStoreSpec:
134
"""
@@ -147,7 +147,7 @@ class OffloadingManager(ABC):
147
"""
148
pass
149
150
- def touch(self, keys: Sequence[OffloadKey]):
150
+ def touch(self, keys: Collection[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.
@@ -157,7 +157,7 @@ class OffloadingManager(ABC):
157
"""
158
return
159
160
- def complete_load(self, keys: Iterable[OffloadKey]):
160
+ def complete_load(self, keys: Collection[OffloadKey]):
161
"""
162
Marks previous blocks that were prepared to load as done loading.
163
@@ -169,7 +169,7 @@ class OffloadingManager(ABC):
169
@abstractmethod
170
def prepare_store(
171
self,
172
- keys: Sequence[OffloadKey],
172
+ keys: Collection[OffloadKey],
173
req_context: ReqContext,
174
) -> PrepareStoreOutput | None:
175
"""
@@ -189,7 +189,7 @@ class OffloadingManager(ABC):
189
"""
190
pass
191
192
- def complete_store(self, keys: Iterable[OffloadKey], success: bool = True):
192
+ def complete_store(self, keys: Collection[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.
vllm/v1/kv_offload/cpu/manager.py
+8
-6
@@ -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 Iterable, Sequence
3
+from collections.abc import Collection, Iterable
4
from typing import Literal
5
6
from vllm.v1.kv_offload.base import (
@@ -90,7 +90,7 @@ class CPUOffloadingManager(OffloadingManager):
90
91
def prepare_load(
92
self,
93
- keys: Sequence[OffloadKey],
93
+ keys: Collection[OffloadKey],
94
req_context: ReqContext,
95
) -> LoadStoreSpec:
96
blocks = []
@@ -102,10 +102,10 @@ class CPUOffloadingManager(OffloadingManager):
102
blocks.append(block)
103
return self._get_load_store_spec(keys, blocks)
104
105
- def touch(self, keys: Sequence[OffloadKey]) -> None:
105
+ def touch(self, keys: Collection[OffloadKey]) -> None:
106
self._policy.touch(keys)
107
108
- def complete_load(self, keys: Iterable[OffloadKey]) -> None:
108
+ def complete_load(self, keys: Collection[OffloadKey]) -> None:
109
for key in keys:
110
block = self._policy.get(key)
111
assert block is not None, f"Block {key!r} not found"
@@ -114,7 +114,7 @@ class CPUOffloadingManager(OffloadingManager):
114
115
def prepare_store(
116
self,
117
- keys: Sequence[OffloadKey],
117
+ keys: Collection[OffloadKey],
118
req_context: ReqContext,
119
) -> PrepareStoreOutput | None:
120
# filter out blocks that are already stored
@@ -167,7 +167,9 @@ class CPUOffloadingManager(OffloadingManager):
167
evicted_keys=to_evict,
168
)
169
170
- def complete_store(self, keys: Iterable[OffloadKey], success: bool = True) -> None:
170
+ def complete_store(
171
+ self, keys: Collection[OffloadKey], success: bool = True
172
+ ) -> None:
173
stored_keys: list[OffloadKey] = []
174
175
if success:
vllm/v1/kv_offload/reuse_manager.py
+8
-6
@@ -8,7 +8,7 @@ FilterReusedOffloadingManager — OffloadingManager decorator that skips
8
"""
9
10
from collections import OrderedDict
11
-from collections.abc import Iterable, Sequence
11
+from collections.abc import Collection, Iterable
12
13
from vllm.v1.kv_offload.base import (
14
LoadStoreSpec,
@@ -79,7 +79,7 @@ class FilterReusedOffloadingManager(OffloadingManager):
79
return self._backing.lookup(key, req_context)
80
81
def prepare_store(
82
- self, keys: Sequence[OffloadKey], req_context: ReqContext
82
+ self, keys: Collection[OffloadKey], req_context: ReqContext
83
) -> PrepareStoreOutput | None:
84
"""Filter out blocks below threshold, then delegate to backing.
85
@@ -101,17 +101,19 @@ class FilterReusedOffloadingManager(OffloadingManager):
101
# ------------------------------------------------------------------
102
103
def prepare_load(
104
- self, keys: Sequence[OffloadKey], req_context: ReqContext
104
+ self, keys: Collection[OffloadKey], req_context: ReqContext
105
) -> LoadStoreSpec:
106
return self._backing.prepare_load(keys, req_context)
107
108
- def touch(self, keys: Sequence[OffloadKey]) -> None:
108
+ def touch(self, keys: Collection[OffloadKey]) -> None:
109
return self._backing.touch(keys)
110
111
- def complete_load(self, keys: Iterable[OffloadKey]) -> None:
111
+ def complete_load(self, keys: Collection[OffloadKey]) -> None:
112
return self._backing.complete_load(keys)
113
114
- def complete_store(self, keys: Iterable[OffloadKey], success: bool = True) -> None:
114
+ def complete_store(
115
+ self, keys: Collection[OffloadKey], success: bool = True
116
+ ) -> None:
117
return self._backing.complete_store(keys, success)
118
119
def take_events(self) -> Iterable[OffloadingEvent]: