diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py index 05343d0deddd..76ab2975fdb1 100644 --- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py +++ b/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() diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 647729a6c0da..4ff9b2af8f25 100644 --- a/tests/integration/test_lists/waives.txt +++ b/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)