[https://nvbugs/6404567][fix] Clamp piecewise cudagraph captures to the reachable ceiling instead of force-appending it - #16256
Conversation
📝 WalkthroughWalkthroughPiecewise CUDA graph capture filtering now clamps requested sizes to the reachable ceiling without appending unrequested ceiling values. Warmup warnings and unit tests were updated to reflect the revised behavior. ChangesPiecewise capture ceiling clamping
Estimated code review effort: 3 (Moderate) | ~20 minutes 🚥 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)
479-485: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winWarning text contradicts the new clamping behavior.
The message still opens with
"Skipping piecewise CUDA graph capture for num_tokens=..."but the revised logic no longer skips these sizes — they are clamped down to the reachable ceiling and a graph is still captured there. The trailing"Clamping them to the ceiling"conflicts with the leading "Skipping", which will read as contradictory to operators triaging logs.📝 Suggested wording
if unrecordable: logger.warning( - f"Skipping piecewise CUDA graph capture for num_tokens=" - f"{unrecordable}: exceeds reachable ceiling " + f"Piecewise CUDA graph capture requested for num_tokens=" + f"{unrecordable} exceeds reachable ceiling " f"max_batch_size*(max_seq_len-1-num_extra_decoding_steps)=" f"{max(0, self.batch_size * (self.max_seq_len - 1 - num_extra_decoding_steps))}. " f"Clamping them to the ceiling; raise max_seq_len for larger graphs." )🤖 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 479 - 485, Update the warning in the piecewise CUDA graph capture logic to describe clamping rather than skipping: revise the message in the relevant capture path to state that the unrecordable token sizes are clamped to the reachable ceiling and captured at that ceiling, while retaining the ceiling calculation and guidance to increase max_seq_len.
🤖 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 479-485: Update the warning in the piecewise CUDA graph capture
logic to describe clamping rather than skipping: revise the message in the
relevant capture path to state that the unrecordable token sizes are clamped to
the reachable ceiling and captured at that ceiling, while retaining the ceiling
calculation and guidance to increase max_seq_len.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 218b23e5-f7fa-4a70-86c3-0ac33f149f52
📒 Files selected for processing (2)
tensorrt_llm/_torch/pyexecutor/model_engine.pytests/unittest/llmapi/test_llm_args.py
|
/bot run --disable-fail-fast |
|
PR_Github #58715 [ run ] triggered by Bot. Commit: |
|
/bot run --disable-fail-fast |
|
PR_Github #58715 [ run ] completed with state
|
|
PR_Github #58733 [ run ] triggered by Bot. Commit: |
|
PR_Github #58733 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #58884 [ run ] triggered by Bot. Commit: |
|
PR_Github #58884 [ run ] completed with state |
ed4ad61 to
cf25857
Compare
|
/bot run --disable-fail-fast |
…he reachable ceiling instead of force-appending it _filter_piecewise_capture_num_tokens force-appended the reachable ceiling (min(max_num_tokens, max_batch_size*(max_seq_len-1))) to the piecewise capture list whenever the user's largest entry was below it. Runtime padding rounds each iteration up to the nearest captured size, so a far appended ceiling (65536 over a user list topping at 13914) made every (13914, 65536]-token iteration execute the full 65536-token graph: up to ~2.8x device time on those iterations, -12% QPS / +146% TTFT on gpt-oss-120b GB300 IFB serving (the rc14->rc15 regression). Replace drop-and-append with clamping: user-requested sizes above the reachable ceiling are lowered to the ceiling (a requested 128 becomes 127 when only 127 is recordable, preserving the NVBug 5615248 fix bit-for-bit), and no capture size beyond the user's list is ever invented. Configs without oversized entries keep exactly their own capture list, i.e. pre-rc15 behavior. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Dongfeng Yu <dongfengy@nvidia.com>
cf25857 to
7e4069f
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #59300 [ run ] triggered by Bot. Commit: |
|
PR_Github #59304 [ run ] triggered by Bot. Commit: |
|
PR_Github #59300 [ run ] completed with state |
brb-nv
left a comment
brb-nv
left a comment
There was a problem hiding this comment.
Thank you for looking into this!
|
/bot run --disable-fail-fast |
|
PR_Github #59442 [ run ] triggered by Bot. Commit: |
|
PR_Github #59304 [ run ] completed with state |
|
/bot run --disable-fail-fast |
|
PR_Github #59444 [ run ] triggered by Bot. Commit: |
|
PR_Github #59442 [ run ] completed with state |
|
/bot run --disable-fail-fast |
|
PR_Github #59444 [ run ] completed with state
|
|
PR_Github #59461 [ run ] triggered by Bot. Commit: |
|
/bot run --disable-fail-fast |
|
PR_Github #59461 [ run ] completed with state
|
|
PR_Github #59495 [ run ] triggered by Bot. Commit: |
|
PR_Github #59495 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #59582 [ run ] triggered by Bot. Commit: |
|
PR_Github #59582 [ run ] completed with state |
4 participants
Summary by CodeRabbit
Bug Fixes
Tests
Description
PR #13574 fixed https://nvbugs/5615248 by (a) dropping piecewise CUDA-graph capture candidates above the reachable ceiling
max_batch_size*(max_seq_len-1-num_extra_decoding_steps)and (b) force-appending the ceiling itself to the capture list.(b) causes a serving perf regression (https://nvbugs/6404567, first shipped in 1.3.0rc15, still present on main): runtime padding rounds each context-bearing iteration up to the nearest captured size, so appending a far ceiling — 65536 over a user
capture_num_tokenslist topping at 13914 in the gpt-oss-120b GB300 serving config — makes every (13914, 65536]-token iteration execute the full 65536-token graph. Measured per-iteration device time is flat ~325-400 ms vs 116-408 ms true-size eager (up to 2.8x), which serializes all in-flight decode streams behind those iterations: -12% QPS, TTFT more than doubles at high concurrency.This PR replaces drop-and-append with clamping: user-requested sizes above the reachable ceiling are lowered to the ceiling (a requested 128 becomes 127 when only 127 is recordable — the resulting capture list for the 5615248 config is preserved bit-for-bit), and no capture size beyond the user's list is ever invented. Configs without oversized entries keep exactly their own capture list, i.e. pre-1.3.0rc15 behavior.
Validation (GB300, gpt-oss-120b, production serving config from the bug, 4x TP1 servers, 896 concurrent streams/server, 24000 requests per arm, steady-window aggregates):
The main A/B used one build (identical native library SHAs in both arms' manifests); only the Python delta of this commit differs.
Test Coverage
tests/unittest/llmapi/test_llm_args.py::TestPiecewiseCudaGraphCaptureDefaults: 2 new tests (a far ceiling is never invented — the exact 6404567 shape; multiple oversized candidates collapse to the clamped ceiling), 3 existing tests updated to the clamp semantics, and the NVBugs 5615248 regression tests pass unchanged (128 is clamped to 127 instead of dropped+re-appended, producing the same list). 13/13 selected tests pass.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.
🤖 Generated with Claude Code