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, 12:30 UTC
[KV Offload] Tighten `keys` type from `Iterable` to `Sequence` in `OffloadingManager` (#41200)
Signed-off-by: Ronen Schaffer <ronen.schaffer@ibm.com>
Ronen Schaffer committed
Apr 29, 2026 at 13:50 UTC
37e288214bc3fa89d974b4d323373f2b2878d604
3 files changed
+14
-17
vllm/v1/kv_offload/abstract.py
+4
-4
@@ -27,7 +27,7 @@ The class provides the following primitives:
27
"""
28
29
from abc import ABC, abstractmethod
30
-from collections.abc import Iterable
30
+from collections.abc import Iterable, Sequence
31
from dataclasses import dataclass
32
from typing import Any, NewType
33
@@ -109,7 +109,7 @@ class OffloadingManager(ABC):
109
@abstractmethod
110
def prepare_load(
111
self,
112
- keys: Iterable[OffloadKey],
112
+ keys: Sequence[OffloadKey],
113
req_context: ReqContext,
114
) -> LoadStoreSpec:
115
"""
@@ -128,7 +128,7 @@ class OffloadingManager(ABC):
128
"""
129
pass
130
131
- def touch(self, keys: Iterable[OffloadKey]):
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.
@@ -150,7 +150,7 @@ class OffloadingManager(ABC):
150
@abstractmethod
151
def prepare_store(
152
self,
153
- keys: Iterable[OffloadKey],
153
+ keys: Sequence[OffloadKey],
154
req_context: ReqContext,
155
) -> PrepareStoreOutput | None:
156
"""
vllm/v1/kv_offload/cpu/manager.py
+6
-8
@@ -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
3
+from collections.abc import Iterable, Sequence
4
from typing import Literal
5
6
from vllm.v1.kv_offload.abstract import (
@@ -90,7 +90,7 @@ class CPUOffloadingManager(OffloadingManager):
90
91
def prepare_load(
92
self,
93
- keys: Iterable[OffloadKey],
93
+ keys: Sequence[OffloadKey],
94
req_context: ReqContext,
95
) -> LoadStoreSpec:
96
blocks = []
@@ -102,7 +102,7 @@ class CPUOffloadingManager(OffloadingManager):
102
blocks.append(block)
103
return self._get_load_store_spec(keys, blocks)
104
105
- def touch(self, keys: Iterable[OffloadKey]) -> None:
105
+ def touch(self, keys: Sequence[OffloadKey]) -> None:
106
self._policy.touch(keys)
107
108
def complete_load(self, keys: Iterable[OffloadKey]) -> None:
@@ -114,13 +114,11 @@ class CPUOffloadingManager(OffloadingManager):
114
115
def prepare_store(
116
self,
117
- keys: Iterable[OffloadKey],
117
+ keys: Sequence[OffloadKey],
118
req_context: ReqContext,
119
) -> PrepareStoreOutput | None:
120
- keys_list = list(keys)
121
-
120
# filter out blocks that are already stored
123
- keys_to_store = [k for k in keys_list if self._policy.get(k) is None]
121
+ keys_to_store = [k for k in keys if self._policy.get(k) is None]
122
123
if not keys_to_store:
124
return PrepareStoreOutput(
@@ -135,7 +133,7 @@ class CPUOffloadingManager(OffloadingManager):
133
if num_blocks_to_evict > 0:
134
# Blocks from the original input are excluded from eviction candidates:
135
# a block that was already stored must remain in the cache after this call.
138
- protected = set(keys_list)
136
+ protected = set(keys)
137
evicted = self._policy.evict(num_blocks_to_evict, protected)
138
if evicted is None:
139
return None
vllm/v1/kv_offload/reuse_manager.py
+4
-5
@@ -8,7 +8,7 @@ FilterReusedOffloadingManager — OffloadingManager decorator that skips
8
"""
9
10
from collections import OrderedDict
11
-from collections.abc import Iterable
11
+from collections.abc import Iterable, Sequence
12
13
from vllm.v1.kv_offload.abstract 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: Iterable[OffloadKey], req_context: ReqContext
82
+ self, keys: Sequence[OffloadKey], req_context: ReqContext
83
) -> PrepareStoreOutput | None:
84
"""Filter out blocks below threshold, then delegate to backing.
85
@@ -87,7 +87,6 @@ class FilterReusedOffloadingManager(OffloadingManager):
87
``prepare_store`` so that blocks that would be skipped do not
88
consume any CPU offload capacity.
89
"""
90
- keys = list(keys)
90
eligible = [
91
key for key in keys if self.counts.get(key, 0) >= self.store_threshold
92
]
@@ -102,11 +101,11 @@ class FilterReusedOffloadingManager(OffloadingManager):
101
# ------------------------------------------------------------------
102
103
def prepare_load(
105
- self, keys: Iterable[OffloadKey], req_context: ReqContext
104
+ self, keys: Sequence[OffloadKey], req_context: ReqContext
105
) -> LoadStoreSpec:
106
return self._backing.prepare_load(keys, req_context)
107
109
- def touch(self, keys: Iterable[OffloadKey]) -> None:
108
+ def touch(self, keys: Sequence[OffloadKey]) -> None:
109
return self._backing.touch(keys)
110
111
def complete_load(self, keys: Iterable[OffloadKey]) -> None: