[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>

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