[https://nvbugs/6475346][fix] Avoid stale CUD… by liji-nv · Pull Request #16952 · NVIDIA/TensorRT-LLM

liji-nv · 2026-07-28T12:37:19Z

…A graph pool reuse

A torch.compile Backend cached a CUDA graph pool handle at class scope and passed that handle to the outer generation CUDA graph runner. When the QA session-reuse path created a second LLM in the same MPI worker, the first LLM had already reset its graphs and released the allocator private pool, but the Backend class still retained the old handle tuple. The next generation-only capture attempted to reuse that stale pool ID and failed in CUDACachingAllocator with use_count > 0.

Piecewise CUDA graphs did not actually use this pre-created handle: piecewise_optimizer immediately replaced the argument with a new graph_pool_handle and shared that independently owned pool among its runners. Remove the unused Backend parameter and class-level registration. Let PCG continue to own its private pool, and initialize the outer generation runner without a pool so its first successful capture creates the pool that subsequent graphs reuse through graph.pool().

Unwaive the affected Gemma3 FP8 torch.compile cases on B300 and GB300, and the affected DeepSeekV3Lite NVFP4 torch.compile cases on RTX PRO 6000 Blackwell.

Validation: reproduced the failure with the DeepSeekV3Lite-to-Gemma3 session-reuse sequence on CMH GB300 (job 2707269), then verified the same sequence completes successfully with this fix (job 2707407).

Fixes https://nvbugs/6475346

Fixes https://nvbugs/6473373

Dev Engineer Review

Removed cached Backend-level CUDA graph pool reuse.
Updated piecewise_optimizer to create its own pool with torch.cuda.graph_pool_handle().
Updated PyTorchModelEngine to initialize _cuda_graph_mem_pool to None and preserve runner-owned pool ownership.
The CUDA graph pool ownership changes are consistent across the backend, piecewise optimizer, and model engine.

QA Engineer Review

Modified tests/integration/test_lists/waives.txt.
Removed the Gemma3 FP8 torch-compile waivers for B300 and GB300.
Reduced RTX PRO 6000 CUTLASS NVFP4 waivers to the cuda_graph=False and overlap_scheduler=False combinations.
No test-db or QA coverage data was provided.
Verdict: needs follow-up.

Description

Test Coverage

PR Checklist

Please review the following before submitting your PR:

PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.
Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

…A graph pool reuse A torch.compile Backend cached a CUDA graph pool handle at class scope and passed that handle to the outer generation CUDA graph runner. When the QA session-reuse path created a second LLM in the same MPI worker, the first LLM had already reset its graphs and released the allocator private pool, but the Backend class still retained the old handle tuple. The next generation-only capture attempted to reuse that stale pool ID and failed in CUDACachingAllocator with use_count > 0. Piecewise CUDA graphs did not actually use this pre-created handle: piecewise_optimizer immediately replaced the argument with a new graph_pool_handle and shared that independently owned pool among its runners. Remove the unused Backend parameter and class-level registration. Let PCG continue to own its private pool, and initialize the outer generation runner without a pool so its first successful capture creates the pool that subsequent graphs reuse through graph.pool(). Unwaive the affected Gemma3 FP8 torch.compile cases on B300 and GB300, and the affected DeepSeekV3Lite NVFP4 torch.compile cases on RTX PRO 6000 Blackwell. Validation: reproduced the failure with the DeepSeekV3Lite-to-Gemma3 session-reuse sequence on CMH GB300 (job 2707269), then verified the same sequence completes successfully with this fix (job 2707407). Fixes https://nvbugs/6475346 Fixes https://nvbugs/6473373 Signed-off-by: Jin Li <59594262+liji-nv@users.noreply.github.com>

coderabbitai · 2026-07-28T12:43:06Z

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

@coderabbitai resume to resume automatic reviews.
@coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

▶️ Resume reviews
🔍 Trigger review

Walkthrough

CUDA graph pool ownership no longer passes through Backend. piecewise_optimizer creates its own pool handle, and PyTorchModelEngine initializes without a shared pool. Integration waivers are removed or narrowed.

Changes

CUDA graph pool ownership

Layer / File(s)	Summary
Piecewise optimizer pool ownership `tensorrt_llm/_torch/compilation/backend.py`, `tensorrt_llm/_torch/compilation/piecewise_optimizer.py`	`Backend` no longer stores or passes a graph pool handle. `piecewise_optimizer` initializes the handle internally.
Engine CUDA graph pool initialization `tensorrt_llm/_torch/pyexecutor/model_engine.py`	`PyTorchModelEngine` initializes `_cuda_graph_mem_pool` to `None` and documents separate pool ownership.
Integration waiver updates `tests/integration/test_lists/waives.txt`	Gemma3 waivers are removed. RTX PRO 6000 CUTLASS NVFP4 waivers are narrowed to four combinations.

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

Suggested reviewers: mikeiovine, reasonsolo, jieli-matrix, mlefeb01

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% 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 clearly identifies the fix for stale CUDA graph pool reuse and follows the required NVBugs and type format.
Description check	✅ Passed	The description explains the issue, solution, affected cases, validation results, and includes the required checklist.
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.}

