[MRV2] Collect KV transfer results before the drafter runs (#62041)

Signed-off-by: Nick Hill <nickhill123@gmail.com>
Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

tests/v1/worker/test_gpu_kv_connector.py

@@ -22,7 +22,12 @@ def _make_connector(
     backend.bind_connector_metadata.side_effect = lambda _: events.append("bind")
     backend.start_load_kv.side_effect = lambda *_args, **_kwargs: events.append("start")
     backend.wait_for_save.side_effect = lambda: events.append("wait")
-    backend.get_transfer_results.return_value = KVConnectorTransferResults()
+
+    def get_transfer_results(_finished_req_ids):
+        events.append("results")
+        return KVConnectorTransferResults()
+
+    backend.get_transfer_results.side_effect = get_transfer_results
     backend.get_block_ids_with_load_errors.return_value = set()
     backend.get_kv_connector_stats.return_value = None
     backend.get_kv_connector_kv_cache_events.return_value = None
@@ -73,7 +78,7 @@ def test_load_start_phase(
     )
 
     connector.post_forward(set())
-    assert events == ["handle", "bind", "start", "wait", "clear"]
+    assert events == ["handle", "bind", "start", "wait", "results", "clear"]
 
     kwargs = connector.kv_connector.start_load_kv.call_args.kwargs
     assert kwargs["request_state_indices"] is request_indices
@@ -94,4 +99,21 @@ def test_no_forward_starts_deferred_load_once(monkeypatch: pytest.MonkeyPatch):
 
     connector.no_forward(_scheduler_output(False))
 
-    assert events == ["handle", "bind", "start", "wait", "clear"]
+    assert events == ["handle", "bind", "start", "wait", "results", "clear"]
+
+
+def test_deferred_finalize_collects_results_before_saving(
+    monkeypatch: pytest.MonkeyPatch,
+):
+    """With a drafter, the transfer results are collected right after the
+    target forward, and the saves and metadata are finalized after the
+    drafter writes its KV."""
+    events: list[str] = []
+    connector = _make_connector(monkeypatch, events)
+    connector.pre_forward(_scheduler_output(False))
+
+    connector.post_forward(set(), defer_finalize=True)
+    assert events == ["handle", "bind", "start", "results"]
+
+    connector.finalize_step()
+    assert events == ["handle", "bind", "start", "results", "wait", "clear"]

vllm/v1/worker/gpu/kv_connector.py

@@ -41,9 +41,14 @@ class KVConnector:
     def finish_forward(self) -> None:
         pass
 
-    def post_forward(self, finished_req_ids: set[str]) -> KVConnectorOutput | None:
+    def post_forward(
+        self, finished_req_ids: set[str], defer_finalize: bool = False
+    ) -> KVConnectorOutput | None:
         return None
 
+    def finalize_step(self) -> None:
+        pass
+
     def no_forward(self, scheduler_output: "SchedulerOutput") -> ModelRunnerOutput:
         return EMPTY_MODEL_RUNNER_OUTPUT
 
@@ -128,7 +133,9 @@ class ActiveKVConnector(KVConnector):
     def reset_capture_state(self) -> None:
         self.kv_connector.reset_capture_state()
 
-    def post_forward(self, finished_req_ids: set[str]) -> KVConnectorOutput | None:
+    def post_forward(
+        self, finished_req_ids: set[str], defer_finalize: bool = False
+    ) -> KVConnectorOutput | None:
         if self._disabled:
             return None
 
@@ -136,7 +143,8 @@ class ActiveKVConnector(KVConnector):
             self._start_load_kv()
 
         output = KVConnectorOutput()
-        self.kv_connector.wait_for_save()
+        if not defer_finalize:
+            self.kv_connector.wait_for_save()
         transfer_results = self.kv_connector.get_transfer_results(finished_req_ids)
         output.finished_sending = transfer_results.finished_sending or None
         output.finished_recving = transfer_results.finished_recving or None
@@ -147,9 +155,16 @@ class ActiveKVConnector(KVConnector):
         output.kv_connector_worker_meta = (
             self.kv_connector.build_connector_worker_meta()
         )
-        self.kv_connector.clear_connector_metadata()
+        if not defer_finalize:
+            self.kv_connector.clear_connector_metadata()
         return output
 
+    def finalize_step(self) -> None:
+        if self._disabled:
+            return
+        self.kv_connector.wait_for_save()
+        self.kv_connector.clear_connector_metadata()
+
     def no_forward(self, scheduler_output: "SchedulerOutput") -> ModelRunnerOutput:
         if self._disabled:
             return EMPTY_MODEL_RUNNER_OUTPUT

vllm/v1/worker/gpu/model_runner.py

@@ -2207,6 +2207,12 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             pending_aux_output=pending_aux_output,
         )
 
+        # Collect KV transfer results now; finalize saves after the drafter runs.
+        defer_kv_finalize = self.speculator is not None
+        model_runner_output.kv_connector_output = self.kv_connector.post_forward(
+            finished_req_ids, defer_finalize=defer_kv_finalize
+        )
+
         mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None
         # The encoder runner exists only on the first PP rank, so later ranks
         # have no cached embeddings to gather.
@@ -2288,9 +2294,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
                     self.req_states.draft_tokens, input_batch
                 )
 
-        # Post-step KV connector related operations.
-        kv_connector_output = self.kv_connector.post_forward(finished_req_ids)
-        model_runner_output.kv_connector_output = kv_connector_output
+        if defer_kv_finalize:
+            self.kv_connector.finalize_step()
         model_runner_output.ec_connector_output = ec_connector_output
 
         return async_output