diff --git a/tests/v1/kv_offload/cpu/test_manager.py b/tests/v1/kv_offload/cpu/test_manager.py index 999f7b38bc..e16d6d378e 100644 --- a/tests/v1/kv_offload/cpu/test_manager.py +++ b/tests/v1/kv_offload/cpu/test_manager.py @@ -897,3 +897,26 @@ def test_evictable_cache_block_count(): manager.complete_store(to_keys([14, 15]), _EMPTY_REQ_CTX) # cache state [10, 11, 14, 15] <- all blocks idle assert manager._num_evictable_cache_blocks == 4 + + +def test_touch_forwards_req_context_to_policy(monkeypatch): + """Regression: CPUOffloadingManager.touch forwards ReqContext to policy.""" + manager = make_cpu_manager(num_blocks=4, cache_policy="lru") + received = [] + + def spy_touch(keys: Iterable[OffloadKey], req_context: ReqContext) -> None: + received.append((list(keys), req_context)) + + monkeypatch.setattr(manager._policy, "touch", spy_touch) + + keys = to_keys([1, 2]) + ctx = make_req_context( + req_id="test-req", + kv_transfer_params={"test_param": "test_value"}, + ) + + manager.touch(keys, ctx) + + assert len(received) == 1 + assert received[0][0] == keys + assert received[0][1] is ctx diff --git a/vllm/v1/kv_offload/cpu/manager.py b/vllm/v1/kv_offload/cpu/manager.py index f0a3d9dead..e7416bf63b 100644 --- a/vllm/v1/kv_offload/cpu/manager.py +++ b/vllm/v1/kv_offload/cpu/manager.py @@ -152,7 +152,7 @@ class CPUOffloadingManager(OffloadingManager): @override def touch(self, keys: Collection[OffloadKey], req_context: ReqContext) -> None: - self._policy.touch(keys) + self._policy.touch(keys, req_context) @override def complete_load( diff --git a/vllm/v1/kv_offload/cpu/policies/arc.py b/vllm/v1/kv_offload/cpu/policies/arc.py index 7d22e51865..f682a47e45 100644 --- a/vllm/v1/kv_offload/cpu/policies/arc.py +++ b/vllm/v1/kv_offload/cpu/policies/arc.py @@ -5,7 +5,7 @@ from collections.abc import Iterable from typing_extensions import override -from vllm.v1.kv_offload.base import OffloadKey +from vllm.v1.kv_offload.base import OffloadKey, ReqContext from vllm.v1.kv_offload.cpu.policies.base import BlockStatus, CachePolicy @@ -72,7 +72,7 @@ class ARCCachePolicy(CachePolicy): self.t2.pop(key, None) @override - def touch(self, keys: Iterable[OffloadKey]) -> None: + def touch(self, keys: Iterable[OffloadKey], req_context: ReqContext) -> None: for key in reversed(list(keys)): if key in self.t1: block = self.t1.pop(key) diff --git a/vllm/v1/kv_offload/cpu/policies/base.py b/vllm/v1/kv_offload/cpu/policies/base.py index f898a60b0f..2b6681e499 100644 --- a/vllm/v1/kv_offload/cpu/policies/base.py +++ b/vllm/v1/kv_offload/cpu/policies/base.py @@ -4,7 +4,7 @@ import ctypes from abc import ABC, abstractmethod from collections.abc import Iterable -from vllm.v1.kv_offload.base import OffloadKey +from vllm.v1.kv_offload.base import OffloadKey, ReqContext class BlockStatus(ctypes.Structure): @@ -57,8 +57,14 @@ class CachePolicy(ABC): """Remove a block (used to clean up after a failed store).""" @abstractmethod - def touch(self, keys: Iterable[OffloadKey]) -> None: - """Mark blocks as recently used.""" + def touch(self, keys: Iterable[OffloadKey], req_context: ReqContext) -> None: + """ + Mark blocks as recently used. + + Args: + keys: Blocks to mark as recently used. + req_context: Per-request context for the request touching these blocks. + """ @abstractmethod def evict( diff --git a/vllm/v1/kv_offload/cpu/policies/lru.py b/vllm/v1/kv_offload/cpu/policies/lru.py index 47e18f5565..efa24fe903 100644 --- a/vllm/v1/kv_offload/cpu/policies/lru.py +++ b/vllm/v1/kv_offload/cpu/policies/lru.py @@ -5,7 +5,7 @@ from collections.abc import Iterable from typing_extensions import override -from vllm.v1.kv_offload.base import OffloadKey +from vllm.v1.kv_offload.base import OffloadKey, ReqContext from vllm.v1.kv_offload.cpu.policies.base import BlockStatus, CachePolicy @@ -39,7 +39,7 @@ class LRUCachePolicy(CachePolicy): self.evictable_blocks.pop(key, None) @override - def touch(self, keys: Iterable[OffloadKey]) -> None: + def touch(self, keys: Iterable[OffloadKey], req_context: ReqContext) -> None: for key in reversed(list(keys)): if key in self.evictable_blocks: self.evictable_blocks.move_to_end(key)