diff --git a/tests/v1/worker/test_gpu_kv_connector.py b/tests/v1/worker/test_gpu_kv_connector.py index 78fe3c185b..dc46c22c7c 100644 --- a/tests/v1/worker/test_gpu_kv_connector.py +++ b/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"] diff --git a/vllm/v1/worker/gpu/kv_connector.py b/vllm/v1/worker/gpu/kv_connector.py index e9b71dae3e..9edcc2831a 100644 --- a/vllm/v1/worker/gpu/kv_connector.py +++ b/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 diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 1b75d73493..b8ad557b68 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/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