From 4b90dc19dca47862fa23f219e383372903140e45 Mon Sep 17 00:00:00 2001
From: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>
Date: Tue, 28 Jul 2026 10:55:45 -0700
Subject: [PATCH] [nvbugs/6523785][fix] Retire CUDA graph pool handle when
 graphs are released

Backend caches one private CUDA graph pool handle per process and hands it
to every CUDAGraphRunner. When the graphs captured into that pool are reset
(KV cache estimation shutdown, executor shutdown), the pool's use_count
drops to zero, but the handle stayed cached on the class. A second LLM in
the same worker process then reused the retired handle, and the first
torch.cuda.graph(pool=...) tripped the CUDACachingAllocator assertion
'it->second->use_count > 0'.

Retire the handle whenever the graphs are released and recreate it lazily,
mirroring the handle rotation clear_piecewise_cuda_graphs() already does
for PiecewiseRunner pools.

Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>
---
 tensorrt_llm/_torch/compilation/backend.py    | 22 +++++++++++++++----
 .../_torch/pyexecutor/model_engine.py         |  8 ++++++-
 2 files changed, 25 insertions(+), 5 deletions(-)

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