have KVCacheManager return empty blocks for nonexistent requests

Signed-off-by: Nick Hill <nhill@redhat.com>

Author: njhill

vllm/v1/core/kv_cache_coordinator.py

@@ -166,6 +166,7 @@ def get_blocks(self, request_id: str) -> list[list[KVCacheBlock]]:
         return [
             manager.req_to_blocks[request_id]
             for manager in self.single_type_managers
+            if request_id in manager.req_to_blocks
         ]
 
     @abstractmethod

vllm/v1/core/kv_cache_manager.py

@@ -388,3 +388,4 @@ def cache_blocks(self, request: Request, block_hashes: list[BlockHash],
     def create_empty_block_list(self) -> KVCacheBlocks:
         """Creates a new KVCacheBlocks instance with no blocks."""
         return KVCacheBlocks([[] for _ in range(self.num_kv_cache_groups)])
+

vllm/v1/core/sched/scheduler.py

@@ -76,6 +76,9 @@ def __init__(
         # KV Connector pushes/pull of remote KVs for P/D and offloading.
         self.connector = None
         if self.vllm_config.kv_transfer_config is not None:
+            assert len(self.kv_cache_config.kv_cache_groups) == 1, (
+                "Multiple KV cache groups are not currently supported "
+                "with KV connectors")
             self.connector = KVConnectorFactory.create_connector_v1(
                 config=self.vllm_config, role=KVConnectorRole.SCHEDULER)
 
@@ -985,16 +988,8 @@ def _connector_finished(
         """
         if self.connector is None:
             return False, None
-        assert len(self.kv_cache_config.kv_cache_groups
-                   ) == 1, "KV connector only supports one KV cache group now"
-        if (request.status == RequestStatus.FINISHED_ABORTED and \
-            request.request_id not in
-            self.kv_cache_manager.single_type_manager.req_to_blocks):
-            # Ensure empty blocks ids are passed to respect connector interface
-            block_ids = KVCacheBlocks.create_empty().get_block_ids()[0]
-        else:
-            block_ids = self.kv_cache_manager.get_block_ids(
-                request.request_id)[0]
+
+        (block_ids, ) = self.kv_cache_manager.get_block_ids(request.request_id)
         return self.connector.request_finished(request, block_ids)
 
     def _update_waiting_for_remote_kv(self, request: Request) -> bool:
@@ -1009,12 +1004,12 @@ def _update_waiting_for_remote_kv(self, request: Request) -> bool:
         and the request state will be moved back to WAITING from
         WAITING_FOR_REMOTE_KV.
         """
+        assert self.connector is not None
         if request.request_id not in self.finished_recving_kv_req_ids:
             return False
-        assert len(self.kv_cache_config.kv_cache_groups
-                   ) == 1, "KV connector only supports one KV cache group now"
+
         # Now that the blocks are ready, actually cache them.
-        block_ids = self.kv_cache_manager.get_block_ids(request.request_id)[0]
+        (block_ids, ) = self.kv_cache_manager.get_block_ids(request.request_id)
         num_computed_tokens = len(block_ids) * self.block_size
         # Handle the case where num request tokens less then one block.
         num_computed_tokens = min(num_computed_tokens, request.num_tokens)