[https://nvbugs/6525011][fix] Store the owning output tensor in _graph_output_refs for each graph's lifetime… - #17010
Conversation
WalkthroughCUDA graph runners now retain captured decoder and encoder outputs until graph cleanup. The integration waiver list removes obsolete skips and adds skips for additional model, hardware, performance, visual-generation, and speculative-decoding configurations. ChangesCUDA graph output lifetime and integration waiver updates
Estimated code review effort: 2 (Simple) | ~10 minutes Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py (1)
142-145: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAvoid
Anyin both strong-reference dictionaries.These maps only retain captured outputs and do not inspect their values; use
objector a precise output type alias at both sites.
tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py#L142-L145: replaceAnyin_graph_output_refs.tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py#L769-L770: apply the same non-Anyannotation to the encoder runner.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py` around lines 142 - 145, Replace the Any value annotations in both strong-reference dictionaries with object or an appropriate precise output type: update _graph_output_refs at tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py:142-145 and the encoder runner dictionary at tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py:769-770. Keep the existing key types and retention behavior unchanged.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py`:
- Around line 142-145: Replace the Any value annotations in both
strong-reference dictionaries with object or an appropriate precise output type:
update _graph_output_refs at
tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py:142-145 and the encoder
runner dictionary at
tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py:769-770. Keep the existing
key types and retention behavior unchanged.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: a4af3cf6-a7f9-44fe-9e52-9d6fb212ecfb
📒 Files selected for processing (2)
tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.pytests/integration/test_lists/waives.txt
💤 Files with no reviewable changes (1)
- tests/integration/test_lists/waives.txt
|
/bot run |
|
PR_Github #62569 [ run ] triggered by Bot. Commit: |
|
PR_Github #62569 [ run ] completed with state
|
|
/bot run |
|
PR_Github #62585 [ run ] triggered by Bot. Commit: |
|
PR_Github #62585 [ run ] completed with state
|
… 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 <handongl@nvidia.com>
Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>
92a4844 to
a34f87d
Compare
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tests/integration/test_lists/waives.txt (1)
37-37: 📐 Maintainability & Code Quality | 🔵 Trivial | 🏗️ Heavy liftConfirm every waiver is backed by a failed test before merging.
Test coverage summary:
- Modified list:
tests/integration/test_lists/waives.txt.- Added 17
SKIPentries across accuracy, visual-generation, performance, H20 LoRA, and speculative-decoding tests.- Removed 8 existing
SKIPentries.- No test functions changed, and no
test-db/orqa/list files were modified.- Coverage verdict: needs follow-up because neither
cbts_touchmap.sqlitenor a CBTS coverage report is provided.Given the two failed L0 pipelines, verify that each new waiver corresponds to an observed failure; otherwise these broad additions may mask regressions. Also split entries unrelated to the CUDA graph fix into a separate waiver update if they are not caused by this PR.
Also applies to: 58-58, 88-96, 139-141, 150-152, 174-175, 194-194, 217-217, 235-235, 252-252, 337-339, 367-368
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/integration/test_lists/waives.txt` at line 37, Verify every newly added SKIP entry in tests/integration/test_lists/waives.txt against a corresponding failed L0 test, and remove any waiver without observed failure evidence. Separate waivers unrelated to the CUDA graph fix into an independent update, preserving only entries caused by this PR.Source: Path instructions
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@tests/integration/test_lists/waives.txt`:
- Line 37: Verify every newly added SKIP entry in
tests/integration/test_lists/waives.txt against a corresponding failed L0 test,
and remove any waiver without observed failure evidence. Separate waivers
unrelated to the CUDA graph fix into an independent update, preserving only
entries caused by this PR.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 9cae8f57-7dda-4187-a4e0-1be57636a832
📒 Files selected for processing (2)
tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.pytests/integration/test_lists/waives.txt
fredricz-20070104
left a comment
fredricz-20070104
left a comment
There was a problem hiding this comment.
Review summary - Approve
Reviewed the full diff; no blocking or major issues found.
Minor, non-blocking notes:
tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py: Any annotation on strong-ref dictstensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py: No per-key cleanup of _graph_output_refs on graph eviction/recapture
Automated review by NVCortex Lite, run by @fredricz-20070104.
|
/bot run |
|
PR_Github #62763 [ run ] triggered by Bot. Commit: |
|
PR_Github #62763 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #62772 [ run ] triggered by Bot. Commit: |
|
PR_Github #62772 [ run ] completed with state
|
|
/bot run |
|
PR_Github #62809 [ run ] triggered by Bot. Commit: |
|
PR_Github #62809 [ run ] completed with state |
7 participants
Summary
Test plan
Links
Dev Engineer Review
CUDAGraphRunnerandEncoderCUDAGraphRunnerby introducing_graph_output_refs: Dict[graph_key, Any]to hold strong references to capture-time outputs for each graph’s replayable lifetime.graph_outputscontinues to store weak refs), while ensuring the underlying storage cannot be returned to/reused by the shared CUDA graph memory pool prematurely.clear()to clear_graph_output_refsbefore resetting captured graphs and releasing the memory pool, preventing allocator/use-count-related failures.QA Engineer Review
tests/integration/test_lists/waives.txt:cute_dsl_nvfp4_4gpusvariants acrossfull:*targets.full:B200/accuracy/test_llm_api_pytorch.py::TestLagunaXS::test_fp8(nvbugs6525011).full:DGX_H100/unittest/scaffolding(nvbugs6529626) andfull:L40S/accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B::test_fp8[latency-torch_compile=True](nvbugs6276841).TestNemotronV3Ultra/TestQwen3NextInstructandTestQwen3_5_397B_A17Bnvfp4/bf16 variants,examples/visual_genWAN22 LPIPS (single- and multi-GPU),H100TestQwen3_30B_A3B_Instructsoftmax-skip,H20MoE multi-Loracudagraph,perf/test_perf_sanity.pyfp4 dep8/mmtp/tp4 scenarios (Qwen3_5_397B),unittest/_torch/speculative/hw_agnostic/test_dflash.pyQwen3_5_4b waivers (both parameterizations).test-db/orqa/files were modified.