[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>
Posedge_Lin committed
Jul 8, 2026 at 10:56 UTC
0ca6eee7433893e94f3d4be00cedb3e9f3d6c44f
5 files changed
+37
-8
tests/v1/kv_offload/cpu/test_manager.py
+23
@@ -897,3 +897,26 @@ def test_evictable_cache_block_count():
897
manager.complete_store(to_keys([14, 15]), _EMPTY_REQ_CTX)
898
# cache state [10, 11, 14, 15] <- all blocks idle
899
assert manager._num_evictable_cache_blocks == 4
900
+
901
+
902
+def test_touch_forwards_req_context_to_policy(monkeypatch):
903
+ """Regression: CPUOffloadingManager.touch forwards ReqContext to policy."""
904
+ manager = make_cpu_manager(num_blocks=4, cache_policy="lru")
905
+ received = []
906
+
907
+ def spy_touch(keys: Iterable[OffloadKey], req_context: ReqContext) -> None:
908
+ received.append((list(keys), req_context))
909
+
910
+ monkeypatch.setattr(manager._policy, "touch", spy_touch)
911
+
912
+ keys = to_keys([1, 2])
913
+ ctx = make_req_context(
914
+ req_id="test-req",
915
+ kv_transfer_params={"test_param": "test_value"},
916
+ )
917
+
918
+ manager.touch(keys, ctx)
919
+
920
+ assert len(received) == 1
921
+ assert received[0][0] == keys
922
+ assert received[0][1] is ctx
vllm/v1/kv_offload/cpu/manager.py
+1
-1
@@ -152,7 +152,7 @@ class CPUOffloadingManager(OffloadingManager):
152
153
@override
154
def touch(self, keys: Collection[OffloadKey], req_context: ReqContext) -> None:
155
- self._policy.touch(keys)
155
+ self._policy.touch(keys, req_context)
156
157
@override
158
def complete_load(
vllm/v1/kv_offload/cpu/policies/arc.py
+2
-2
@@ -5,7 +5,7 @@ from collections.abc import Iterable
5
6
from typing_extensions import override
7
8
-from vllm.v1.kv_offload.base import OffloadKey
8
+from vllm.v1.kv_offload.base import OffloadKey, ReqContext
9
from vllm.v1.kv_offload.cpu.policies.base import BlockStatus, CachePolicy
10
11
@@ -72,7 +72,7 @@ class ARCCachePolicy(CachePolicy):
72
self.t2.pop(key, None)
73
74
@override
75
- def touch(self, keys: Iterable[OffloadKey]) -> None:
75
+ def touch(self, keys: Iterable[OffloadKey], req_context: ReqContext) -> None:
76
for key in reversed(list(keys)):
77
if key in self.t1:
78
block = self.t1.pop(key)
vllm/v1/kv_offload/cpu/policies/base.py
+9
-3
@@ -4,7 +4,7 @@ import ctypes
4
from abc import ABC, abstractmethod
5
from collections.abc import Iterable
6
7
-from vllm.v1.kv_offload.base import OffloadKey
7
+from vllm.v1.kv_offload.base import OffloadKey, ReqContext
8
9
10
class BlockStatus(ctypes.Structure):
@@ -57,8 +57,14 @@ class CachePolicy(ABC):
57
"""Remove a block (used to clean up after a failed store)."""
58
59
@abstractmethod
60
- def touch(self, keys: Iterable[OffloadKey]) -> None:
61
- """Mark blocks as recently used."""
60
+ def touch(self, keys: Iterable[OffloadKey], req_context: ReqContext) -> None:
61
+ """
62
+ Mark blocks as recently used.
63
+
64
+ Args:
65
+ keys: Blocks to mark as recently used.
66
+ req_context: Per-request context for the request touching these blocks.
67
+ """
68
69
@abstractmethod
70
def evict(
vllm/v1/kv_offload/cpu/policies/lru.py
+2
-2
@@ -5,7 +5,7 @@ from collections.abc import Iterable
5
6
from typing_extensions import override
7
8
-from vllm.v1.kv_offload.base import OffloadKey
8
+from vllm.v1.kv_offload.base import OffloadKey, ReqContext
9
from vllm.v1.kv_offload.cpu.policies.base import BlockStatus, CachePolicy
10
11
@@ -39,7 +39,7 @@ class LRUCachePolicy(CachePolicy):
39
self.evictable_blocks.pop(key, None)
40
41
@override
42
- def touch(self, keys: Iterable[OffloadKey]) -> None:
42
+ def touch(self, keys: Iterable[OffloadKey], req_context: ReqContext) -> None:
43
for key in reversed(list(keys)):
44
if key in self.evictable_blocks:
45
self.evictable_blocks.move_to_end(key)