[https://nvbugs/6546909][fix] Make the handle's lifetime match its graphs — lazy… - #17313
[https://nvbugs/6546909][fix] Make the handle's lifetime match its graphs — lazy…#17313trtllm-agent wants to merge 1 commit into
Conversation
|
Note Reviews pausedIt 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 Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
WalkthroughThe GB300 skip waiver is removed for the Llama 3.1 8B FP8 integration test configuration using ChangesIntegration test waiver update
Estimated code review effort: 1 (Trivial) | ~2 minutes Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 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/compilation/backend.py (1)
124-138: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd docstrings for the graph-pool lifecycle methods.
get_graph_pool_handle()andretire_graph_pool_handle()are used across module boundaries. Replace the implementation comments with Google-style docstrings that define their lifecycle contract.🤖 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/compilation/backend.py` around lines 124 - 138, Replace the inline comments in get_graph_pool_handle() and retire_graph_pool_handle() with Google-style docstrings documenting lazy handle creation, replacement after retirement, and the requirement to retire only after captured graphs are reset. Preserve the existing implementation behavior.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/compilation/backend.py`:
- Around line 124-138: Replace the inline comments in get_graph_pool_handle()
and retire_graph_pool_handle() with Google-style docstrings documenting lazy
handle creation, replacement after retirement, and the requirement to retire
only after captured graphs are reset. Preserve the existing implementation
behavior.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 0b58bdab-da34-4805-bbb4-fc4b75c759e0
📒 Files selected for processing (3)
tensorrt_llm/_torch/compilation/backend.pytensorrt_llm/_torch/pyexecutor/model_engine.pytests/integration/test_lists/waives.txt
💤 Files with no reviewable changes (1)
- tests/integration/test_lists/waives.txt
brnguyen2
left a comment
brnguyen2
left a comment
There was a problem hiding this comment.
Root cause analysis looks right, but the retirement is at the wrong granularity. Backend._graph_pool_handle is class-level and shared by every Backend/PyTorchModelEngine in the process, while retire_graph_pool_handle() is called from a per-engine _release_cuda_graphs(). If two engines coexist (e.g. draft + target model for spec decoding, or an encoder engine), tearing down one nulls the handle that the other's still-live graphs were captured into; the next capture in the surviving engine mints a fresh pool and silently loses memory sharing, and any code still holding the old handle in _cuda_graph_mem_pool passes a retired pool to capture_begin — the exact assert this PR fixes.
Either refcount the handle (number of engines that took it) and retire on the last release, or move the handle to per-Backend-instance state.
Also: the description says self._cuda_graph_mem_pool is nulled, but the diff never does that. And there's no regression test for the build → shutdown → rebuild sequence, which is the whole failure mode; the unwaived integration test only covers a single-engine run.
| # (e.g. a second LLM built after this one shuts down) passes it to | ||
| # capture_begin. | ||
| Backend.retire_graph_pool_handle() | ||
|
|
There was a problem hiding this comment.
This retires a class-level handle from a per-engine teardown. With more than one engine alive in the process (spec-decode draft+target, encoder+decoder), releasing one engine's graphs nulls the pool that the other engine's still-captured graphs own. The survivor's next capture then mints a new pool (losing sharing), and anything that cached the old value — e.g. self._cuda_graph_mem_pool on the other engine — now holds a handle whose graphs may have been reset.
Suggest refcounting: Backend.acquire_graph_pool_handle() on engine init, release_graph_pool_handle() here, and only null when the count hits zero.
| self.spec_metadata = None | ||
| self.iter_states = {} | ||
| self._cuda_graph_mem_pool = self._torch_compile_backend._graph_pool_handle if self._torch_compile_enabled else None | ||
| self._cuda_graph_mem_pool = (Backend.get_graph_pool_handle() |
There was a problem hiding this comment.
The PR description says _release_cuda_graphs() also nulls self._cuda_graph_mem_pool, but it doesn't. After retirement this field still holds the retired handle, and any subsequent capture path that reads it (rather than re-calling get_graph_pool_handle()) will pass the dead pool to capture_begin. Please either set it to None in _release_cuda_graphs() or make it a property that always delegates to the accessor.
| torch.cuda.Event() for _ in range(num_events - len(self.events)) | ||
| ] | ||
|
|
||
| @classmethod |
There was a problem hiding this comment.
No locking here. If two Backend instances are constructed or optimize concurrently (torch.compile worker threads, or the num_streams path), this can mint two pools and leak one. Probably fine in practice given the GIL and current call sites, but worth confirming that optimize()/engine init are single-threaded.
| @@ -246,7 +246,6 @@ full:GB300/accuracy/test_llm_api_pytorch.py::TestDeepSeekV4FlashBase::test_fp8_4 | |||
| full:GB300/accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_4gpus[v1_kv_cache-dp4-trtllm-fp8] SKIP (https://nvbugs/6474894) | |||
There was a problem hiding this comment.
Unwaiving is right if the fix holds, but the failure mode is a second LLM built in a reused MPI worker — a single-engine accuracy test won't catch a regression here. Worth adding a cheap unit test that constructs an engine with torch_compile, calls _release_cuda_graphs(), then asserts Backend.get_graph_pool_handle() returns a different handle and a fresh capture succeeds.
| return cls._graph_pool_handle | ||
|
|
||
| @classmethod | ||
| def retire_graph_pool_handle(cls) -> None: |
There was a problem hiding this comment.
clear_piecewise_cuda_graphs() already implements this same "the handle is dead, mint a new one" rule by calling torch.cuda.graph_pool_handle() itself, and the two mint different ids — so after a release the surviving runners sit on one pool while the next optimize() and the CUDAGraphRunner land on another, silently losing the backend-wide pool sharing this handle exists for.
Could clear_piecewise_cuda_graphs() instead call retire_graph_pool_handle() and then assign get_graph_pool_handle() to all runners? That keeps the rotation rule in one place and preserves the sharing.
9a4ea0a to
55ddcea
Compare
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
1a2a18f to
737ce9a
Compare
…eleased Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>
737ce9a to
c4e6599
Compare
4 participants
Summary
Backend._graph_pool_handleis a process-wide class attribute never retired when_release_cuda_graphs()resets the graphs captured into it, so a later engine in the same reused MPI worker passes the now-dead private pool tocapture_beginand trips the allocator'suse_count > 0assert.Backend.get_graph_pool_handle()plusBackend.retire_graph_pool_handle()invoked at the end of_release_cuda_graphs()(also nullingself._cuda_graph_mem_pool), so the next capture generation mints a fresh handle; removed the obsolete waiver.pytest "tests/integration/defs/accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8[fp8kv=False-attn_backend=FLASHINFER-torch_compile=True]" "tests/integration/defs/accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8[fp8kv=False-attn_backend=TRTLLM-torch_compile=True]" -vTest plan
Links
Dev Engineer Review
TestLlama3_1_8BInstruct::test_fp8withfp8kv=False, TRTLLM attention, and Torch Compile.QA Engineer Review
tests/integration/test_lists/waives.txt.