From e629145de22025b9447817f0b14950b6b815de09 Mon Sep 17 00:00:00 2001
From: Ovidiu Mara <ovidium@nvidia.com>
Date: Tue, 6 Oct 2026 19:06:06 +0200
Subject: [PATCH] [Bugfix][NIXL] Release unreachable prefill peers on failed
 heartbeat

In pull mode, the decoder does not release a remote prefill's NIXL state
on failed heartbeat (which is detectable from UCX disconnection error,
but current code does not check it).

When a prefill node crashes, its rkeys stay loaded on the decode nodes
it recently communicated with, which results in MNNVL stale mappings.
When the prefill node attempts to restart, it is unable to remap memory.

The current implementation of heartbeat could detect this, but they are only
sent by decoders when they need to extend KV leases for pending requests.

This patch adds periodic sending of heartbeats, with detection and cleanup
of disconnected peers.

Signed-off-by: Ovidiu Mara <ovidium@nvidia.com>
---
 .../kv_connector/v1/nixl/base_worker.py       | 70 ++++++++++++++++++-
 1 file changed, 69 insertions(+), 1 deletion(-)

diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py
index fa29d41a825e..02b9de3ad255 100644
--- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py
@@ -162,6 +162,10 @@ def _uses_dense_virtual_transfer_pages(
     )
 
 
+def _is_remote_disconnect(error: Exception) -> bool:
+    return type(error).__name__ == "nixlRemoteDisconnectError"
+
+
 class NixlBaseConnectorWorker:
     """Base implementation of Worker side methods shared by pull and push."""
 
@@ -884,6 +888,14 @@ def __init__(
         )
         # Map of pull peer (host, port) -> engine_id currently serving it.
         self._engine_by_address: dict[tuple[str, int], EngineId] = {}
+        # Peers whose NIXL endpoints failed UCX keepalive.
+        self._unreachable_remote_engines: set[EngineId] = set()
+        self._engine_heartbeat_interval: float = float(
+            vllm_config.kv_transfer_config.get_from_extra_config(
+                "engine_heartbeat_interval", 5.0
+            )
+        )
+        self._next_engine_heartbeat: float = 0.0
 
         self.model_config = vllm_config.model_config
 
@@ -3041,6 +3053,7 @@ def get_transfer_results(self) -> KVConnectorTransferResults:
         self._reap_expired_send_leases(done_sending)
 
         self._cleanup_remote_engines()
+        self._release_unreachable_remote_engines()
         return KVConnectorTransferResults(
             finished_sending=done_sending,
             finished_recving=done_recving,
@@ -3259,7 +3272,10 @@ def _send_heartbeats(self, metadata: NixlConnectorMetadata) -> None:
             for agent_name in self._remote_agents[engine_id].values():
                 try:
                     self.nixl_wrapper.send_notif(agent_name, notif_msg=hb_msg)
-                except Exception:
+                except Exception as e:
+                    if _is_remote_disconnect(e):
+                        self._mark_remote_engine_unreachable(engine_id, agent_name)
+                        break
                     logger.debug(
                         "Failed to send heartbeat to engine %s",
                         engine_id,
@@ -3735,6 +3751,58 @@ def _mark_remote_engine_invalid(self, engine_id: EngineId, now: float) -> None:
             f" and a {backoff:.0f}s backoff" if backoff else "",
         )
 
+    def _release_unreachable_remote_engines(self) -> None:
+        """Release peers that failed heartbeat, once no reads are in flight."""
+        if self._TRANSFER_MODE != "pull":
+            return
+        self._heartbeat_remote_engines()
+        with self._handshake_lock:
+            if self._handshake_futures:
+                return
+            unreachable = self._unreachable_remote_engines & self._remote_agents.keys()
+        self._unreachable_remote_engines = set(unreachable)
+        for engine_id in unreachable - self._engines_with_inflight_transfers():
+            self._cleanup_remote_engine(engine_id, log_eviction=False)
+            self._unreachable_remote_engines.discard(engine_id)
+            logger.info(
+                "Released NIXL state for unreachable remote engine %s.", engine_id
+            )
+
+    def _heartbeat_remote_engines(self) -> None:
+        """Send an empty heartbeat to every peer; it fails once UCX keepalive
+        has failed the peer's endpoints."""
+        if self._engine_heartbeat_interval <= 0:
+            return
+        now = time.perf_counter()
+        if now < self._next_engine_heartbeat:
+            return
+        self._next_engine_heartbeat = now + self._engine_heartbeat_interval
+        with self._handshake_lock:
+            engines = {
+                engine_id: list(agents.values())
+                for engine_id, agents in self._remote_agents.items()
+            }
+        for engine_id, agent_names in engines.items():
+            for agent_name in agent_names:
+                try:
+                    self.nixl_wrapper.send_notif(agent_name, notif_msg=b"HB:")
+                except Exception as e:
+                    if _is_remote_disconnect(e):
+                        self._mark_remote_engine_unreachable(engine_id, agent_name)
+                        break
+
+    def _mark_remote_engine_unreachable(
+        self, engine_id: EngineId, agent_name: str
+    ) -> None:
+        if engine_id in self._unreachable_remote_engines:
+            return
+        logger.warning(
+            "Remote engine %s (agent %s) failed heartbeat; releasing its NIXL state.",
+            engine_id,
+            agent_name,
+        )
+        self._unreachable_remote_engines.add(engine_id)
+
     def _cleanup_remote_engine(
         self, engine_id: EngineId, *, log_eviction: bool = True
     ) -> None: