[KV Offload] Tighten keys type from Iterable to Sequence in OffloadingManager (#41200)

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

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:
         """

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

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: