[Core] Pass request context to CPU offload cache policy touch (#47744)

Signed-off-by: jacklin78911-collab <jacklin78911@gmail.com>
Co-authored-by: Or Ozeri <oro@il.ibm.com>

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

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(

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)

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(

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)