Skip to content

[https://nvbugs/6567403][fix] Revert #17010 to unblock DGX_H100 PyTorch-Others-1 stage timeout #17352

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
QiJune merged 1 commit into from
Aug 6, 2026
+2 −15
Merged

[https://nvbugs/6567403][fix] Revert #17010 to unblock DGX_H100 PyTorch-Others-1 stage timeout #17352

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 0 additions & 12 deletions tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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] = {}
Expand Down Expand Up @@ -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)

@BowenFu BowenFu Aug 6, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is what the revert restores: graph_outputs[key] goes back to weak-ref-only, with no owning copy retained.

The pinning that #17010 added was one output buffer per graph key held for the runner's lifetime — with a batch-size × draft-len key space that is not a small number of buffers, which is why a tp=2 mpirun CUDA-graph workload was the first thing to fall over.

If the hazard #17010 described is real for this runner (and not just for the VisualGen one it appears to be modelled on), the fix wants to be capture() returning the owning output rather than a lifetime-long strong-ref dict — otherwise this reopens nvbugs/6525011 with the memory cost simply traded back.

self.graph_outputs[key] = graph_output
self.memory_pool = graph.pool()
Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -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

Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -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()
Expand Down
5 changes: 2 additions & 3 deletions tests/integration/test_lists/waives.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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)
Expand All @@ -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)
Expand Down
Loading