v1: Support worker -> scheduler KV connector MD

This commit makes the following changes:
1. Add a new kv_connector_metadata to ModelRunnerOutput
    to allow arbitrary connector metadata flow from workers to the scheduler.
2. Add a new worker-side connector API to build the above metadata.
3. Change MultiprocExecutor to get ModelRunnerOutput from all workers,
    and aggregate the kv_connector_metadata from all.
3. Move the get_finished connector API from the worker side to
    the scheduler side.
4. Change the nixl and multi connectors to match the above API changes.

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

Author: orozery

tests/v1/kv_connector/unit/test_remote_decode_lifecycle.py

@@ -1,12 +1,11 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-import copy
-
 from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT
 from vllm.v1.request import FinishReason, RequestStatus
 
-from .utils import (assert_scheduler_empty, create_model_runner_output,
-                    create_request, create_scheduler, create_vllm_config)
+from .utils import (assert_scheduler_empty, create_empty_model_runner_output,
+                    create_model_runner_output, create_request,
+                    create_scheduler, create_vllm_config)
 
 
 def test_basic_lifecycle():
@@ -85,8 +84,8 @@ def test_basic_lifecycle():
     assert len(scheduler.finished_req_ids) == 0
 
     # (3b): execute_model()
-    model_runner_output = copy.deepcopy(EMPTY_MODEL_RUNNER_OUTPUT)
-    model_runner_output.finished_sending = [request_id]
+    model_runner_output = create_empty_model_runner_output(
+        finished_sending=[request_id])
 
     # (3c): update_from_output()
     scheduler.update_from_output(scheduler_output, model_runner_output)
@@ -175,8 +174,8 @@ def test_prefix_cache_lifecycle():
     # STEP (2): Ensure it is freed.
     scheduler_output = scheduler.schedule()
     scheduler.schedule()
-    model_runner_output = copy.deepcopy(EMPTY_MODEL_RUNNER_OUTPUT)
-    model_runner_output.finished_sending = [request_remote.request_id]
+    model_runner_output = create_empty_model_runner_output(
+        finished_sending=[request_remote.request_id])
     scheduler.update_from_output(scheduler_output, model_runner_output)
     _ = scheduler.schedule()
     assert_scheduler_empty(scheduler)

tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py

@@ -1,12 +1,11 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-import copy
-
 from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT
 from vllm.v1.request import FinishReason, RequestStatus
 
-from .utils import (assert_scheduler_empty, create_model_runner_output,
-                    create_request, create_scheduler, create_vllm_config)
+from .utils import (assert_scheduler_empty, create_empty_model_runner_output,
+                    create_model_runner_output, create_request,
+                    create_scheduler, create_vllm_config)
 
 
 def test_basic_lifecycle():
@@ -71,8 +70,8 @@ def test_basic_lifecycle():
     assert len(scheduler.running) == 0
 
     # (2b): forward(): request finishes recv.
