[None][fix] avoid attention workspace resize during CUDA graph capture - #16319
Conversation
|
/bot run --disable-fail-fast |
|
PR_Github #58950 [ run ] triggered by Bot. Commit: |
|
/bot run --disable-fail-fast |
|
PR_Github #58980 [ run ] triggered by Bot. Commit: |
|
PR_Github #58950 [ run ] completed with state
|
📝 WalkthroughWalkthroughCUDA graph execution now uses explicit warmup-only and capture-only phases for decoder and encoder paths, validates prepared graph metadata, captures and returns outputs, and distinguishes direct warmup outputs from graph replay. ChangesCUDA graph capture flow
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant ModelEngine
participant CUDAGraphRunner
participant EncoderCUDAGraphRunner
participant CUDAGraph
ModelEngine->>CUDAGraphRunner: run decoder warmup-only pass
CUDAGraphRunner->>CUDAGraph: prepare graph metadata
ModelEngine->>CUDAGraphRunner: run decoder capture-only pass
CUDAGraphRunner->>CUDAGraph: capture and store graph output
ModelEngine->>EncoderCUDAGraphRunner: run encoder warmup-only pass
EncoderCUDAGraphRunner->>CUDAGraph: prepare graph metadata
ModelEngine->>EncoderCUDAGraphRunner: run encoder capture-only pass
EncoderCUDAGraphRunner->>CUDAGraph: capture and store graph output
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.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tensorrt_llm/_torch/pyexecutor/model_engine.py (1)
5267-5292: 🩺 Stability & Availability | 🔴 Critical | ⚡ Quick winWarmup-only execution can read an uninitialized capture result.
Existing or duplicate graph keys make
needs_capturefalse, but both warmup-only branches still consumecapture_outputs, causingUnboundLocalError.
tensorrt_llm/_torch/pyexecutor/model_engine.py#L5267-L5292: invokecapture()whenever encoder warmup-only mode is active.tensorrt_llm/_torch/pyexecutor/model_engine.py#L5466-L5488: apply the same condition to decoder capture.🤖 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 5267 - 5292, The encoder capture flow around encoder_cuda_graph_runner.capture and the corresponding decoder flow must invoke capture whenever warmup-only mode is active, even when needs_capture is false; update both sites in tensorrt_llm/_torch/pyexecutor/model_engine.py (5267-5292 and 5466-5488) so capture runs when needs_capture or is_warmup_only is true, while preserving normal replay behavior otherwise.
🤖 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 `@tensorrt_llm/_torch/pyexecutor/model_engine.py`:
- Around line 5267-5292: The encoder capture flow around
encoder_cuda_graph_runner.capture and the corresponding decoder flow must invoke
capture whenever warmup-only mode is active, even when needs_capture is false;
update both sites in tensorrt_llm/_torch/pyexecutor/model_engine.py (5267-5292
and 5466-5488) so capture runs when needs_capture or is_warmup_only is true,
while preserving normal replay behavior otherwise.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: d1feb614-cdd3-4ddf-9677-211ccaed2e72
📒 Files selected for processing (2)
tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.pytensorrt_llm/_torch/pyexecutor/model_engine.py
|
/bot run --disable-fail-fast |
|
PR_Github #59109 [ run ] triggered by Bot. Commit: |
|
PR_Github #58980 [ run ] completed with state |
|
/bot run --disable-fail-fast |
|
PR_Github #59128 [ run ] triggered by Bot. Commit: |
|
PR_Github #59109 [ run ] completed with state |
|
/bot run --disable-fail-fast |
|
PR_Github #59137 [ run ] triggered by Bot. Commit: |
|
PR_Github #59128 [ run ] completed with state |
Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
|
/bot run --disable-fail-fast |
|
PR_Github #59413 [ run ] triggered by Bot. Commit: |
|
PR_Github #59413 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #59469 [ run ] triggered by Bot. Commit: |
|
PR_Github #59469 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #59684 [ run ] triggered by Bot. Commit: |
|
PR_Github #59684 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #59875 [ run ] triggered by Bot. Commit: |
|
PR_Github #59875 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #59962 [ run ] triggered by Bot. Commit: |
|
PR_Github #59962 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #60240 [ run ] triggered by Bot. Commit: |
|
PR_Github #60240 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #60340 [ run ] triggered by Bot. Commit: |
|
PR_Github #60340 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #60573 [ run ] triggered by Bot. Commit: |
|
PR_Github #60573 [ run ] completed with state |
6 participants
Summary by CodeRabbit
Description
CUDA graph warmup processes generation graph shapes from larger to smaller batch sizes so later captures can reuse the attention workspace. This assumption does not always hold: an attention kernel can switch implementations at a smaller batch size and require a larger workspace. Resizing the shared workspace after an earlier graph was captured invalidates the workspace address stored in that graph.
Separate eager warmup from graph capture for both generation and encoder CUDA graph runners.
PyTorchModelEnginefirst runs every graph shape in warmup-only mode. This performs the normal eager warmup iterations, grows the shared attention workspace to the maximum required size, and retains the prepared graph metadata, but does not create CUDA graphs or a graph memory pool. It then runs the same shapes in capture-only mode, reuses the prepared metadata, and captures the runtime graphs without repeating the eager warmups or resizing the workspace. Piecewise CUDA graphs are captured only during the capture-only pass.This keeps the same two eager warmups plus one capture per graph shape as before the PR while avoiding the previous double-capture approach. It also removes the need to disable replay, clear captured graph state, or replace and restore the CUDA graph memory pool, avoiding allocator pool lifetime conflicts.
Test Coverage
tests/integration/defs/accuracy/test_llm_api_pytorch.pyon H200.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.