[https://nvbugs/6525011][fix] Store the owning output tensor in _graph_output_refs for each graph's lifetime… by trtllm-agent · Pull Request #17010 · NVIDIA/TensorRT-LLM

trtllm-agent · 2026-07-29T17:10:12Z

Summary

Root cause: capture() kept the graph output only as a non-owning weak ref, so its storage was free while the graph stayed replayable, letting a shared-graph-pool allocation land on top of a live graph's output.
Fix: Store the owning output tensor in _graph_output_refs for each graph's lifetime in both runners, cleared in clear() before graph.reset() so buffers drain into a still-live pool; callers still get the weak ref.
Automated fix generated by repair-bot

Test plan

Verify fix on the same GPU type as the original failure
Check for regressions in related tests

Links

Bug: https://nvbugs/6525011

Dev Engineer Review

Added strong output references in both CUDA graph runners.
Preserved weak references for callers.
Cleared strong references before graph reset and memory-pool teardown.
Removed two full:* TestLagunaXS::test_fp8 SKIP waivers for B200 and B300.
Both removed waivers referenced nvbug 6525011.
The changes are scoped to the CUDA graph lifetime fix and waiver cleanup.

QA Engineer Review

Updated tests/integration/test_lists/waives.txt.
Removed two full:* TestLagunaXS::test_fp8 SKIP waivers.
No test-db/ or qa/ files were modified.
Verdict: needs follow-up because CBTS coverage data is unavailable.

coderabbitai · 2026-07-29T17:13:08Z

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: a7bffcaf-1b30-4958-83cd-5a151e617c03

📥 Commits

Reviewing files that changed from the base of the PR and between 4ad991f and e839f1d.

📒 Files selected for processing (2)

tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py
tests/integration/test_lists/waives.txt

💤 Files with no reviewable changes (1)

tests/integration/test_lists/waives.txt

🚧 Files skipped from review as they are similar to previous changes (1)

tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py

Walkthrough

CUDA graph runners now retain captured decoder and encoder outputs until cleanup. The integration waiver list removes the B200 and B300 TestLagunaXS::test_fp8 skip entries associated with nvbug 6525011.

Changes

CUDA graph output lifetime and integration waiver updates

Layer / File(s)	Summary
Retain captured graph outputs `tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py`	Both graph runners store captured outputs in strong-reference dictionaries alongside weak references.
Release outputs during graph cleanup `tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py`	Cleanup clears strong output references before graph reset and CUDA memory-pool teardown.
Remove integration test waivers `tests/integration/test_lists/waives.txt`	Removes the B200 and B300 `TestLagunaXS::test_fp8` skip entries.

Estimated code review effort: 2 (Simple) | ~10 minutes

Possibly related PRs

NVIDIA/TensorRT-LLM#16949: Modifies the same TestLagunaXS::test_fp8 waiver entry.
NVIDIA/TensorRT-LLM#17108: Directly relates to the same TestLagunaXS::test_fp8 waiver entries.

Suggested reviewers: niukuo, crazydemo, allisonlim-nv

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 42.86% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title identifies the bug, fix type, and primary change to retain owning CUDA graph output tensors.
Description check	✅ Passed	The description explains the root cause, fix, affected runners, and test plan, but it does not reproduce the repository checklist sections.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

_{Comment @coderabbitai help to get the list of available commands.}

coderabbitai

🧹 Nitpick comments (1)

tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py (1)
142-145: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Avoid Any in both strong-reference dictionaries.

These maps only retain captured outputs and do not inspect their values; use object or a precise output type alias at both sites.

tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py#L142-L145: replace Any in _graph_output_refs.

tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py#L769-L770: apply the same non-Any annotation 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

📥 Commits

Reviewing files that changed from the base of the PR and between 2e1a997 and 92a4844.

📒 Files selected for processing (2)

tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py
tests/integration/test_lists/waives.txt

💤 Files with no reviewable changes (1)

tests/integration/test_lists/waives.txt

DomBrown · 2026-07-29T17:38:02Z

/bot run

tensorrt-cicd · 2026-07-29T17:48:04Z

PR_Github #62569 [ run ] triggered by Bot. Commit: 92a4844 Link to invocation

tensorrt-cicd · 2026-07-29T18:56:39Z

PR_Github #62569 [ run ] completed with state SUCCESS. Commit: 92a4844
/LLM/main/L0_MergeRequest_PR pipeline #50713 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

DomBrown · 2026-07-29T19:28:23Z

/bot run

tensorrt-cicd · 2026-07-29T19:34:40Z

PR_Github #62585 [ run ] triggered by Bot. Commit: 92a4844 Link to invocation

tensorrt-cicd · 2026-07-29T20:09:23Z

PR_Github #62585 [ run ] completed with state FAILURE. Commit: 92a4844
/LLM/main/L0_MergeRequest_PR pipeline #50727 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

coderabbitai

🧹 Nitpick comments (1)

tests/integration/test_lists/waives.txt (1)
37-37: 📐 Maintainability & Code Quality | 🔵 Trivial | 🏗️ Heavy lift

Confirm every waiver is backed by a failed test before merging.

Test coverage summary:

Modified list: tests/integration/test_lists/waives.txt.

Added 17 SKIP entries across accuracy, visual-generation, performance, H20 LoRA, and speculative-decoding tests.

Removed 8 existing SKIP entries.

No test functions changed, and no test-db/ or qa/ list files were modified.

Coverage verdict: needs follow-up because neither cbts_touchmap.sqlite nor 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

📥 Commits

Reviewing files that changed from the base of the PR and between 92a4844 and a34f87d.

📒 Files selected for processing (2)

tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py
tests/integration/test_lists/waives.txt

fredricz-20070104

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 dicts
tensorrt_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.

DomBrown · 2026-07-30T10:47:52Z

/bot run

tensorrt-cicd · 2026-07-30T10:54:10Z

PR_Github #62763 [ run ] triggered by Bot. Commit: a34f87d Link to invocation

tensorrt-cicd · 2026-07-30T11:49:57Z

PR_Github #62763 [ run ] completed with state SUCCESS. Commit: a34f87d
/LLM/main/L0_MergeRequest_PR pipeline #50894 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

DomBrown · 2026-07-30T12:39:46Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-30T12:46:22Z

PR_Github #62772 [ run ] triggered by Bot. Commit: a34f87d Link to invocation

tensorrt-cicd · 2026-07-30T16:12:10Z

PR_Github #62772 [ run ] completed with state SUCCESS. Commit: a34f87d
/LLM/main/L0_MergeRequest_PR pipeline #50901 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

DomBrown · 2026-07-30T16:22:53Z

/bot run

tensorrt-cicd · 2026-07-30T16:29:40Z

PR_Github #62809 [ run ] triggered by Bot. Commit: a34f87d Link to invocation

tensorrt-cicd · 2026-07-30T18:23:52Z

PR_Github #62809 [ run ] completed with state SUCCESS. Commit: a34f87d
/LLM/main/L0_MergeRequest_PR pipeline #50934 completed with status: 'SUCCESS'

CI Report

Link to invocation

DomBrown · 2026-08-03T09:02:44Z

/bot run

zhaoyangwang-nvidia

LGTM

coderabbitai

Actionable comments posted: 1

🤖 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.

Inline comments:
In `@tests/integration/test_lists/waives.txt`:
- Around line 157-158: Add the 2026 NVIDIA SPDX copyright header at the
beginning of waives.txt, prefixing every header line with the file’s comment
syntax so waiver readers ignore it; leave the existing waiver entries unchanged.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

Push a commit to this branch (recommended)
Create a new PR with the fixes

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: f50c1e63-2d05-45cd-b987-d1daf1ecba68

📥 Commits

Reviewing files that changed from the base of the PR and between a34f87d and 4ad991f.

📒 Files selected for processing (2)

tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py
tests/integration/test_lists/waives.txt

🚧 Files skipped from review as they are similar to previous changes (1)

tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py

tensorrt-cicd · 2026-08-03T09:09:01Z

PR_Github #63402 [ run ] triggered by Bot. Commit: 4ad991f Link to invocation

tensorrt-cicd · 2026-08-03T09:41:53Z

PR_Github #63402 [ run ] completed with state FAILURE. Commit: 4ad991f
/LLM/main/L0_MergeRequest_PR pipeline #51381 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

… 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>

DomBrown · 2026-08-03T10:54:28Z

/bot run

tensorrt-cicd · 2026-08-03T11:00:43Z

PR_Github #63430 [ run ] triggered by Bot. Commit: e839f1d Link to invocation

tensorrt-cicd · 2026-08-03T12:39:24Z

PR_Github #63430 [ run ] completed with state SUCCESS. Commit: e839f1d
/LLM/main/L0_MergeRequest_PR pipeline #51403 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

DomBrown · 2026-08-03T12:40:16Z

/bot run --disable-fail-fast

trtllm-agent requested review from a team as code owners July 29, 2026 17:10

trtllm-agent assigned DomBrown Jul 29, 2026

trtllm-agent requested review from BowenFu, JunyiXu-nv, LarryXFly, ZhanruiSunCh, allisonlim-nv, brnguyen2, crazydemo and niukuo July 29, 2026 17:10

github-actions Bot assigned trtllm-agent Jul 29, 2026

coderabbitai Bot reviewed Jul 29, 2026

View reviewed changes

LarryXFly approved these changes Jul 30, 2026

View reviewed changes

BowenFu approved these changes Jul 30, 2026

View reviewed changes

trtllm-agent force-pushed the repair-bot-bug6525011 branch from 92a4844 to a34f87d Compare July 30, 2026 08:33

coderabbitai Bot reviewed Jul 30, 2026

View reviewed changes

fredricz-20070104 approved these changes Jul 30, 2026

View reviewed changes

coderabbitai Bot mentioned this pull request Jul 31, 2026

[None][perf] optimize encoder-decoder PyTorch performance #16706

Open

trtllm-agent force-pushed the repair-bot-bug6525011 branch from a34f87d to 4ad991f Compare August 3, 2026 08:58

lori-ren approved these changes Aug 3, 2026

View reviewed changes

zhaoyangwang-nvidia approved these changes Aug 3, 2026

View reviewed changes

DomBrown enabled auto-merge (squash) August 3, 2026 09:05

coderabbitai Bot reviewed Aug 3, 2026

View reviewed changes

Comment thread tests/integration/test_lists/waives.txt

HandongLi-01 and others added 2 commits August 3, 2026 11:54

[nvbugs/6525011][chore] Remove stale waiver after fix

e839f1d

Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>

DomBrown force-pushed the repair-bot-bug6525011 branch from 4ad991f to e839f1d Compare August 3, 2026 10:54

Uh oh!

Conversation

trtllm-agent commented Jul 29, 2026 • edited by coderabbitai Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Test plan

Links

Uh oh!

coderabbitai Bot commented Jul 29, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

❌ Failed checks (1 warning)

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

DomBrown commented Jul 29, 2026

Uh oh!

tensorrt-cicd commented Jul 29, 2026

Uh oh!

tensorrt-cicd commented Jul 29, 2026

Uh oh!

DomBrown commented Jul 29, 2026

Uh oh!

tensorrt-cicd commented Jul 29, 2026

Uh oh!

tensorrt-cicd commented Jul 29, 2026

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

fredricz-20070104 left a comment

Choose a reason for hiding this comment

Review summary - Approve

Uh oh!

DomBrown commented Jul 30, 2026

Uh oh!

tensorrt-cicd commented Jul 30, 2026

Uh oh!

tensorrt-cicd commented Jul 30, 2026

Uh oh!

DomBrown commented Jul 30, 2026

Uh oh!

tensorrt-cicd commented Jul 30, 2026

Uh oh!

tensorrt-cicd commented Jul 30, 2026

Uh oh!

DomBrown commented Jul 30, 2026

Uh oh!

tensorrt-cicd commented Jul 30, 2026

Uh oh!

tensorrt-cicd commented Jul 30, 2026

Uh oh!

DomBrown commented Aug 3, 2026

Uh oh!

zhaoyangwang-nvidia left a comment

Choose a reason for hiding this comment

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

tensorrt-cicd commented Aug 3, 2026

Uh oh!

tensorrt-cicd commented Aug 3, 2026

Uh oh!

DomBrown commented Aug 3, 2026

Uh oh!

tensorrt-cicd commented Aug 3, 2026

Uh oh!

tensorrt-cicd commented Aug 3, 2026

Uh oh!

DomBrown commented Aug 3, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

trtllm-agent commented Jul 29, 2026 •

edited by coderabbitai Bot

Loading

coderabbitai Bot commented Jul 29, 2026 •

edited

Loading