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()