v1: Pass KVConnectorOutput to scheduler-side (#22157)

Signed-off-by: Or Ozeri <[email protected]>

Author: orozery

vllm/distributed/kv_transfer/kv_connector/v1/base.py

@@ -12,6 +12,8 @@
             times for a given request and should be side-effect free.
         update_state_after_alloc() - update KVConnector state after
             temporary buffer alloc by the CacheManager.
+        update_connector_output() - update KVConnector state after
+            output is received from worker-side connectors.
         request_finished() - called when a request is finished, with
             the computed kv cache blocks for the request.
             Returns whether KV cache should be freed now or will be
@@ -38,6 +40,7 @@
 
 from vllm.logger import init_logger
 from vllm.v1.core.sched.output import SchedulerOutput
+from vllm.v1.outputs import KVConnectorOutput
 
 if TYPE_CHECKING:
     from vllm.attention.backends.abstract import AttentionMetadata
@@ -283,6 +286,16 @@ def build_connector_meta(
         """
         pass
 
+    def update_connector_output(self, connector_output: KVConnectorOutput):
+        """
+        Update KVConnector state from worker-side connectors output.
+
+        Args:
+            connector_output (KVConnectorOutput): the worker-side
+                connectors output.
+        """
+        return
+
     def request_finished(
         self,
         request: "Request",

vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py

@@ -14,6 +14,7 @@
 from vllm.logger import init_logger
 from vllm.v1.core.kv_cache_manager import KVCacheBlocks
 from vllm.v1.core.sched.output import SchedulerOutput
+from vllm.v1.outputs import KVConnectorOutput
 
 if TYPE_CHECKING:
     from vllm.attention.backends.abstract import AttentionMetadata
@@ -177,6 +178,10 @@ def build_connector_meta(
             self._extra_async_saves = {}
         return metadata
 
+    def update_connector_output(self, connector_output: KVConnectorOutput):
+        for c in self._connectors:
+            c.update_connector_output(connector_output)
+
     def request_finished(
         self,
         request: "Request",

vllm/v1/core/sched/scheduler.py

@@ -1150,6 +1150,10 @@ def _update_from_kv_xfer_finished(self,
         # if finished_recving: add to state so we can
             scheduler the request during the next step.
         """
+
+        assert self.connector is not None
+        self.connector.update_connector_output(kv_connector_output)
+
         # KV Connector:: update recv and send status from last step.
         for req_id in (kv_connector_output.finished_recving or ()):
             logger.debug("Finished recving KV transfer for request %s", req_id)