From 8ce5d3198d00631a76e1aa02a57947b46bc7218c Mon Sep 17 00:00:00 2001 From: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com> Date: Mon, 29 Sep 2025 00:46:30 -0400 Subject: [PATCH] [P/D] NIXL Updates (#25844) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sage Moore Signed-off-by: simon-mo Signed-off-by: rentianyue-jk Signed-off-by: Russell Bryant Signed-off-by: Isotr0py Signed-off-by: Chenheli Hua Signed-off-by: mgoin Signed-off-by: Tyler Michael Smith Signed-off-by: NickLucche Signed-off-by: Roger Wang Signed-off-by: Robert Shaw Co-authored-by: Sage Moore Co-authored-by: Russell Bryant Co-authored-by: rentianyue-jk Co-authored-by: Isotr0py Co-authored-by: Chenheli Hua Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Co-authored-by: Michael Goin Co-authored-by: Tyler Michael Smith Co-authored-by: Nicolò Lucchesi Co-authored-by: Roger Wang Co-authored-by: Robert Shaw Signed-off-by: simon-mo --- .../kv_transfer/kv_connector/v1/nixl_connector.py | 7 +++++-- vllm/v1/core/sched/scheduler.py | 7 ++++++- 2 files changed, 11 insertions(+), 3 deletions(-) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py index 5af2b33f029c..c11189d7ec10 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py @@ -474,8 +474,11 @@ class NixlConnectorWorker: "backends", ["UCX"]) # Agent. non_ucx_backends = [b for b in self.nixl_backends if b != "UCX"] - config = nixl_agent_config(backends=self.nixl_backends) if len( - non_ucx_backends) > 0 and nixl_agent_config is not None else None + if nixl_agent_config is None: + config = None + else: + config = nixl_agent_config(backends=self.nixl_backends) if len( + non_ucx_backends) > 0 else nixl_agent_config(num_threads=8) self.nixl_wrapper = NixlWrapper(str(uuid.uuid4()), config) # Map of engine_id -> {rank0: agent_name0, rank1: agent_name1..}. diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 7fc4776b0261..2b2cd63c2e20 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -1288,4 +1288,9 @@ class Scheduler(SchedulerInterface): self.finished_recving_kv_req_ids.add(req_id) for req_id in (kv_connector_output.finished_sending or ()): logger.debug("Finished sending KV transfer for request %s", req_id) - self._free_blocks(self.requests[req_id]) + if req_id not in self.requests: + logger.warning( + "Got finished sending KV transfer for request %s," + "but the request is already freed.", req_id) + else: + self._free_blocks(self.requests[req_id])