liji-nv · 2026-07-28T12:45:06Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-28T12:51:10Z

PR_Github #62197 [ run ] triggered by Bot. Commit: 9eb0439 Link to invocation

chienchunhung

Thanks for the PR!

tensorrt-cicd · 2026-07-28T18:50:25Z

PR_Github #62197 [ run ] completed with state FAILURE. Commit: 9eb0439
/LLM/main/L0_MergeRequest_PR pipeline #50367 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

BowenFu

Verified the default path is unchanged: `torch_compile_config` defaults to None, so `_cuda_graph_mem_pool` was already None there — this only converges the torch.compile path onto it. The dropped `piecewise_optimizer` arg was dead (overwritten at piecewise_optimizer.py:281), and all 6 waiver removals match exactly one line each on main with nothing re-added.

Not blocking, but worth noting: the validation you cite (jobs 2707269/2707407) is the GB300 session-reuse sequence, which covers 6475346; the four RTX PRO 6000 unwaives for 6473373 ride on same-root-cause reasoning rather than a cited run on that platform. Also the L0 failure on this head is unrelated — the two distinct tests are sampler logprobs and NemotronH LoRA pp1, neither of which touches torch.compile (and test_lora_pp2_sanity is already waived under 6428124).

Signed-off-by: Jin Li <59594262+liji-nv@users.noreply.github.com>

liji-nv · 2026-07-29T05:07:32Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-29T05:13:41Z

PR_Github #62422 [ run ] triggered by Bot. Commit: 5bd0663 Link to invocation

tensorrt-cicd · 2026-07-29T09:03:20Z

PR_Github #62422 [ run ] completed with state FAILURE. Commit: 5bd0663
/LLM/main/L0_MergeRequest_PR pipeline #50584 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

liji-nv · 2026-07-30T02:54:06Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-30T02:59:45Z

PR_Github #62667 [ run ] triggered by Bot. Commit: 5bd0663 Link to invocation

Signed-off-by: Jin Li <59594262+liji-nv@users.noreply.github.com>

coderabbitai

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

tests/integration/test_lists/waives.txt (1)
91-104: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win

Separate unrelated waiver additions.

These new SKIPs cover unrelated Nemotron, Qwen, Wan, perf, and DFlash cases; move them to their owning fixes so this CUDA-graph PR remains attributable and reversible. As per coding guidelines, “Keep each pull request focused on one concern; split unrelated changes into separate pull requests.”

