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

[trtllm-agent]

Summary

• Root cause: Backend._graph_pool_handle is a process-wide class attribute that was kept after _release_cuda_graphs() reset every graph captured into that private pool, so a second LLM in the same reused MPI worker got a retired handle whose use_count was already 0, tripping CUDACachingAllocator's ensureExistsAndIncrefPool assert.
• Fix: Added Backend.get_graph_pool_handle() / retire_graph_pool_handle() and retire the shared handle (plus the engine's cached copy) in _release_cuda_graphs(), so the next capture lazily allocates a fresh pool.
• 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/6523785

Dev Engineer Review

• Updated Backend to manage the process-wide CUDA graph pool handle via two new lifecycle APIs: Backend.get_graph_pool_handle() for lazy creation and Backend.retire_graph_pool_handle() for clearing the cached handle after graph teardown.
• Adjusted Backend.optimize()’s piecewise_cuda_graph path to pass self.get_graph_pool_handle() into piecewise_optimizer (instead of using a pre-initialized cached field).
• Updated PyTorchModelEngine.__init__ to initialize _cuda_graph_mem_pool from Backend.get_graph_pool_handle() when torch.compile is enabled, and updated PyTorchModelEngine._release_cuda_graphs() to retire the shared pool handle (Backend.retire_graph_pool_handle()) and set _cuda_graph_mem_pool to None after clearing runners.
• Ensures graphs reset in one engine/worker sequence cannot accidentally reuse a now-dead CUDA graph pool handle in later engines within the same process.

QA Engineer Review

No test changes.

[coderabbitai]

Walkthrough

The change makes CUDA graph pool handle creation lazy at the backend level and adds explicit retirement. PyTorchModelEngine uses the shared accessor during initialization and clears both the backend handle and local pool reference during CUDA graph release.

Changes

CUDA graph pool lifecycle

Layer / File(s)	Summary
Backend handle lifecycle tensorrt_llm/_torch/compilation/backend.py	Backend initialization no longer eagerly creates the graph pool handle. Class methods now lazily acquire or retire it, and piecewise CUDA graph optimization uses the accessor.
Model engine integration tensorrt_llm/_torch/pyexecutor/model_engine.py	Model engine initialization obtains the shared handle through Backend.get_graph_pool_handle(), while CUDA graph release retires the handle and clears the local pool reference.

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

Suggested reviewers: juney-nvidia

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title is clearly related to the main change: adding lazy graph pool handle management and retirement.
Description check	✅ Passed	The summary and test plan cover the issue, fix, and verification, but the template's Description and PR Checklist sections are missing.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
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]

#16952 Fixed the issue.

[coderabbitai]

🧹 Nitpick comments (1)

tensorrt_llm/_torch/pyexecutor/model_engine.py (1)

2750-2754: 🩺 Stability & Availability | 🔵 Trivial | 🏗️ Heavy lift

Add a regression test for pool rotation.

Exercise first-engine capture → _release_cuda_graphs() → second-engine capture in the same MPI worker, verifying that the second capture uses a fresh pool and does not hit the allocator assertion.

🤖 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/model_engine.py` around lines 2750 - 2754, Add
a regression test covering two engines in the same MPI worker: capture CUDA
graphs with the first engine, call its _release_cuda_graphs() method, then
capture with the second engine. Assert that the second capture uses a
different/fresh graph memory pool and completes without the allocator assertion,
using the existing engine and CUDA-graph test helpers.

🤖 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/model_engine.py`:
- Around line 2750-2754: Add a regression test covering two engines in the same
MPI worker: capture CUDA graphs with the first engine, call its
_release_cuda_graphs() method, then capture with the second engine. Assert that
the second capture uses a different/fresh graph memory pool and completes
without the allocator assertion, using the existing engine and CUDA-graph test
helpers.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 7ec5762d-16e0-41d1-826b-9475ba998939

📥 Commits

Reviewing files that changed from the base of the PR and between 5523ee2 and 4b90dc1.

📒 Files selected for processing (2)

• tensorrt_llm/_torch/compilation/backend.py
• tensorrt_llm/_torch/pyexecutor/model_engine.py