-    model_runner_output = copy.deepcopy(EMPTY_MODEL_RUNNER_OUTPUT)
-    model_runner_output.finished_recving = [request_id]
+    model_runner_output = create_empty_model_runner_output(
+        finished_recving=[request_id])
 
     # (2c): update_from_output():
     engine_core_outputs = scheduler.update_from_output(scheduler_output,
@@ -308,8 +307,8 @@ def test_full_block_prompt():
 
     # # STEP (2): Recv.
     scheduler_output = scheduler.schedule()
-    model_runner_output = copy.deepcopy(EMPTY_MODEL_RUNNER_OUTPUT)
-    model_runner_output.finished_recving = [request_id]
+    model_runner_output = create_empty_model_runner_output(
+        finished_recving=[request_id])
     scheduler.update_from_output(scheduler_output, model_runner_output)
     assert len(scheduler.waiting) == 1
     assert (request_id in scheduler.finished_recving_kv_req_ids)

tests/v1/kv_connector/unit/utils.py

@@ -1,16 +1,19 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import copy
 from typing import Any, Optional
 
 import torch
 
 from vllm import SamplingParams
 from vllm.config import (CacheConfig, DeviceConfig, KVTransferConfig,
                          ModelConfig, SchedulerConfig, VllmConfig)
+from vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector import (
+    NixlWorkerConnectorMetadata)
 from vllm.v1.core.sched.scheduler import Scheduler
 from vllm.v1.kv_cache_interface import (FullAttentionSpec, KVCacheConfig,
                                         KVCacheGroupSpec)
-from vllm.v1.outputs import ModelRunnerOutput
+from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT, ModelRunnerOutput
 from vllm.v1.request import Request
 from vllm.v1.structured_output import StructuredOutputManager
 
@@ -175,6 +178,14 @@ def create_model_runner_output(
     sampled_token = EOS_TOKEN_ID if use_eos else 0
     sampled_token_ids = [[sampled_token] for _ in req_ids]
 
+    # Make worker connector metadata
+    kv_connector_metadata = None
+    if finished_sending or finished_recving:
+        kv_connector_metadata = [
+            NixlWorkerConnectorMetadata(finished_sending or [],
+                                        finished_recving or [])
+        ]
+
     # Make output data structure.
     return ModelRunnerOutput(
         req_ids=req_ids,
@@ -183,6 +194,21 @@ def create_model_runner_output(
         spec_token_ids=None,
         logprobs=None,
         prompt_logprobs_dict={},
-        finished_sending=finished_sending,
-        finished_recving=finished_recving,
+        kv_connector_metadata=kv_connector_metadata,
     )
+
+
+def create_empty_model_runner_output(
+    finished_sending: Optional[list[str]] = None,
+    finished_recving: Optional[list[str]] = None,
+) -> ModelRunnerOutput:
+    """Make dummy empty model runner output for testing."""
+    model_runner_output = copy.deepcopy(EMPTY_MODEL_RUNNER_OUTPUT)
+
+    kv_connector_metadata = [
+        NixlWorkerConnectorMetadata(finished_sending or [], finished_recving
+                                    or [])
+    ]
+    model_runner_output.kv_connector_metadata = kv_connector_metadata
+
+    return model_runner_output

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

@@ -17,6 +17,8 @@
             Returns whether KV cache should be freed now or will be
             freed asynchronously and optionally returns KV transfer
             params.
+        get_finished() - returns ids of requests that have completed
+            async sending/recving.
 
     Worker-side: runs in each worker, loads/saves KV cache to/from
     the Connector based on the metadata.
@@ -26,8 +28,8 @@
         save_kv_layer() - starts saving KV for layer i (maybe async)
         wait_for_save() - blocks until all saves are done
 
-        get_finished() - called with ids of finished requests, returns
-            ids of requests that have completed async sending/recving.
+        build_worker_connector_meta() - builds metadata to be sent
+            back to the scheduler.
 """
 
 import enum
@@ -38,6 +40,7 @@
 
 from vllm.logger import init_logger
 from vllm.v1.core.sched.output import SchedulerOutput
+from vllm.v1.outputs import ModelRunnerOutput
 
 if TYPE_CHECKING:
     from vllm.attention.backends.abstract import AttentionMetadata
@@ -185,21 +188,21 @@ def wait_for_save(self):
         """
         pass
 
-    def get_finished(
-        self, finished_req_ids: set[str]
-    ) -> tuple[Optional[set[str]], Optional[set[str]]]:
+    def build_worker_connector_meta(
+        self, scheduler_output: SchedulerOutput,
+        model_runner_output: ModelRunnerOutput
+    ) -> Optional[KVConnectorMetadata]:
         """
-        Notifies worker-side connector ids of requests that have
-        finished generating tokens.
+        Build the worker->scheduler connector metadata for this step.
 
-        Returns:
-            ids of requests that have finished asynchronous transfer
-            (requests that previously returned True from request_finished()),
-            tuple of (sending/saving ids, recving/loading ids).
-            The finished saves/sends req ids must belong to a set provided in a
-            call to this method (this call or a prior one).
+        This function should NOT modify fields of its arguments.
+
+        Args:
+            scheduler_output (SchedulerOutput): the scheduler output object.
+            model_runner_output (ModelRunnerOutput):
+                the model runner (worker) output object.
         """
-        return None, None
+        return None
 
     # ==============================
     # Scheduler-side methods
@@ -281,3 +284,21 @@ def request_finished(
             returned by the engine.
         """
         return False, None
+
+    def get_finished(
+        self,
+        model_runner_output: ModelRunnerOutput,
+    ) -> tuple[Optional[set[str]], Optional[set[str]]]:
+        """
+        Get request IDs that recently finished async transfer.
+
+        Args:
+            model_runner_output (ModelRunnerOutput):
+                the model runner (worker) output object.
+
+        Returns:
+            ids of requests that have finished asynchronous transfer
+            (requests that previously returned True from request_finished()),
+            tuple of (sending/saving ids, recving/loading ids).
+        """
+        return None, None

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 ModelRunnerOutput
 
 if TYPE_CHECKING:
     from vllm.attention.backends.abstract import AttentionMetadata
@@ -26,7 +27,11 @@
 @dataclass
 class MultiKVConnectorMetadata(KVConnectorMetadata):
     metadata: tuple[KVConnectorMetadata, ...]
-    extra_async_saves: Optional[dict[str, int]] = None
+
+
+@dataclass
+class MultiKVWorkerConnectorMetadata(KVConnectorMetadata):
+    metadata: tuple[Optional[KVConnectorMetadata], ...]
 
 
 class MultiConnector(KVConnectorBase_V1):
@@ -58,7 +63,6 @@ def __init__(self, vllm_config: "VllmConfig", role: KVConnectorRole):
         # Keeps track of *additional* remaining async saves (beyond 1) to be
         # finished per request. Not needed for async loads since we only allow
         # a single connector to load.
-        # Propagated from scheduler to worker side via the connector metadata.
         self._extra_async_saves: dict[str, int] = {}
 
     def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
@@ -71,9 +75,6 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
     def bind_connector_metadata(
             self, connector_metadata: KVConnectorMetadata) -> None:
         assert isinstance(connector_metadata, MultiKVConnectorMetadata)
-        if connector_metadata.extra_async_saves:
-            self._extra_async_saves.update(
-                connector_metadata.extra_async_saves)
         for c, cm in zip(self._connectors, connector_metadata.metadata):
             c.bind_connector_metadata(cm)
 
@@ -102,32 +103,14 @@ def wait_for_save(self):
         for c in self._connectors:
             c.wait_for_save()
 
-    def get_finished(
-        self, finished_req_ids: set[str]
-    ) -> tuple[Optional[set[str]], Optional[set[str]]]:
-        finished_sending: set[str] = set()
-        finished_recving: set[str] = set()
-        for c in self._connectors:
-            sending, recving = c.get_finished(finished_req_ids)
-            if not recving and not sending:
-                continue
-            # Aggregate finished recving request ids.
-            finished_recving.update(recving or ())
-            # Aggregate finished sending request ids - only include
-            # once we've drained the "extra" count (for cases where
-            # more than one connector is async-saving the same request).
-            for req_id in sending or ():
-                extra_pending = self._extra_async_saves.get(req_id)
-                if extra_pending is None:
-                    finished_sending.add(req_id)
-                    continue
-                assert extra_pending > 0
-                if extra_pending == 1:
-                    del self._extra_async_saves[req_id]
-                else:
-                    self._extra_async_saves[req_id] = extra_pending - 1
-
-        return finished_sending or None, finished_recving or None
+    def build_worker_connector_meta(
+        self, scheduler_output: SchedulerOutput,
+        model_runner_output: ModelRunnerOutput
+    ) -> Optional[MultiKVWorkerConnectorMetadata]:
+        return MultiKVWorkerConnectorMetadata(metadata=tuple(
+            c.build_worker_connector_meta(scheduler_output,
+                                          model_runner_output)
+            for c in self._connectors))
 
     # ==============================
     # Scheduler-side methods
@@ -169,9 +152,6 @@ def build_connector_meta(
         metadata = MultiKVConnectorMetadata(metadata=tuple(
             c.build_connector_meta(scheduler_output)
             for c in self._connectors))
-        if self._extra_async_saves:
-            metadata.extra_async_saves = self._extra_async_saves
-            self._extra_async_saves = {}
         return metadata
 
     def request_finished(
@@ -199,3 +179,30 @@ def request_finished(
         self._requests_to_connector.pop(request.request_id, None)
 
         return async_saves > 0, kv_txfer_params
+
+    def get_finished(
+        self, model_runner_output: ModelRunnerOutput
+    ) -> tuple[Optional[set[str]], Optional[set[str]]]:
+        finished_sending: set[str] = set()
+        finished_recving: set[str] = set()
+        for c in self._connectors:
+            sending, recving = c.get_finished(model_runner_output)
+            if not recving and not sending:
+                continue
+            # Aggregate finished recving request ids.
+            finished_recving.update(recving or ())
+            # Aggregate finished sending request ids - only include
+            # once we've drained the "extra" count (for cases where
+            # more than one connector is async-saving the same request).
+            for req_id in sending or ():
+                extra_pending = self._extra_async_saves.get(req_id)
+                if extra_pending is None:
+                    finished_sending.add(req_id)
+                    continue
+                assert extra_pending > 0
+                if extra_pending == 1:
+                    del self._extra_async_saves[req_id]
+                else:
+                    self._extra_async_saves[req_id] = extra_pending - 1
+
+        return finished_sending or None, finished_recving or None