From 37e288214bc3fa89d974b4d323373f2b2878d604 Mon Sep 17 00:00:00 2001 From: Ronen Schaffer Date: Wed, 29 Apr 2026 13:50:42 +0300 Subject: [PATCH] [KV Offload] Tighten `keys` type from `Iterable` to `Sequence` in `OffloadingManager` (#41200) Signed-off-by: Ronen Schaffer --- vllm/v1/kv_offload/abstract.py | 8 ++++---- vllm/v1/kv_offload/cpu/manager.py | 14 ++++++-------- vllm/v1/kv_offload/reuse_manager.py | 9 ++++----- 3 files changed, 14 insertions(+), 17 deletions(-) diff --git a/vllm/v1/kv_offload/abstract.py b/vllm/v1/kv_offload/abstract.py index 8f809ceaa0..be54611356 100644 --- a/vllm/v1/kv_offload/abstract.py +++ b/vllm/v1/kv_offload/abstract.py @@ -27,7 +27,7 @@ The class provides the following primitives: """ from abc import ABC, abstractmethod -from collections.abc import Iterable +from collections.abc import Iterable, Sequence from dataclasses import dataclass from typing import Any, NewType @@ -109,7 +109,7 @@ class OffloadingManager(ABC): @abstractmethod def prepare_load( self, - keys: Iterable[OffloadKey], + keys: Sequence[OffloadKey], req_context: ReqContext, ) -> LoadStoreSpec: """ @@ -128,7 +128,7 @@ class OffloadingManager(ABC): """ pass - def touch(self, keys: Iterable[OffloadKey]): + def touch(self, keys: Sequence[OffloadKey]): """ Mark the given blocks as recently used. This could in practice mean moving them to the end of an LRU list. @@ -150,7 +150,7 @@ class OffloadingManager(ABC): @abstractmethod def prepare_store( self, - keys: Iterable[OffloadKey], + keys: Sequence[OffloadKey], req_context: ReqContext, ) -> PrepareStoreOutput | None: """ diff --git a/vllm/v1/kv_offload/cpu/manager.py b/vllm/v1/kv_offload/cpu/manager.py index fcfaa919a3..fb03846105 100644 --- a/vllm/v1/kv_offload/cpu/manager.py +++ b/vllm/v1/kv_offload/cpu/manager.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from collections.abc import Iterable +from collections.abc import Iterable, Sequence from typing import Literal from vllm.v1.kv_offload.abstract import ( @@ -90,7 +90,7 @@ class CPUOffloadingManager(OffloadingManager): def prepare_load( self, - keys: Iterable[OffloadKey], + keys: Sequence[OffloadKey], req_context: ReqContext, ) -> LoadStoreSpec: blocks = [] @@ -102,7 +102,7 @@ class CPUOffloadingManager(OffloadingManager): blocks.append(block) return self._get_load_store_spec(keys, blocks) - def touch(self, keys: Iterable[OffloadKey]) -> None: + def touch(self, keys: Sequence[OffloadKey]) -> None: self._policy.touch(keys) def complete_load(self, keys: Iterable[OffloadKey]) -> None: @@ -114,13 +114,11 @@ class CPUOffloadingManager(OffloadingManager): def prepare_store( self, - keys: Iterable[OffloadKey], + keys: Sequence[OffloadKey], req_context: ReqContext, ) -> PrepareStoreOutput | None: - keys_list = list(keys) - # filter out blocks that are already stored - keys_to_store = [k for k in keys_list if self._policy.get(k) is None] + keys_to_store = [k for k in keys if self._policy.get(k) is None] if not keys_to_store: return PrepareStoreOutput( @@ -135,7 +133,7 @@ class CPUOffloadingManager(OffloadingManager): if num_blocks_to_evict > 0: # Blocks from the original input are excluded from eviction candidates: # a block that was already stored must remain in the cache after this call. - protected = set(keys_list) + protected = set(keys) evicted = self._policy.evict(num_blocks_to_evict, protected) if evicted is None: return None diff --git a/vllm/v1/kv_offload/reuse_manager.py b/vllm/v1/kv_offload/reuse_manager.py index 96b8f969e7..751f3a93df 100644 --- a/vllm/v1/kv_offload/reuse_manager.py +++ b/vllm/v1/kv_offload/reuse_manager.py @@ -8,7 +8,7 @@ FilterReusedOffloadingManager — OffloadingManager decorator that skips """ from collections import OrderedDict -from collections.abc import Iterable +from collections.abc import Iterable, Sequence from vllm.v1.kv_offload.abstract import ( LoadStoreSpec, @@ -79,7 +79,7 @@ class FilterReusedOffloadingManager(OffloadingManager): return self._backing.lookup(key, req_context) def prepare_store( - self, keys: Iterable[OffloadKey], req_context: ReqContext + self, keys: Sequence[OffloadKey], req_context: ReqContext ) -> PrepareStoreOutput | None: """Filter out blocks below threshold, then delegate to backing. @@ -87,7 +87,6 @@ class FilterReusedOffloadingManager(OffloadingManager): ``prepare_store`` so that blocks that would be skipped do not consume any CPU offload capacity. """ - keys = list(keys) eligible = [ key for key in keys if self.counts.get(key, 0) >= self.store_threshold ] @@ -102,11 +101,11 @@ class FilterReusedOffloadingManager(OffloadingManager): # ------------------------------------------------------------------ def prepare_load( - self, keys: Iterable[OffloadKey], req_context: ReqContext + self, keys: Sequence[OffloadKey], req_context: ReqContext ) -> LoadStoreSpec: return self._backing.prepare_load(keys, req_context) - def touch(self, keys: Iterable[OffloadKey]) -> None: + def touch(self, keys: Sequence[OffloadKey]) -> None: return self._backing.touch(keys) def complete_load(self, keys: Iterable[OffloadKey]) -> None: