diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py index 3fa8edd8cdd1..07618d87a1a9 100644 --- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py +++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py @@ -139,10 +139,6 @@ 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] = {} @@ -541,7 +537,6 @@ 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() @@ -850,9 +845,6 @@ 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() @@ -908,8 +900,6 @@ 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 @@ -1229,7 +1219,6 @@ 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() @@ -1301,7 +1290,6 @@ 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 b2971766eacd..c3ad766a084a 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -153,6 +153,7 @@ 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) @@ -186,6 +187,7 @@ 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) @@ -198,9 +200,6 @@ full:B300/llmapi/test_llm_api_pytorch_moe_lora.py::test_qwen_moe_routed_expert_m full:B300/test_e2e.py::test_qwen_e2e_cpprunner_large_new_tokens[DeepSeek-R1-Distill-Qwen-1.5B-DeepSeek-R1-Distill-Qwen-1.5B] SKIP (https://nvbugs/6414760) full:DGX_B200/accuracy/test_disaggregated_serving.py::TestQwen3NextInstruct::test_auto_dtype[use_py_transceiver=True] SKIP (https://nvbugs/6501837) full:DGX_B200/accuracy/test_llm_api_pytorch.py::TestDeepSeekV4ProDSpark::test_gsm8k_dep8_megamoe_deepgemm SKIP (https://nvbugs/6506920) -full:DGX_H100/unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k "CUTLASS and W4A16_MXFP4" SKIP (https://nvbugs/6567403) -full:DGX_H100/unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k "CUTLASS and W8A16" SKIP (https://nvbugs/6567403) -full:DGX_H100/unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_multi_gpu_eplb SKIP (https://nvbugs/6567403) full:DGX_H200/accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_fp8_blockscale[disable_skip_indexer] SKIP (https://nvbugs/6476233) full:DGX_H200/accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_fp8_blockscale[latency_default] SKIP (https://nvbugs/6476233) full:GB200/accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_auto_dtype[False-True-True-True] SKIP (https://nvbugs/6525893)