[https://nvbugs/6475346][fix] Avoid stale CUD… - #16952
Conversation
…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>
liji-nv
changed the title
liji-nv
changed the title
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (1)
💤 Files with no reviewable changes (1)
WalkthroughCUDA graph pool ownership is no longer shared through ChangesCUDA graph and MRoPE execution updates
Estimated code review effort: 3 (Moderate) | ~20 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 |
|
/bot run --disable-fail-fast |
|
PR_Github #62197 [ run ] triggered by Bot. Commit: |
chienchunhung
left a comment
chienchunhung
left a comment
There was a problem hiding this comment.
Thanks for the PR!
|
PR_Github #62197 [ run ] completed with state
|
BowenFu
left a comment
BowenFu
left a comment
There was a problem hiding this comment.
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>
|
/bot run --disable-fail-fast |
|
PR_Github #62422 [ run ] triggered by Bot. Commit: |
|
PR_Github #62422 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #62667 [ run ] triggered by Bot. Commit: |
Signed-off-by: Jin Li <59594262+liji-nv@users.noreply.github.com>
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
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 winSeparate 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
📒 Files selected for processing (2)
tensorrt_llm/_torch/pyexecutor/model_engine.pytests/integration/test_lists/waives.txt
|
PR_Github #62667 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #62774 [ run ] triggered by Bot. Commit: |
6 participants
…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
tensorrt_llm/_torch/compilation/backend.py:Backend._graph_pool_handleand the associated lazy initialization/registration logic.piecewise_optimizeris no longer passed a shared graph pool handle; piecewise orchestration (runner updates, event generation, and returning the compiledGraphModule) remains intact.tensorrt_llm/_torch/compilation/piecewise_optimizer.pyAPI to remove the externally providedgraph_pool_handleparameter:piecewise_optimizernow creates/initializes its own internal pool viatorch.cuda.graph_pool_handle().tensorrt_llm/_torch/pyexecutor/model_engine.pyto reinforce pool ownership boundaries:_cuda_graph_mem_poolis no longer derived from the torch-compile backend pool handle and is initialized toNone(first capture creates a private pool; piecewise graphs keep private pool ownership; outer generation usesgraph.pool()to create/reuse its own pool)._prepare_tp_inputs:has_gen_mrope_deltato detect whether any generation request supplies non-zeromrope_position_deltas.mrope_config/deltas, it records a reserved “zero delta” slot (mrope_dummy_seq_slot) and clearsmrope_delta_read_seq_slotswhen no real deltas exist to preserve the steady-state fast path.mrope_position_idsby only appending when present.QA Engineer Review
Test-list change:
tests/integration/test_lists/waives.txtfull:L40S/...::TestQwen3_30B_A3B::test_fp8[latency-torch_compile=True]perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-b200_deepseek-r1-fp4_...]full:B300/...::TestGemma3_1BInstruct::test_fp8_prequantized[torch_compile=True]full:GB300/...::TestGemma3_1BInstruct::test_fp8_prequantized[torch_compile=True]full:RTX_PRO_6000_Blackwell_Server_Edition/...::TestDeepSeekV3Lite::test_nvfp4mtp_nextn=0andmtp_nextn=2with a narrower set that keeps onlycuda_graph=Falseandoverlap_scheduler=Falsecombinations (dropping broader variants includingcuda_graph=True/ other parameter combos).Verdict: needs follow-up (CBTS coverage data not provided).
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-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin 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.