diff --git a/tensorrt_llm/_torch/compilation/backend.py b/tensorrt_llm/_torch/compilation/backend.py index 64e078264a3c..169de9c67759 100644 --- a/tensorrt_llm/_torch/compilation/backend.py +++ b/tensorrt_llm/_torch/compilation/backend.py @@ -78,9 +78,6 @@ def __init__( self.events = Backend.Events() inductor_config.enable_auto_functionalized_v2 = False - if Backend._graph_pool_handle is None: - Backend._graph_pool_handle = torch.cuda.graph_pool_handle() - self.match_count = [] self.match_count_by_pass = OrderedDict() @@ -112,6 +109,23 @@ def build_custom_passes(cls, enable_userbuffers, mapping: Mapping): register_add_norm(custom_passes[-1]) return custom_passes + @classmethod + def get_graph_pool_handle(cls) -> tuple[int, int]: + """Return the pool handle shared by every graph runner in this process.""" + if cls._graph_pool_handle is None: + cls._graph_pool_handle = torch.cuda.graph_pool_handle() + return cls._graph_pool_handle + + @classmethod + def retire_graph_pool_handle(cls) -> None: + """Drop the cached handle once its graphs have been reset. + + A private pool cannot outlive its graphs: CUDACachingAllocator asserts + on any attempt to incref a pool whose use_count already dropped to + zero, so the next caller has to allocate a fresh handle. + """ + cls._graph_pool_handle = None + def bypass_optimization(self): self.no_optimization = True @@ -179,7 +193,7 @@ def optimize( self.enable_inductor, self.input_num_tokens, self.capture_num_tokens, - self._graph_pool_handle, + self.get_graph_pool_handle(), self.num_streams, ) self._piecewise_runners.update(runners) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index b498154d4781..66427c6a3e65 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -650,7 +650,8 @@ def __init__( self.encoder_attn_metadata = None self.spec_metadata = None self.iter_states = {} - self._cuda_graph_mem_pool = self._torch_compile_backend._graph_pool_handle if self._torch_compile_enabled else None + self._cuda_graph_mem_pool = (Backend.get_graph_pool_handle() + if self._torch_compile_enabled else None) self._cuda_graph_padding_enabled = cuda_graph_padding_enabled @@ -2746,6 +2747,11 @@ def _release_cuda_graphs(self): if hasattr(self, 'encoder_cuda_graph_runner' ) and self.encoder_cuda_graph_runner is not None: self.encoder_cuda_graph_runner.clear() + # The graphs reset above were captured into Backend's process-wide pool, + # so a later engine in this process (e.g. a second LLM sharing the + # worker) must not reuse that now-dead handle. + Backend.retire_graph_pool_handle() + self._cuda_graph_mem_pool = None def get_max_num_sequences(self) -> int: """