[https://nvbugs/7023112][fix] Release profiling KV pools before NIXL deregistration

tensorrt_llm/_torch/disaggregation/kv_cache_transceiver.py

@@ -484,8 +484,21 @@ def get_status_dump(self) -> str:
         """Return a human-readable dump of transceiver state for debugging hangs."""
         return ""
 
+    def prepare_shutdown(self) -> None:
+        """Stop transfers before cache release; keep registrations when supported.
+
+        Backends without separate drain and deregistration phases shut down
+        completely here. Raise if transfers may still access the cache; the caller
+        retains the cache and may retry.
+        """
+        self.shutdown()
+
     def shutdown(self) -> None:
-        """Shut down the transceiver and release registered resources."""
+        """Idempotently stop the transceiver and release registered resources.
+
+        May follow prepare_shutdown() after cache release, so completed
+        preparation must prevent further cache-memory access during shutdown.
+        """
 
 
 class BindKvCacheTransceiver(KvCacheTransceiver):

tensorrt_llm/_torch/disaggregation/native/transfer.py

@@ -4540,6 +4540,7 @@ class TransferWorkerConfig:
 
 class TransferWorker:
     def __init__(self, config: TransferWorkerConfig):
+        self._shutdown_prepared = False
         self._config = config
         self._retirement_watchdog = None
         if config.quiescence_fatal_callback is not None:
@@ -4728,8 +4729,9 @@ def page_table(self):
         assert self._rank_info is not None
         return self._rank_info.page_table
 
-    def shutdown(self):
-        if getattr(self, "_shutdown", False):
+    def prepare_shutdown(self) -> None:
+        """Stop listeners and transfers, refusing active owners; retain registrations."""
+        if getattr(self, "_shutdown_prepared", False):
             return
         self.request_shutdown()
         watchdog = getattr(self, "_retirement_watchdog", None)
@@ -4749,9 +4751,17 @@ def shutdown(self):
             receiver.shutdown()
         if rank_info_server is not None:
             rank_info_server.shutdown()
+        self._shutdown_prepared = True
+
+    def shutdown(self) -> None:
+        """Finish shutdown by releasing registrations and the transfer agent."""
+        if getattr(self, "_shutdown", False):
+            return
+        self.prepare_shutdown()
         self._shutdown = True
-        # Deregister NIXL memory before agent.shutdown so pinned GPU memory is released
-        # (e.g. when the KV cache manager is recreated after profiling).
+        # For prompt unpinning, callers unmap KV pools before deregistration.
+        # UCX shared GDRCopy caching otherwise retains the pinned pages until
+        # a later registration. prepare_shutdown() retains registrations.
         agent = getattr(self, "_agent", None)
         if agent is not None:
             try:

tensorrt_llm/_torch/disaggregation/transceiver.py

@@ -689,23 +689,39 @@ def summarize(
             f"waiting_for_peer_info={len(self._wait_reqs)}"
         )
 
+    def _get_shutdown_lock(self) -> threading.Lock:
+        """Return the shutdown lock, including after partial initialization."""
+        lock = getattr(self, "_shutdown_lock", None)
+        if lock is None:
+            lock = threading.Lock()
+            self._shutdown_lock = lock
+        return lock
+
+    def prepare_shutdown(self) -> None:
+        """Retire sessions and stop transfers without deregistering the KV pools."""
+        with self._get_shutdown_lock():
+            if getattr(self, "_shutdown_complete", False):
+                return
+            self._close_sessions_for_shutdown()
+            self._transfer_worker.prepare_shutdown()
+
     def shutdown(self) -> None:
-        shutdown_lock = getattr(self, "_shutdown_lock", None)
-        if shutdown_lock is None:
-            shutdown_lock = threading.Lock()
-            self._shutdown_lock = shutdown_lock
-        with shutdown_lock:
+        with self._get_shutdown_lock():
             self._shutdown_once()
 
     def _shutdown_once(self) -> None:
         if getattr(self, "_shutdown_complete", False):
             return
+        self._close_sessions_for_shutdown()
+        self._transfer_worker.shutdown()
+        # Record completion only after every owner and worker has shut down.
+        # A refusal leaves shutdown retryable after the physical operation drains.
+        self._shutdown_complete = True
+
+    def _close_sessions_for_shutdown(self) -> None:
         # Preflight may refuse closure while DMA remains active. Start the
-        # bounded drain before that refusal, not only in worker.shutdown().
+        # bounded drain before that refusal, not only in worker.prepare_shutdown().
         self._transfer_worker.request_shutdown()
-        # This flag records completed teardown, not an attempted shutdown. If
-        # an active owner refuses closure, leave it false so shutdown can be
-        # retried after the physical operation drains.
         # Close receive owners before touching send-side or worker state. The
         # separate drain snapshot would not be sufficient: late evidence can make an
         # owner unretirable before close() acquires its lock.
@@ -721,8 +737,6 @@ def _shutdown_once(self) -> None:
         self._send_reqs.clear()
         self._recv_sessions.clear()
         self._recv_reqs.clear()
-        self._transfer_worker.shutdown()
-        self._shutdown_complete = True
 
     def __enter__(self):
         return self

tensorrt_llm/_torch/pyexecutor/py_executor.py

@@ -2479,8 +2479,9 @@ def release_resources(self) -> None:
         """Stop a profiling phase and detach resources before cache replacement.
 
         Call collectively after consuming profiling responses. Graph destruction
-        precedes workspace release, and transfer registrations precede cache
-        release. Model engines, weights and the execution stream remain valid.
+        precedes workspace release. Transfers drain before cache release, while
+        deferred registrations are released after the pools are unmapped. Model
+        engines, weights and the execution stream remain valid.
         """
         if self._closed or self._resources_released:
             raise RuntimeError("Executor resources are already released")
@@ -2644,11 +2645,14 @@ def _release_execution_resources(self,
         if torch.cuda.is_available():
             torch.cuda.synchronize()
         if self.kv_cache_transceiver is not None:
-            self.kv_cache_transceiver.shutdown()
+            self.kv_cache_transceiver.prepare_shutdown()
         if self.resource_manager is not None:
             for manager in self.resource_manager.resource_managers.values():
                 if manager:
                     manager.shutdown()
+        if self.kv_cache_transceiver is not None:
+            # Backends that defer deregistration release it after the pools.
+            self.kv_cache_transceiver.shutdown()
         # Retry a later failure without repeating request publication, graph
         # release, or resource-manager shutdown.
         self._shutdown_runtime_complete = True

tests/unittest/_torch/pyexecutor/test_executor_resource_rebind.py

@@ -44,7 +44,10 @@ def release_engine_resources(*, release_model_owned_graphs):
 
     engine.release_executor_resources = Mock(side_effect=release_engine_resources)
     cache = SimpleNamespace(shutdown=lambda: events.append("cache"))
-    transceiver = SimpleNamespace(shutdown=lambda: events.append("unregister"))
+    transceiver = SimpleNamespace(
+        prepare_shutdown=lambda: events.append("drain"),
+        shutdown=lambda: events.append("unregister"),
+    )
     ex = object.__new__(PyExecutor)
     ex._closed = False
     ex.enable_attention_dp = False
@@ -138,9 +141,10 @@ def shutdown_slots():
         "graphs",
         "draft_graphs",
         "synchronize",
-        "unregister",
+        "drain",
         "slots",
         "cache",
+        "unregister",
     ]
     assert ex.draft_model_engine is draft_engine
     assert draft_engine.attn_metadata is None
@@ -175,7 +179,7 @@ def test_release_exits_a_bound_dwdp_manager_after_the_caches(
 
     ex.release_resources()
 
-    assert events[-2:] == ["cache", field]
+    assert events[-3:] == ["cache", "unregister", field]
     assert getattr(ex, field) is None
     ex.shutdown()
     manager.__exit__.assert_called_once_with(None, None, None)
@@ -186,7 +190,7 @@ def test_release_before_worker_start_does_not_wait_for_shutdown(executor):
     ex.worker_started = False
     ex.shutdown_event.clear()
     ex.release_resources()
-    assert events == ["graphs", "synchronize", "unregister", "cache"]
+    assert events == ["graphs", "synchronize", "drain", "cache", "unregister"]
     ex.worker_thread.join.assert_not_called()
 
 
@@ -422,13 +426,15 @@ def test_hung_worker_does_not_release_live_resources(executor):
     assert slots.slot_manager.get_slot(17) is not None
 
 
-def test_transfer_shutdown_failure_keeps_cache_owned(executor):
+def test_transfer_drain_failure_keeps_cache_owned(executor):
     ex, events = executor
     manager = ex.resource_manager
     slots = manager.resource_managers[ResourceManagerType.SEQ_SLOT_MANAGER]
-    ex.kv_cache_transceiver.shutdown = Mock(side_effect=RuntimeError("DMA still active"))
+    ex.kv_cache_transceiver.prepare_shutdown = Mock(side_effect=RuntimeError("DMA still active"))
+
     with pytest.raises(RuntimeError, match="DMA still active"):
         ex.release_resources()
+    assert "unregister" not in events
     assert "cache" not in events
     assert ex.resource_manager is manager
     assert slots.slot_manager.get_slot(17) is not None
@@ -456,3 +462,11 @@ def test_release_stops_pipeline_progress_before_gpu_teardown(executor):
     assert ex._pp_mpi_progress_quiesced.is_set()
     assert ex._mpi_progress_thread is None
     assert ex._mpi_progress_thread_exit_hook is None
+
+
+@pytest.mark.parametrize("terminal", [False, True])
+def test_both_release_paths_unmap_before_deregistration(executor, terminal):
+    """Profiling and terminal release must use the same two-phase order."""
+    ex, events = executor
+    assert ex._release_execution_resources(release_model_owned_graphs=terminal)
+    assert events.index("drain") < events.index("cache") < events.index("unregister")

tests/unittest/disaggregated/test_profiling_cache_release.py

@@ -0,0 +1,143 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Transfer drain must precede pool release, and deregistration must follow it."""
+
+import threading
+from collections.abc import Iterator
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import pytest
+
+from tensorrt_llm._torch.disaggregation.kv_cache_transceiver import KvCacheTransceiver
+from tensorrt_llm._torch.disaggregation.native.transfer import Sender, TransferWorker
+from tensorrt_llm._torch.disaggregation.transceiver import KvCacheTransceiverV2
+
+pytestmark = pytest.mark.cpu_only
+
+
+@pytest.fixture
+def worker() -> Iterator[TransferWorker]:
+    worker = object.__new__(TransferWorker)
+    worker._retirement_watchdog = Mock()
+    worker._sender = Mock()
+    worker._receiver = Mock()
+    worker._rank_info_server = Mock()
+    worker._agent = Mock(bounce_enabled=False)
+    worker._registered_mem = [object()]
+    yield worker
+    worker._retirement_watchdog.require_retired.side_effect = None
+    worker.shutdown()
+
+
+def test_worker_drain_retains_registration_until_shutdown(worker: TransferWorker) -> None:
+    agent = worker._agent
+    registered = worker._registered_mem.copy()
+    worker.prepare_shutdown()
+    worker.prepare_shutdown()
+
+    worker._sender.shutdown.assert_called_once_with()
+    worker._receiver.shutdown.assert_called_once_with()
+    worker._rank_info_server.shutdown.assert_called_once_with()
+    worker._retirement_watchdog.require_retired.assert_called_once_with()
+    assert worker._agent is agent
+    assert worker._registered_mem == registered
+    agent.deregister_memory.assert_not_called()
+    agent.shutdown.assert_not_called()
+
+    worker.shutdown()
+    worker.shutdown()
+    agent.deregister_memory.assert_called_once_with(registered[0])
+    agent.shutdown.assert_called_once_with()
+    worker._sender.shutdown.assert_called_once_with()
+    assert worker._registered_mem == []
+
+
+def test_worker_drain_refusal_preserves_registration_and_can_retry(worker: TransferWorker) -> None:
+    agent = worker._agent
+    registered = worker._registered_mem.copy()
+    worker._retirement_watchdog.require_retired.side_effect = RuntimeError("DMA still active")
+    with pytest.raises(RuntimeError, match="DMA still active"):
+        worker.prepare_shutdown()
+    assert not getattr(worker, "_shutdown_prepared", False)
+    assert worker._registered_mem == registered
+    worker._sender.shutdown.assert_not_called()
+    agent.deregister_memory.assert_not_called()
+    agent.shutdown.assert_not_called()
+
+    worker._retirement_watchdog.require_retired.side_effect = None
+    worker.prepare_shutdown()
+    assert worker._shutdown_prepared
+    agent.deregister_memory.assert_not_called()
+
+
+def test_transceiver_drain_refusal_preserves_sessions_and_registration(
+    worker: TransferWorker,
+) -> None:
+    session = SimpleNamespace(
+        _enforce_physical_ownership=True,
+        resources_drained=Mock(return_value=False),
+        close=Mock(return_value=True),
+    )
+    transceiver = object.__new__(KvCacheTransceiverV2)
+    transceiver._send_sessions = {}
+    transceiver._send_reqs = {}
+    transceiver._recv_sessions = {1: session}
+    transceiver._recv_reqs = {1: object()}
+    transceiver._transfer_worker = worker
+    agent = worker._agent
+
+    with pytest.raises(RuntimeError, match="physical resources remain active"):
+        transceiver.prepare_shutdown()
+    assert transceiver._recv_sessions[1] is session
+    session.close.assert_not_called()
+    worker._receiver.shutdown.assert_not_called()
+    agent.deregister_memory.assert_not_called()
+
+    session.resources_drained.return_value = True
+    transceiver.prepare_shutdown()
+    transceiver.prepare_shutdown()
+    assert transceiver._recv_sessions == {}
+    session.close.assert_called_once_with()
+    worker._receiver.shutdown.assert_called_once_with()
+    agent.deregister_memory.assert_not_called()
+    assert not getattr(transceiver, "_shutdown_complete", False)
+
+    transceiver.shutdown()
+    assert transceiver._shutdown_complete
+    agent.deregister_memory.assert_called_once()
+    agent.shutdown.assert_called_once_with()
+
+
+def test_legacy_transceiver_prepare_uses_full_shutdown() -> None:
+    transceiver = SimpleNamespace(shutdown=Mock(side_effect=RuntimeError("DMA still active")))
+    with pytest.raises(RuntimeError, match="DMA still active"):
+        KvCacheTransceiver.prepare_shutdown(transceiver)
+    transceiver.shutdown.assert_called_once_with()
+
+
+def test_receiver_refusal_retries_after_real_sender_shutdown(worker: TransferWorker) -> None:
+    """Retry a partial stop without deregistering memory or stopping sender twice."""
+    sender = object.__new__(Sender)
+    sender._sessions_lock = threading.Lock()
+    sender._shutdown = False
+    sender._sessions = {}
+    sender._pending_settlements = []
+    sender._messenger = Mock()
+    sender._send_task_queues = []
+    sender._worker_threads = []
+    sender._loaded_remote_agents_lock = threading.Lock()
+    sender._loaded_remote_agents = set()
+    sender._dealers = {}
+    worker._sender = sender
+    agent = worker._agent
+    worker._receiver.shutdown.side_effect = [RuntimeError("receiver busy"), None]
+    with pytest.raises(RuntimeError, match="receiver busy"):
+        worker.prepare_shutdown()
+    assert sender._shutdown
+    assert not getattr(worker, "_shutdown_prepared", False)
+    agent.deregister_memory.assert_not_called()
+    worker.prepare_shutdown()
+    assert worker._shutdown_prepared
+    sender._messenger.stop.assert_called_once_with()
+    agent.deregister_memory.assert_not_called()