diff --git a/tests/v1/kv_offload/cpu/test_manager.py b/tests/v1/kv_offload/cpu/test_manager.py index e043590a418..86cec459c3e 100644 --- a/tests/v1/kv_offload/cpu/test_manager.py +++ b/tests/v1/kv_offload/cpu/test_manager.py @@ -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) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py index 773fe8f056a..137eaef9788 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -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. diff --git a/vllm/v1/kv_offload/base.py b/vllm/v1/kv_offload/base.py index 3d403ea5083..fed47cf5993 100644 --- a/vllm/v1/kv_offload/base.py +++ b/vllm/v1/kv_offload/base.py @@ -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 diff --git a/vllm/v1/kv_offload/cpu/manager.py b/vllm/v1/kv_offload/cpu/manager.py index 80bcb568f99..dc501d74568 100644 --- a/vllm/v1/kv_offload/cpu/manager.py +++ b/vllm/v1/kv_offload/cpu/manager.py @@ -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] = [] diff --git a/vllm/v1/kv_offload/reuse_manager.py b/vllm/v1/kv_offload/reuse_manager.py index 6cb0a5f7591..d8c20b9dec7 100644 --- a/vllm/v1/kv_offload/reuse_manager.py +++ b/vllm/v1/kv_offload/reuse_manager.py @@ -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()