From 1838cb1f1a8a7b067b6ba1a621b7f9cfe9548e83 Mon Sep 17 00:00:00 2001 From: handongl Date: Wed, 29 Jul 2026 09:37:07 -0700 Subject: [PATCH 1/2] [nvbugs/6525011][fix] Keep an owning reference to captured CUDA graph outputs CUDAGraphRunner.capture() retained the captured output only through make_weak_ref, which is a non-owning pointer/shape view. The output storage was therefore free the instant capture() returned, even though the graph stayed replayable and its kernels still wrote to those addresses. Any allocation on the shared graph pool could then be placed on top of a live graph's output -- including the persistent, graph-read workspaces Buffers.get_buffer takes from get_shared_pool() -- surfacing as an illegal memory access at graph.replay() during generation CUDA graph capture on B200. Hold the output tensors for the lifetime of each graph and release them in clear() before the graphs and their pool are torn down, so the buffers drain into a still-live pool. Callers still receive the weak reference, so the graph output contract is unchanged. This restores the same ownership invariant visual_gen/cuda_graph_runner.py already keeps. The retained buffers scale with (number of captured graph shapes x output size), and are allocated during the KV-cache estimation dry run, so they are accounted for in the profiled peak rather than escaping it: peak moved 36.70 -> 36.89 GiB and estimated KV cache 128.04 -> 127.87 GiB on this config. Large-vocab, large-max-batch models will see a proportionally larger shift. Unwaives B200 TestLagunaXS::test_fp8. That test is not in l0_b200.yml (only test_nvfp4 is; test_fp8 lives in qa/llm_function_core.txt), so unwaiving does not restore it to pre-merge CI. Verified directly on a B200: TestLagunaXS::test_fp8 passes (GSM8K 84.344 vs threshold 83.947) and TestLagunaXS_2_1::test_bf16_dflash passes (87.718 vs 87.000). Signed-off-by: handongl --- tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) 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() From 859b7d2f211a4fa984943d49c44d92667ea0c019 Mon Sep 17 00:00:00 2001 From: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> Date: Thu, 30 Jul 2026 01:32:26 -0700 Subject: [PATCH 2/2] [nvbugs/6525011][chore] Remove stale waiver after fix Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> --- tests/integration/test_lists/waives.txt | 2 -- 1 file changed, 2 deletions(-) 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)