From 194c4a92760d22796e69644813b9967d1f398c5c Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Fri, 6 Jun 2025 10:21:02 -0700 Subject: [PATCH 1/2] [BugFix] Fix MultiConnector test after HMA changes Signed-off-by: Nick Hill --- .../v1/kv_connector/unit/test_multi_connector.py | 15 ++++++++------- .../kv_connector/v1/multi_connector.py | 4 ++-- vllm/v1/core/kv_cache_manager.py | 4 ++++ 3 files changed, 14 insertions(+), 9 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_multi_connector.py b/tests/v1/kv_connector/unit/test_multi_connector.py index d90ddcdbbbc8..d48037941dfc 100644 --- a/tests/v1/kv_connector/unit/test_multi_connector.py +++ b/tests/v1/kv_connector/unit/test_multi_connector.py @@ -60,7 +60,8 @@ def wrapper(*args, **kwargs): if isinstance(arg, int): to_log.append(str(arg)) elif isinstance(arg, KVCacheBlocks): - to_log.append(f"num_blocks={len(arg.blocks)}") + to_log.append(f"num_blocks={ + [len(b) for b in arg.blocks]}") # Log the event as a line to the file try: @@ -176,7 +177,7 @@ def test_multi_shared_storage_connector_consistency(): # on each connector in turn. assert events["storage1-SCHEDULER"][:3] == [ 'get_num_new_matched_tokens 0', - 'update_state_after_alloc num_blocks=0 0', 'build_connector_meta' + 'update_state_after_alloc num_blocks=[0] 0', 'build_connector_meta' ] assert events["storage1-WORKER"][:5] == [ 'register_kv_caches', 'bind_connector_metadata', 'start_load_kv', @@ -184,7 +185,7 @@ def test_multi_shared_storage_connector_consistency(): ] assert events["storage2-SCHEDULER"][:3] == [ 'get_num_new_matched_tokens 0', - 'update_state_after_alloc num_blocks=0 0', 'build_connector_meta' + 'update_state_after_alloc num_blocks=[0] 0', 'build_connector_meta' ] assert events["storage2-WORKER"][:5] == [ 'register_kv_caches', 'bind_connector_metadata', 'start_load_kv', @@ -205,11 +206,11 @@ def test_multi_shared_storage_connector_consistency(): # chosen). assert events["storage1-SCHEDULER"][:3] == [ 'get_num_new_matched_tokens 0', - 'update_state_after_alloc num_blocks=7 96', 'build_connector_meta' + 'update_state_after_alloc num_blocks=[7] 96', 'build_connector_meta' ] assert events["storage2-SCHEDULER"][:3] == [ 'get_num_new_matched_tokens 0', - 'update_state_after_alloc num_blocks=0 0', 'build_connector_meta' + 'update_state_after_alloc num_blocks=[0] 0', 'build_connector_meta' ] # Delete storage1 connector state @@ -229,11 +230,11 @@ def test_multi_shared_storage_connector_consistency(): # blocks for the second connector. assert events["storage1-SCHEDULER"][:3] == [ 'get_num_new_matched_tokens 0', - 'update_state_after_alloc num_blocks=0 0', 'build_connector_meta' + 'update_state_after_alloc num_blocks=[0] 0', 'build_connector_meta' ] assert events["storage2-SCHEDULER"][:3] == [ 'get_num_new_matched_tokens 0', - 'update_state_after_alloc num_blocks=7 96', 'build_connector_meta' + 'update_state_after_alloc num_blocks=[7] 96', 'build_connector_meta' ] # Clean up diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py index b2cf88f5b83d..be3c23399419 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py @@ -153,6 +153,7 @@ def update_state_after_alloc(self, request: "Request", num_external_tokens: int): chosen_connector = self._requests_to_connector.get( request.request_id, -1) + empty_blocks = blocks.new_empty() for i, c in enumerate(self._connectors): if i == chosen_connector: # Forward call to the chosen connector (if any). @@ -160,8 +161,7 @@ def update_state_after_alloc(self, request: "Request", num_external_tokens) else: # Call with empty blocks for other connectors. - c.update_state_after_alloc(request, - KVCacheBlocks.create_empty(), 0) + c.update_state_after_alloc(request, empty_blocks, 0) def build_connector_meta( self, diff --git a/vllm/v1/core/kv_cache_manager.py b/vllm/v1/core/kv_cache_manager.py index fc701215ba5d..35fb189fda34 100644 --- a/vllm/v1/core/kv_cache_manager.py +++ b/vllm/v1/core/kv_cache_manager.py @@ -61,6 +61,10 @@ def get_unhashed_block_ids(self) -> list[int]: if block.block_hash is None ] + def new_empty(self) -> "KVCacheBlocks": + """Creates a new KVCacheBlocks instance with no blocks.""" + return KVCacheBlocks([[] for _ in range(len(self.blocks))]) + class KVCacheManager: From ee31f97b6b1d4907a6fe64e04391f30e9cce31dc Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Fri, 6 Jun 2025 10:39:24 -0700 Subject: [PATCH 2/2] yapf error Signed-off-by: Nick Hill --- tests/v1/kv_connector/unit/test_multi_connector.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_multi_connector.py b/tests/v1/kv_connector/unit/test_multi_connector.py index d48037941dfc..72848c1a706e 100644 --- a/tests/v1/kv_connector/unit/test_multi_connector.py +++ b/tests/v1/kv_connector/unit/test_multi_connector.py @@ -60,8 +60,8 @@ def wrapper(*args, **kwargs): if isinstance(arg, int): to_log.append(str(arg)) elif isinstance(arg, KVCacheBlocks): - to_log.append(f"num_blocks={ - [len(b) for b in arg.blocks]}") + to_log.append( + f"num_blocks={[len(b) for b in arg.blocks]}") # Log the event as a line to the file try: