[None][fix] Retire the shared CUDA graph pool handle on engine teardown - #15973
[None][fix] Retire the shared CUDA graph pool handle on engine teardown#15973yizhang-nv wants to merge 2 commits 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:
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 (3)
💤 Files with no reviewable changes (1)
🚧 Files skipped from review as they are similar to previous changes (2)
WalkthroughThe backend now creates a fresh CUDA graph pool handle after graph teardown. Model engine cleanup invokes this rotation after clearing graph runners. One waived DeepSeek integration test entry was removed. ChangesCUDA graph pool lifecycle
Estimated code review effort: 2 (Simple) | ~10 minutes Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
|
/bot run --disable-fail-fast |
|
PR_Github #57916 [ run ] triggered by Bot. Commit: |
| if (self.drafter is None | ||
| and getattr(self.model_engine, 'spec_config', None) is None): | ||
| return | ||
|
|
There was a problem hiding this comment.
Should also early return if use_spec_decode is False
|
PR_Github #57916 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #58156 [ run ] triggered by Bot. Commit: |
|
/bot run --disable-fail-fast |
|
PR_Github #58158 [ run ] triggered by Bot. Commit: |
|
PR_Github #58156 [ run ] completed with state |
|
PR_Github #58158 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #58362 [ run ] triggered by Bot. Commit: |
|
PR_Github #58362 [ run ] completed with state |
…upstream PR NVIDIA#15973) Squash of upstream PR NVIDIA#15973 (4 commits). One-engine speculative decoders (MTP / DSpark) keep the draft model inside ModelEngine, so PyExecutor.drafter is None; under the overlap scheduler the next draft tokens are produced by the previous batch and their count may not be copied to the C++ request before the next scheduling pass, so the scheduler can budget a gen request as 1 token though the model consumes 1+K -> oversized mixed context/generation batch (total_num_tokens > max_num_tokens). Run scheduler-facing draft-token reservation for every speculative decoder immediately before scheduling (_prepare_scheduler_draft_tokens); placeholder token values are not consumed, their length reserves the correct micro-batch capacity. Includes: skip when spec decode is disabled, handle lightweight executor fixtures, and unwaive the DeepSeekV3 nvfp4 mtp3_fp8kv_chunked test that exposed the bug. Signed-off-by: chungen28 <chung-en@deepinfra.com>
b71f9f0 to
1639857
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. |
yizhang-nv
changed the title
|
/bot run --disable-fail-fast --stage-list "DGX_B200-8_GPUs-PyTorch-1, DGX_B200-8_GPUs-PyTorch-2, DGX_B200-8_GPUs-PyTorch-3, DGX_B200-8_GPUs-PyTorch-4" |
|
Removed the "ci: full pre-merge approved" label because @yizhang-nv could not be verified as an active member of NVIDIA/trt-llm-ci-approvers. Ask a member of that team to apply it. |
|
PR_Github #63972 [ run ] triggered by Bot. Commit: |
|
PR_Github #63968 [ run ] completed with state |
|
PR_Github #63972 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #63992 [ run ] triggered by Bot. Commit: |
|
PR_Github #63992 [ run ] completed with state
|
yizhang-nv
changed the title
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tests/unittest/_torch/compilation/test_graph_pool_handle_rotation.py (1)
45-66: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winAssert teardown order and cover encoder cleanup.
The allocator fix depends on clearing every runner before rotating the pool. This test sets
encoder_cuda_graph_runnertoNoneand checks only that calls occurred. It does not detect a future call-order regression or an encoder cleanup regression.Add an encoder mock and assert this order: piecewise cleanup, decoder cleanup, encoder cleanup, then pool rotation.
Based on the supplied teardown and test code, call order is part of the stale-pool fix.
🤖 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/unittest/_torch/compilation/test_graph_pool_handle_rotation.py` around lines 45 - 66, Update test_release_cuda_graphs_rotates_shared_pool to use an encoder CUDA graph runner mock, then assert the teardown sequence is piecewise graph cleanup, decoder runner cleanup, encoder runner cleanup, and finally graph-pool handle rotation. Preserve the existing stale-handle assertion while making the test fail if any cleanup occurs after rotation or encoder cleanup is omitted.
🤖 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/unittest/_torch/compilation/test_graph_pool_handle_rotation.py`:
- Around line 45-66: Update test_release_cuda_graphs_rotates_shared_pool to use
an encoder CUDA graph runner mock, then assert the teardown sequence is
piecewise graph cleanup, decoder runner cleanup, encoder runner cleanup, and
finally graph-pool handle rotation. Preserve the existing stale-handle assertion
while making the test fail if any cleanup occurs after rotation or encoder
cleanup is omitted.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: dd395d9f-f888-40a6-8a00-7f803714700f
📒 Files selected for processing (3)
tensorrt_llm/_torch/compilation/backend.pytensorrt_llm/_torch/pyexecutor/model_engine.pytests/unittest/_torch/compilation/test_graph_pool_handle_rotation.py
|
/bot run --disable-fail-fast |
b32b829 to
c4a5959
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. |
|
/bot run --disable-fail-fast |
|
PR_Github #64056 [ run ] triggered by Bot. Commit: |
|
PR_Github #64057 [ run ] triggered by Bot. Commit: |
|
PR_Github #64056 [ run ] completed with state |
|
PR_Github #64057 [ run ] completed with state
|
Unwaive TestDeepSeekV32::test_nvfp4_multi_gpus_piecewise_cuda_graph[mtp3_fp8kv_chunked] (nvbugs/5989920) to re-enable it in CI. Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
Backend._graph_pool_handle is class state, so one private pool is shared by every engine in the process. With torch.compile enabled the generation CUDAGraphRunner captures into that pool. Tearing the engine down resets those graphs, dropping the pool's use_count to zero, but the caching allocator only erases the entry once the pool has no blocks left. A second engine built in the same worker process therefore captured into a retired handle and tripped beginAllocateToPool's use_count > 0 assert. clear_piecewise_cuda_graphs already rotates the per-runner handles for exactly this reason; the class-level handle -- whose only consumer is the generation runner, since piecewise_optimizer allocates its own -- was missed. Rotate it in _release_cuda_graphs once every runner has been cleared. Surfaced by back-to-back DeepSeek-V3.2 runs in one MPI worker: test_nvfp4_multi_gpus_piecewise_cuda_graph[mtp3_fp8kv_chunked] passed in isolation but failed after [baseline] in the same session. Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
|
/bot run --disable-fail-fast |
c4a5959 to
4faea3e
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. |
|
PR_Github #64174 [ run ] triggered by Bot. Commit: |
| so the next engine built in this process must start from a fresh id. | ||
| Same rationale as the per-runner rotation above. | ||
| """ | ||
| cls._graph_pool_handle = torch.cuda.graph_pool_handle() |
There was a problem hiding this comment.
I have removed the unnecessary _graph_pool_handle used inside torch compile in #16952.
6 participants
Description
accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_nvfp4_multi_gpus_piecewise_cuda_graph[mtp3_fp8kv_chunked]has been waived under NVBug 5989920. The scheduler token-budget failure it was originally waived for (total_num_tokens (8195) should be <= max_num_tokens (8192)) was fixed onmainby #16101. Re-enabling the test surfaced a second, unrelated failure, which this PR fixes.The remaining failure
All 8 ranks aborted during executor init:
That assert lives in
beginAllocateToPool: the private pool id is still present ingraph_pools, but itsuse_counthas already fallen to zero.Backend._graph_pool_handleis class state, so a single private pool is shared by every engine in the process, and with torch.compile enabled the generationCUDAGraphRunnercaptures into it (model_engine.py). Engine teardown resets those graphs, droppinguse_countto zero, but the allocator only erases the entry once the pool has no blocks left. The next engine built in the same worker process then captures into that retired handle and trips the assert.clear_piecewise_cuda_graphsalready rotates the per-runner handles for exactly this reason ("CUDACachingAllocator does not allow a private pool handle to be reused after its last graph is reset", added in #15923). The class-level handle was missed — and it is the generation runner's only consumer, sincepiecewise_optimizerallocates its own pool and ignores the handle passed to it. This PR rotates it in_release_cuda_graphsonce every runner has been cleared.This is why the test only fails in CI: it passes standalone, but the stage runs
[baseline]first and the MPI workers are reused across both.Test Coverage
8x B200,
LLM_MODELS_ROOTset, both params in one pytest session (the ordering that reproduces the bug):1 failed, 1 passed+ allocator assert on all 8 ranks[mtp3_fp8kv_chunked]alone1 passed(isolates the ordering dependency)2 passed in 1180.31s, zero assertsPost-fix accuracy:
[mtp3_fp8kv_chunked]MMLU 87.865 (ref 87.200), GSM8K 95.072 (ref 95.600);[baseline]GSM8K 95.375 (ref 95.600).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.Dev Engineer Review
Backend.rotate_graph_pool_handle()to replace the shared CUDA graph pool handle after graph teardown._release_cuda_graphsto clear TorchInductor and CUDA graph runners before rotating the handle.QA Engineer Review
tests/integration/test_lists/waives.txt.TestDeepSeekV32::test_nvfp4_multi_gpus_piecewise_cuda_graph[mtp3_fp8kv_chunked].test-db/orqa/files were modified.