karylab_agents/vllm
1
0
Fork 0
forked from Karylab-cklius/vllm
Code Pull Requests 2 Actions 1 Packages Activity

[KV Offload] Pass ReqContext to touch(), complete_load(), and complete_store() (#41366)

Browse Source 
Signed-off-by: Ronen Schaffer <ronen.schaffer@ibm.com>
This commit is contained in:
Ronen Schaffer
2026-05-10 15:09:25 +03:00
committed by GitHub
parent a54f0d1049
commit e175192d33
5 changed files with 70 additions and 50 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/v1/kv_offload/cpu/test_manager.py
+32 -32
View File
@@ -117,10 +117,10 @@ def test_already_stored_block_not_evicted_during_prepare_store(eviction_policy):
# store [1, 2] and complete
manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX)
manager.complete_store(to_keys([1, 2]))
manager.complete_store(to_keys([1, 2]), _EMPTY_REQ_CTX)
# touch [1] to make block 2 the LRU candidate
manager.touch(to_keys([1]))
manager.touch(to_keys([1]), _EMPTY_REQ_CTX)
# prepare_store([2, 3, 4, 5]):
# - block 2 is already stored -> filtered out of keys_to_store
@@ -137,7 +137,7 @@ def test_already_stored_block_not_evicted_during_prepare_store(eviction_policy):
)
# complete_store must not silently drop block 2
manager.complete_store(to_keys([2, 3, 4, 5]))
manager.complete_store(to_keys([2, 3, 4, 5]), _EMPTY_REQ_CTX)
# block 2 must still be present in the cache
assert manager.lookup(to_key(2), _EMPTY_REQ_CTX) is True
@@ -171,7 +171,7 @@ def test_cpu_manager():
assert list(cpu_manager.take_events()) == []
# complete store [1, 2]
cpu_manager.complete_store(to_keys([1, 2]))
cpu_manager.complete_store(to_keys([1, 2]), _EMPTY_REQ_CTX)
verify_events(cpu_manager.take_events(), expected_stores=({1, 2},))
# lookup [1, 2]
@@ -199,7 +199,7 @@ def test_cpu_manager():
assert cpu_manager.prepare_store(to_keys([1, 6]), _EMPTY_REQ_CTX) is None
# complete store [2, 3, 4, 5]
cpu_manager.complete_store(to_keys([2, 3, 4, 5]))
cpu_manager.complete_store(to_keys([2, 3, 4, 5]), _EMPTY_REQ_CTX)
# lookup (now that we have [2, 3, 4, 5])
assert cpu_manager.lookup(to_key(1), _EMPTY_REQ_CTX) is False
@@ -217,7 +217,7 @@ def test_cpu_manager():
assert cpu_manager.prepare_store(to_keys([6, 7, 8]), _EMPTY_REQ_CTX) is None
# complete load [2, 3]
cpu_manager.complete_load(to_keys([2, 3]))
cpu_manager.complete_load(to_keys([2, 3]), _EMPTY_REQ_CTX)
# prepare store [6, 7, 8] -> evicts [2, 3, 4] (oldest)
prepare_store_output = cpu_manager.prepare_store(to_keys([6, 7, 8]), _EMPTY_REQ_CTX)
@@ -231,10 +231,10 @@ def test_cpu_manager():
)
# complete store [6, 7, 8]
cpu_manager.complete_store(to_keys([6, 7, 8]))
cpu_manager.complete_store(to_keys([6, 7, 8]), _EMPTY_REQ_CTX)
# touch [5, 6, 7] (move to end of LRU order)
cpu_manager.touch(to_keys([5, 6, 7]))
cpu_manager.touch(to_keys([5, 6, 7]), _EMPTY_REQ_CTX)
# prepare store [7, 9] -> evicts [8] (oldest following previous touch)
prepare_store_output = cpu_manager.prepare_store(to_keys([9]), _EMPTY_REQ_CTX)
@@ -248,7 +248,7 @@ def test_cpu_manager():
)
# complete store [7, 9] with failure
cpu_manager.complete_store(to_keys([7, 9]), success=False)
cpu_manager.complete_store(to_keys([7, 9]), _EMPTY_REQ_CTX, success=False)
# assert [7] is still stored, but [9] is not
assert cpu_manager.lookup(to_key(7), _EMPTY_REQ_CTX) is True
@@ -304,7 +304,7 @@ class TestARCPolicy:
assert list(cpu_manager.take_events()) == []
# complete store [1, 2]
cpu_manager.complete_store(to_keys([1, 2]))
cpu_manager.complete_store(to_keys([1, 2]), _EMPTY_REQ_CTX)
verify_events(cpu_manager.take_events(), expected_stores=({1, 2},))
# lookup [1, 2]
@@ -325,14 +325,14 @@ class TestARCPolicy:
# store and complete block 1
cpu_manager.prepare_store(to_keys([1]), _EMPTY_REQ_CTX)
cpu_manager.complete_store(to_keys([1]))
cpu_manager.complete_store(to_keys([1]), _EMPTY_REQ_CTX)
# block 1 starts in T1 (recent)
assert to_keys([1])[0] in arc_policy.t1
assert to_keys([1])[0] not in arc_policy.t2
# touch block 1 (simulate second access)
cpu_manager.touch(to_keys([1]))
cpu_manager.touch(to_keys([1]), _EMPTY_REQ_CTX)
# block 1 should now be in T2 (frequent)
assert to_keys([1])[0] not in arc_policy.t1
@@ -357,7 +357,7 @@ class TestARCPolicy:
evicted_keys=[],
),
)
cpu_manager.complete_store(to_keys([1, 2, 3, 4]))
cpu_manager.complete_store(to_keys([1, 2, 3, 4]), _EMPTY_REQ_CTX)
# prepare load [2, 3] (increases ref_cnt)
prepare_load_output = cpu_manager.prepare_load(to_keys([2, 3]), _EMPTY_REQ_CTX)
@@ -368,7 +368,7 @@ class TestARCPolicy:
assert cpu_manager.prepare_store(to_keys([5, 6, 7]), _EMPTY_REQ_CTX) is None
# complete load [2, 3]
cpu_manager.complete_load(to_keys([2, 3]))
cpu_manager.complete_load(to_keys([2, 3]), _EMPTY_REQ_CTX)
# now prepare store [5, 6, 7] should succeed
# ARC will evict blocks one at a time from T1 as needed
@@ -389,20 +389,20 @@ class TestARCPolicy:
# store blocks 1, 2 (fills cache)
cpu_manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX)
cpu_manager.complete_store(to_keys([1, 2]))
cpu_manager.complete_store(to_keys([1, 2]), _EMPTY_REQ_CTX)
initial_target = arc_policy.target_t1_size
# store block 3, evicting block 1 (moves to B1 ghost list)
cpu_manager.prepare_store(to_keys([3]), _EMPTY_REQ_CTX)
cpu_manager.complete_store(to_keys([3]))
cpu_manager.complete_store(to_keys([3]), _EMPTY_REQ_CTX)
# block 1 should be in B1 (ghost list)
assert to_keys([1])[0] in arc_policy.b1
# touch block 1 (cache miss, but in B1)
# this should increase target_t1_size (favor recency)
cpu_manager.touch(to_keys([1]))
cpu_manager.touch(to_keys([1]), _EMPTY_REQ_CTX)
# target should have increased
assert arc_policy.target_t1_size > initial_target
@@ -416,10 +416,10 @@ class TestARCPolicy:
# store blocks 1, 2, 3, 4
cpu_manager.prepare_store(to_keys([1, 2, 3, 4]), _EMPTY_REQ_CTX)
cpu_manager.complete_store(to_keys([1, 2, 3, 4]))
cpu_manager.complete_store(to_keys([1, 2, 3, 4]), _EMPTY_REQ_CTX)
# promote blocks 3, 4 to T2 by touching them
cpu_manager.touch(to_keys([3, 4]))
cpu_manager.touch(to_keys([3, 4]), _EMPTY_REQ_CTX)
# now: T1 = {1, 2}, T2 = {3, 4}
assert len(arc_policy.t1) == 2
@@ -434,7 +434,7 @@ class TestARCPolicy:
assert output is not None
assert to_keys([1]) == output.evicted_keys
cpu_manager.complete_store(to_keys([5]))
cpu_manager.complete_store(to_keys([5]), _EMPTY_REQ_CTX)
# block 1 should be in B1 (ghost list)
assert to_keys([1])[0] in arc_policy.b1
@@ -450,12 +450,12 @@ class TestARCPolicy:
# fill cache with blocks 1, 2
cpu_manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX)
cpu_manager.complete_store(to_keys([1, 2]))
cpu_manager.complete_store(to_keys([1, 2]), _EMPTY_REQ_CTX)
# store many blocks to fill ghost lists
for i in range(3, 20):
cpu_manager.prepare_store(to_keys([i]), _EMPTY_REQ_CTX)
cpu_manager.complete_store(to_keys([i]))
cpu_manager.complete_store(to_keys([i]), _EMPTY_REQ_CTX)
# ghost lists should not exceed cache_capacity
assert len(arc_policy.b1) <= arc_policy.cache_capacity
@@ -470,14 +470,14 @@ class TestARCPolicy:
# store blocks 1, 2, 3, 4
cpu_manager.prepare_store(to_keys([1, 2, 3, 4]), _EMPTY_REQ_CTX)
cpu_manager.complete_store(to_keys([1, 2, 3, 4]))
cpu_manager.complete_store(to_keys([1, 2, 3, 4]), _EMPTY_REQ_CTX)
# promote 3, 4 to T2
cpu_manager.touch(to_keys([3, 4]))
cpu_manager.touch(to_keys([3, 4]), _EMPTY_REQ_CTX)
# T1 = {1, 2}, T2 = {3, 4}
# touch [1, 3, 4] - should promote 1 to T2, and move 3,4 to end of T2
cpu_manager.touch(to_keys([1, 3, 4]))
cpu_manager.touch(to_keys([1, 3, 4]), _EMPTY_REQ_CTX)
# T1 = {2}, T2 = {1, 3, 4} (in that order, with 4 most recent)
assert len(arc_policy.t1) == 1
@@ -503,7 +503,7 @@ class TestARCPolicy:
# store blocks 1, 2, 3, 4
cpu_manager.prepare_store(to_keys([1, 2, 3, 4]), _EMPTY_REQ_CTX)
cpu_manager.complete_store(to_keys([1, 2, 3, 4]))
cpu_manager.complete_store(to_keys([1, 2, 3, 4]), _EMPTY_REQ_CTX)
# prepare store block 5 (will evict block 1)
prepare_store_output = cpu_manager.prepare_store(to_keys([5]), _EMPTY_REQ_CTX)
@@ -511,7 +511,7 @@ class TestARCPolicy:
assert len(prepare_store_output.evicted_keys) == 1
# complete store with failure
cpu_manager.complete_store(to_keys([5]), success=False)
cpu_manager.complete_store(to_keys([5]), _EMPTY_REQ_CTX, success=False)
# block 5 should not be in cache
assert cpu_manager.lookup(to_key(5), _EMPTY_REQ_CTX) is False
@@ -532,7 +532,7 @@ class TestARCPolicy:
# store [1, 2]
cpu_manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX)
cpu_manager.complete_store(to_keys([1, 2]))
cpu_manager.complete_store(to_keys([1, 2]), _EMPTY_REQ_CTX)
# store [3, 4, 5] -> evicts [1]
prepare_store_output = cpu_manager.prepare_store(
@@ -540,10 +540,10 @@ class TestARCPolicy:
)
assert prepare_store_output is not None
assert len(prepare_store_output.evicted_keys) == 1
cpu_manager.complete_store(to_keys([3, 4, 5]))
cpu_manager.complete_store(to_keys([3, 4, 5]), _EMPTY_REQ_CTX)
# promote some blocks to T2
cpu_manager.touch(to_keys([2, 3]))
cpu_manager.touch(to_keys([2, 3]), _EMPTY_REQ_CTX)
# T1 has {4, 5}, T2 has {2, 3}
assert len(arc_policy.t1) == 2
@@ -552,7 +552,7 @@ class TestARCPolicy:
# store [6] -> should evict from T1 (4 is oldest in T1)
prepare_store_output = cpu_manager.prepare_store(to_keys([6]), _EMPTY_REQ_CTX)
assert prepare_store_output is not None
cpu_manager.complete_store(to_keys([6]))
cpu_manager.complete_store(to_keys([6]), _EMPTY_REQ_CTX)
# verify blocks 2, 3 (in T2) are still present
assert cpu_manager.lookup(to_key(2), _EMPTY_REQ_CTX) is True
@@ -609,4 +609,4 @@ def test_filter_reused_manager():
assert prepare_store_output is not None
assert prepare_store_output.keys_to_store == []
manager.complete_store(to_keys([1]))
manager.complete_store(to_keys([1]), _EMPTY_REQ_CTX)
vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py
+8 -6
View File
@@ -291,7 +291,7 @@ class OffloadingConnectorScheduler:
self.config.kv_group_configs, req_status.group_states
):
if group_config.sliding_window_size_in_blocks is None:
self.manager.touch(group_state.offload_keys)
self.manager.touch(group_state.offload_keys, req_status.req_context)
else:
# we aim to keep just blocks that are necessary to hit
# the original request (+ decoded blocks)
@@ -300,7 +300,10 @@ class OffloadingConnectorScheduler:
group_state.num_hit_blocks
- group_config.sliding_window_size_in_blocks,
)
self.manager.touch(group_state.offload_keys[blocks_to_skip:])
self.manager.touch(
group_state.offload_keys[blocks_to_skip:],
req_status.req_context,
)
def _lookup(self, req_status: RequestOffloadState) -> int | None:
"""
@@ -802,14 +805,13 @@ class OffloadingConnectorScheduler:
continue
assert job_status.pending_count == 0
req_status = self._req_status[job_status.req_id]
if job_status.is_store:
self.manager.complete_store(job_status.keys)
self.manager.complete_store(job_status.keys, req_status.req_context)
else:
self.manager.complete_load(job_status.keys)
self.manager.complete_load(job_status.keys, req_status.req_context)
if self._blocks_being_loaded:
self._blocks_being_loaded.difference_update(job_status.keys)
req_status = self._req_status[job_status.req_id]
if self._block_id_to_pending_jobs:
# Sliding window blocks are tracked from store creation
# and must be cleaned up unconditionally.
vllm/v1/kv_offload/base.py
+11 -3
View File
@@ -147,22 +147,24 @@ class OffloadingManager(ABC):
"""
pass
def touch(self, keys: Collection[OffloadKey]):
def touch(self, keys: Collection[OffloadKey], req_context: ReqContext):
"""
Mark the given blocks as recently used.
This could in practice mean moving them to the end of an LRU list.
Args:
keys: the keys identifying the blocks.
req_context: per-request context (e.g. kv_transfer_params).
"""
return
def complete_load(self, keys: Collection[OffloadKey]):
def complete_load(self, keys: Collection[OffloadKey], req_context: ReqContext):
"""
Marks previous blocks that were prepared to load as done loading.
Args:
keys: the keys identifying the blocks.
req_context: per-request context (e.g. kv_transfer_params).
"""
return
@@ -189,7 +191,12 @@ class OffloadingManager(ABC):
"""
pass
def complete_store(self, keys: Collection[OffloadKey], success: bool = True):
def complete_store(
self,
keys: Collection[OffloadKey],
req_context: ReqContext,
success: bool = True,
):
"""
Marks blocks which were previously prepared to be stored, as stored.
Following this call, the blocks become loadable.
@@ -198,6 +205,7 @@ class OffloadingManager(ABC):
Args:
keys: the keys identifying the blocks.
req_context: per-request context (e.g. kv_transfer_params).
success: whether the blocks were stored successfully.
"""
return
vllm/v1/kv_offload/cpu/manager.py
+8 -3
View File
@@ -106,10 +106,12 @@ class CPUOffloadingManager(OffloadingManager):
blocks.append(block)
return self._get_load_store_spec(keys, blocks)
def touch(self, keys: Collection[OffloadKey]) -> None:
def touch(self, keys: Collection[OffloadKey], req_context: ReqContext) -> None:
self._policy.touch(keys)
def complete_load(self, keys: Collection[OffloadKey]) -> None:
def complete_load(
self, keys: Collection[OffloadKey], req_context: ReqContext
) -> None:
for key in keys:
block = self._policy.get(key)
assert block is not None, f"Block {key!r} not found"
@@ -172,7 +174,10 @@ class CPUOffloadingManager(OffloadingManager):
)
def complete_store(
self, keys: Collection[OffloadKey], success: bool = True
self,
keys: Collection[OffloadKey],
req_context: ReqContext,
success: bool = True,
) -> None:
stored_keys: list[OffloadKey] = []
vllm/v1/kv_offload/reuse_manager.py
+11 -6
View File
@@ -105,16 +105,21 @@ class FilterReusedOffloadingManager(OffloadingManager):
) -> LoadStoreSpec:
return self._backing.prepare_load(keys, req_context)
def touch(self, keys: Collection[OffloadKey]) -> None:
return self._backing.touch(keys)
def touch(self, keys: Collection[OffloadKey], req_context: ReqContext) -> None:
return self._backing.touch(keys, req_context)
def complete_load(self, keys: Collection[OffloadKey]) -> None:
return self._backing.complete_load(keys)
def complete_load(
self, keys: Collection[OffloadKey], req_context: ReqContext
) -> None:
return self._backing.complete_load(keys, req_context)
def complete_store(
self, keys: Collection[OffloadKey], success: bool = True
self,
keys: Collection[OffloadKey],
req_context: ReqContext,
success: bool = True,
) -> None:
return self._backing.complete_store(keys, success)
return self._backing.complete_store(keys, req_context, success)
def take_events(self) -> Iterable[OffloadingEvent]:
return self._backing.take_events()