[https://nvbugs/6525011][fix] Store the owning output tensor in _graph_output_refs for each graph's lifetime…

tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py

@@ -139,6 +139,10 @@ def __init__(self, config: CUDAGraphRunnerConfig):
         self.graphs: Dict[KeyType, torch.cuda.CUDAGraph] = {}
         self.graph_outputs: Dict[KeyType,
                                  Callable[[], Optional[torch.Tensor]]] = {}
+        # graph_outputs holds only non-owning weak refs, so these strong refs are
+        # what stop the capture-time output storage from returning to the shared
+        # graph pool and being reused while the graph is still replayable.
+        self._graph_output_refs: Dict[KeyType, Any] = {}
         self.graph_metadata: Dict[KeyType, Dict[str, Any]] = {}
         self.memory_pool = config.cuda_graph_mem_pool
         self.padding_dummy_requests: Dict[int, LlmRequest] = {}
@@ -537,6 +541,7 @@ def _setup_spec_decoding_and_forward(key: KeyType, forward_fn: Callable,
                                                saved_kv_lens_cuda)
 
         self.graphs[key] = graph
+        self._graph_output_refs[key] = output
         graph_output = make_weak_ref(output)
         self.graph_outputs[key] = graph_output
         self.memory_pool = graph.pool()
@@ -763,6 +768,9 @@ def pad_batch(self,
 
     def clear(self):
         """Releases all captured graphs and the associated memory pool."""
+        # Drop the output buffers while the pool that backs them is still alive;
+        # freeing them after graph.reset() trips the allocator's use_count check.
+        self._graph_output_refs.clear()
         for graph in self.graphs.values():
             graph.reset()
         self.graphs.clear()
@@ -818,6 +826,8 @@ def __init__(self, config: EncoderCUDAGraphRunnerConfig):
         self.graphs: Dict[EncoderKeyType, torch.cuda.CUDAGraph] = {}
         self.graph_outputs: Dict[EncoderKeyType, Callable[[],
                                                           Optional[Any]]] = {}
+        # See CUDAGraphRunner._graph_output_refs.
+        self._graph_output_refs: Dict[EncoderKeyType, Any] = {}
         self.graph_metadata: Dict[EncoderKeyType, Dict[str, Any]] = {}
         self.memory_pool = config.cuda_graph_mem_pool
 
@@ -1137,6 +1147,7 @@ def capture(
                 "Encoder CUDA graph does not support nested tensor outputs. "
                 "Disable encoder CUDA graphs for models with ragged outputs.")
         self.graphs[key] = graph
+        self._graph_output_refs[key] = output
         graph_output = make_weak_ref(output)
         self.graph_outputs[key] = graph_output
         self.memory_pool = graph.pool()
@@ -1208,6 +1219,7 @@ def get_graph_pool(self):
         return self.memory_pool
 
     def clear(self):
+        self._graph_output_refs.clear()
         for graph in self.graphs.values():
             graph.reset()
         self.graphs.clear()

tests/integration/test_lists/waives.txt

@@ -153,7 +153,6 @@ full:B200/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4g
 full:B200/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_cute_dsl_nvfp4_4gpus[tp4-fp8kv=False-attention_dp=False-cuda_graph=False-overlap_scheduler=False-torch_compile=False] SKIP (https://nvbugs/6526186)
 full:B200/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_cute_dsl_nvfp4_4gpus[tp4-fp8kv=False-attention_dp=False-cuda_graph=False-overlap_scheduler=False-torch_compile=True] SKIP (https://nvbugs/6474888)
 full:B200/accuracy/test_llm_api_pytorch.py::TestDeepSeekV4FlashBase::test_fp8_4gpus_static_eplb[moe_backend=WIDEEP] SKIP (https://nvbugs/6546609)
-full:B200/accuracy/test_llm_api_pytorch.py::TestLagunaXS::test_fp8 SKIP (https://nvbugs/6525011)
 full:B200/accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8[fp8kv=True-attn_backend=TRTLLM-torch_compile=True] SKIP (https://nvbugs/6473161)
 full:B200/accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8_4gpus[tp4-fp8kv=True-attn_backend=TRTLLM-torch_compile=True] SKIP (https://nvbugs/6473161)
 full:B200/accuracy/test_llm_api_pytorch.py::TestLlama3_3_70BInstruct::test_nvfp4_tp4[torch_compile=True] SKIP (https://nvbugs/6525010)
@@ -189,7 +188,6 @@ full:B300/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16[mt
 full:B300/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=0-ep4-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True-low_precision_combine=False-torch_compile=True] SKIP (https://nvbugs/6474888)
 full:B300/accuracy/test_llm_api_pytorch.py::TestDeepSeekV4FlashBase::test_fp8_4gpus_static_eplb[moe_backend=WIDEEP] SKIP (https://nvbugs/6546609)
 full:B300/accuracy/test_llm_api_pytorch.py::TestGemma3_1BInstruct::test_fp8_prequantized[torch_compile=True] SKIP (https://nvbugs/6475346)
-full:B300/accuracy/test_llm_api_pytorch.py::TestLagunaXS::test_fp8 SKIP (https://nvbugs/6525011)
 full:B300/accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_auto_dtype[tp_size=8-ep_size=8] SKIP (https://nvbugs/6445375)
 full:B300/accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_mxfp8[use_msa=False] SKIP (https://nvbugs/6424188)
 full:B300/accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4[use_msa=False] SKIP (https://nvbugs/6445375)