From 1592ab72b10cbd2143c22302cb51bc82868eb5c4 Mon Sep 17 00:00:00 2001 From: Dongfeng Yu Date: Sat, 31 Oct 2026 16:04:37 -0700 Subject: [PATCH] fix: release startup profiling cache before NIXL shutdown Keep NIXL alive while initial profiling KV pools are released so their VMM allocations are reclaimed before the final cache is built. Opt in only at unpublished startup profiling; retain default and terminal cleanup ordering. Update the existing mock expectation for the explicit argument. Signed-off-by: Dongfeng Yu --- tensorrt_llm/_torch/pyexecutor/_util.py | 2 +- tensorrt_llm/_torch/pyexecutor/py_executor.py | 20 +++++++++++++------ .../kv_cache/test_kv_cache_estimation.py | 2 +- 3 files changed, 16 insertions(+), 8 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 8b0668b1bba9..15c768241588 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -1899,7 +1899,7 @@ def configure_kv_cache_capacity(self, kv_stats_draft.allocated_bytes if kv_stats_draft is not None else 0) warmup.close() - py_executor.release_resources() + py_executor.release_resources(startup_profiling=True) py_executor.enable_iter_perf_stats = origin_iter_stats py_executor.set_gather_responses(False) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 43efe52eca16..53b14906c659 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -2192,16 +2192,20 @@ def _stop_mpi_progress_thread(self) -> None: self._mpi_progress_thread = None self._mpi_progress_thread_exit_hook = None - def release_resources(self) -> None: + def release_resources(self, *, startup_profiling: bool = False) -> 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 + precedes workspace release, and registrations normally precede cache release. Model engines, weights and the execution stream remain valid. + + ``startup_profiling`` is only for the unpublished initial sizing phase, + whose dummy requests do not transfer KV cache. """ if self._closed or self._resources_released: raise RuntimeError("Executor resources are already released") - if not self._release_execution_resources(): + if not self._release_execution_resources( + startup_profiling=startup_profiling): raise RuntimeError( "Cannot release resources after an executor hang") self.resource_manager = None @@ -2267,8 +2271,8 @@ def shutdown(self) -> None: def _release_execution_resources(self, *, - release_model_owned_graphs: bool = False - ) -> bool: + release_model_owned_graphs: bool = False, + startup_profiling: bool = False) -> bool: """Stop execution and release GPU resources; retain model ownership.""" if getattr(self, "_shutdown_runtime_complete", False): return self._release_execution_pools() @@ -2360,12 +2364,16 @@ def _release_execution_resources(self, # resource managers start freeing GPU-backed workspaces. if torch.cuda.is_available(): torch.cuda.synchronize() - if self.kv_cache_transceiver is not None: + if not startup_profiling and self.kv_cache_transceiver is not None: self.kv_cache_transceiver.shutdown() if self.resource_manager is not None: for manager in self.resource_manager.resource_managers.values(): if manager: manager.shutdown() + if startup_profiling and self.kv_cache_transceiver is not None: + # This unpublished phase has no KV transfers. Keep NIXL alive until + # the temporary pools are released so their VMM memory is reclaimed. + self.kv_cache_transceiver.shutdown() # Retry a later failure without repeating request publication, graph # release, or resource-manager shutdown. self._shutdown_runtime_complete = True diff --git a/tests/unittest/_torch/pyexecutor/kv_cache/test_kv_cache_estimation.py b/tests/unittest/_torch/pyexecutor/kv_cache/test_kv_cache_estimation.py index 5c06f0766e20..67220a2b7580 100644 --- a/tests/unittest/_torch/pyexecutor/kv_cache/test_kv_cache_estimation.py +++ b/tests/unittest/_torch/pyexecutor/kv_cache/test_kv_cache_estimation.py @@ -2089,7 +2089,7 @@ def test_profiling_failure_releases_resources_and_restores_response_state(): ): creator.configure_kv_cache_capacity(executor) - executor.release_resources.assert_called_once_with() + executor.release_resources.assert_called_once_with(startup_profiling=True) executor.shutdown.assert_not_called() executor.warmup_scope.assert_called_once_with() warmup_scope.__enter__.assert_called_once()