[https://nvbugs/6523785][fix] Added Backend.get_graph_pool_handle() / retire_graph_pool_handle() and…

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)

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:
         """