Also applies to: 142-144, 237-237, 254-254, 336-338, 366-367
🤖 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` around lines 91 - 104, Remove the
unrelated waiver entries added alongside the listed Nemotron, Qwen, Wan,
performance, and DFlash cases from the current change, leaving only waivers
owned by this CUDA-graph fix. Move each removed SKIP entry to the pull request
or change that addresses its corresponding failure, preserving the existing
waiver format and attribution.
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.

Outside diff comments:
In `@tests/integration/test_lists/waives.txt`:
- Around line 91-104: Remove the unrelated waiver entries added alongside the
listed Nemotron, Qwen, Wan, performance, and DFlash cases from the current
change, leaving only waivers owned by this CUDA-graph fix. Move each removed
SKIP entry to the pull request or change that addresses its corresponding
failure, preserving the existing waiver format and attribution.

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 3659d7ec-75a6-49e5-bddb-bce7602c17d5

📥 Commits

Reviewing files that changed from the base of the PR and between 5bd0663 and d1c46ec.

📒 Files selected for processing (2)

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

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

PR_Github #62667 [ run ] completed with state SUCCESS. Commit: 5bd0663
/LLM/main/L0_MergeRequest_PR pipeline #50808 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

liji-nv · 2026-07-30T12:43:02Z

/bot run --disable-fail-fast

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

PR_Github #62774 [ run ] triggered by Bot. Commit: f2869b8 Link to invocation

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

PR_Github #62774 [ run ] completed with state FAILURE. Commit: f2869b8
/LLM/main/L0_MergeRequest_PR pipeline #50904 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

liji-nv · 2026-07-31T02:40:01Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-31T02:46:21Z

PR_Github #62922 [ run ] triggered by Bot. Commit: f2869b8 Link to invocation

tensorrt-cicd · 2026-07-31T07:25:59Z

PR_Github #62922 [ run ] completed with state FAILURE. Commit: f2869b8
/LLM/main/L0_MergeRequest_PR pipeline #51044 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

liji-nv · 2026-07-31T07:48:01Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-31T07:53:20Z

PR_Github #62993 [ run ] triggered by Bot. Commit: f2869b8 Link to invocation

tensorrt-cicd · 2026-07-31T13:58:00Z

PR_Github #62993 [ run ] completed with state FAILURE. Commit: f2869b8
/LLM/main/L0_MergeRequest_PR pipeline #51102 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

Signed-off-by: Jin Li <59594262+liji-nv@users.noreply.github.com>

liji-nv · 2026-08-03T02:39:17Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-08-03T02:45:31Z

PR_Github #63337 [ run ] triggered by Bot. Commit: 3178c9c Link to invocation

tensorrt-cicd · 2026-08-03T06:23:17Z

PR_Github #63337 [ run ] completed with state FAILURE. Commit: 3178c9c
/LLM/main/L0_MergeRequest_PR pipeline #51330 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

Signed-off-by: Jin Li <59594262+liji-nv@users.noreply.github.com>

liji-nv · 2026-08-03T07:11:01Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-08-03T07:16:28Z

PR_Github #63387 [ run ] triggered by Bot. Commit: 38143cf Link to invocation

liji-nv requested review from a team as code owners July 28, 2026 12:37

liji-nv requested review from BowenFu, StanleySun639, brnguyen2, chienchunhung, jieli-matrix, lori-ren, mikeiovine, mlefeb01, mzweilz and reasonsolo July 28, 2026 12:37

github-actions Bot assigned liji-nv Jul 28, 2026

liji-nv changed the title ~~[https://nvbugs/6475346][https://nvbugs/6473373][fix] Avoid stale CUD…~~ [https://nvbugs/6475346 https://nvbugs/6473373][fix] Avoid stale CUD… Jul 28, 2026

liji-nv changed the title ~~[https://nvbugs/6475346 https://nvbugs/6473373][fix] Avoid stale CUD…~~ [https://nvbugs/6475346][fix] Avoid stale CUD… Jul 28, 2026

chienchunhung approved these changes Jul 28, 2026

View reviewed changes

coderabbitai Bot mentioned this pull request Jul 28, 2026

[https://nvbugs/6523785][fix] Added Backend.get_graph_pool_handle() / retire_graph_pool_handle() and… #16968

Open

2 tasks

BowenFu approved these changes Jul 29, 2026

View reviewed changes

Merge branch 'main' into agent/fix-pcg-cuda-graph-pool-reuse

5bd0663

Signed-off-by: Jin Li <59594262+liji-nv@users.noreply.github.com>

StanleySun639 approved these changes Jul 29, 2026

View reviewed changes

liji-nv mentioned this pull request Jul 29, 2026

[None][fix] Fix Backend's CUDA graph pool handle lifetime #16995

Closed

1 task

Merge branch 'main' into agent/fix-pcg-cuda-graph-pool-reuse

d1c46ec

Signed-off-by: Jin Li <59594262+liji-nv@users.noreply.github.com>

lori-ren approved these changes Jul 30, 2026

View reviewed changes

coderabbitai Bot reviewed Jul 30, 2026

View reviewed changes

Merge branch 'main' into agent/fix-pcg-cuda-graph-pool-reuse

f2869b8

liji-nv enabled auto-merge (squash) July 30, 2026 12:43

coderabbitai Bot mentioned this pull request Jul 30, 2026

[https://nvbugs/6305404][chore] Unwaive DeepSeek V3 Lite L0 test #16975

Merged

Merge branch 'main' into agent/fix-pcg-cuda-graph-pool-reuse

3178c9c

Signed-off-by: Jin Li <59594262+liji-nv@users.noreply.github.com>

Merge branch 'main' into agent/fix-pcg-cuda-graph-pool-reuse

38143cf

Signed-off-by: Jin Li <59594262+liji-nv@users.noreply.github.com>

Uh oh!

Conversation

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

Uh oh!

Dev Engineer Review

QA Engineer Review

Description

Test Coverage

PR Checklist

GitHub Bot Help

Uh oh!

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

Uh oh!

Reviews paused

Walkthrough

Changes

❌ Failed checks (1 warning)

Uh oh!

liji-nv commented Jul 28, 2026

Uh oh!

tensorrt-cicd commented Jul 28, 2026

Uh oh!

chienchunhung left a comment

Choose a reason for hiding this comment

Uh oh!

tensorrt-cicd commented Jul 28, 2026

Uh oh!

BowenFu left a comment

Choose a reason for hiding this comment

Uh oh!

liji-nv commented Jul 29, 2026

Uh oh!

tensorrt-cicd commented Jul 29, 2026

Uh oh!

tensorrt-cicd commented Jul 29, 2026

Uh oh!

liji-nv commented Jul 30, 2026

Uh oh!

tensorrt-cicd commented Jul 30, 2026

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

tensorrt-cicd commented Jul 30, 2026

Uh oh!

liji-nv commented Jul 30, 2026

Uh oh!

tensorrt-cicd commented Jul 30, 2026

Uh oh!

tensorrt-cicd commented Jul 30, 2026

Uh oh!

liji-nv commented Jul 31, 2026

Uh oh!

tensorrt-cicd commented Jul 31, 2026

Uh oh!

tensorrt-cicd commented Jul 31, 2026

Uh oh!

liji-nv commented Jul 31, 2026

Uh oh!

tensorrt-cicd commented Jul 31, 2026

Uh oh!

tensorrt-cicd commented Jul 31, 2026

Uh oh!

liji-nv commented Aug 3, 2026

Uh oh!

tensorrt-cicd commented Aug 3, 2026

Uh oh!

tensorrt-cicd commented Aug 3, 2026

Uh oh!

liji-nv commented Aug 3, 2026

Uh oh!

tensorrt-cicd commented Aug 3, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

liji-nv commented Jul 28, 2026 •

edited by coderabbitai Bot

Loading

coderabbitai Bot commented Jul 28, 2026 •

edited

Loading