[#11422][feat] AutoDeploy: Piecewise cudagraph support Prototype - #11515
Merged
nvchenghaoz merged 19 commits intoMar 6, 2026
Conversation
Contributor
📝 WalkthroughWalkthroughThis pull request introduces a dual-mode CUDA graph execution system with piecewise graph splitting support. It adds graph dispatch logic to select between monolithic (decode-only) and piecewise (prefill/mixed) paths, includes utilities to partition graphs at dynamic ops, and integrates piecewise compilation into the backend pipeline with configurable token buckets. Changes
Sequence DiagramsequenceDiagram
participant User as User Code
participant Dispatcher as DualModeCapturedGraph
participant Monolithic as Monolithic<br/>CapturedGraph
participant Piecewise as PiecewiseCapturedGraph
participant Runner as ADPiecewiseRunner
participant CUDA as CUDA Graph<br/>Execution
User->>Dispatcher: forward(args, kwargs)
Note over Dispatcher: Inspect batch_info_host<br/>and token counts
alt Decode-only path
Dispatcher->>Monolithic: forward(args)
Monolithic->>CUDA: Replay monolithic<br/>graph
CUDA-->>Monolithic: outputs
Monolithic-->>Dispatcher: outputs
else Prefill/Mixed path
Dispatcher->>Piecewise: forward(args)
Piecewise->>Piecewise: Find matching<br/>bucket
Piecewise->>Runner: forward(submodule_args)
Note over Runner: Determine phase<br/>(warmup/capture/replay)
alt Warmup Phase
Runner->>Runner: Track tensor ptrs
Runner->>Runner: Run eagerly
else Capture Phase
Runner->>Runner: Identify dynamic inputs
Runner->>CUDA: Capture graph
CUDA-->>Runner: Graph handle
Runner->>Runner: Register static outputs
else Replay Phase
Runner->>Runner: Copy dynamic inputs
Runner->>CUDA: Replay graph
CUDA-->>Runner: outputs
end
Runner-->>Piecewise: outputs
Piecewise-->>Dispatcher: outputs
end
Dispatcher-->>User: outputs
Estimated Code Review Effort🎯 4 (Complex) | ⏱️ ~50 minutes 🚥 Pre-merge checks | ✅ 1 | ❌ 3❌ Failed checks (3 warnings)
✅ Passed checks (1 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 10
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (3)
tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/cuda_backend_causal_conv.py (1)
1-1:⚠️ Potential issue | 🟡 MinorUpdate copyright year to 2026.
The file is being modified in 2026, but the copyright header still reads
2022-2025. As per coding guidelines, the year should be updated on modified files.-# SPDX-FileCopyrightText: Copyright (c) 2022-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/triton_backend_mamba.py (1)
1-1:⚠️ Potential issue | 🟡 MinorCopyright year should be updated to 2026.
The header says
2022-2025but this file is being modified in 2026. As per coding guidelines, "update year on modified files."Proposed fix
-# SPDX-FileCopyrightText: Copyright (c) 2022-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.tensorrt_llm/_torch/auto_deploy/transform/library/compile_model.py (1)
1-1:⚠️ Potential issue | 🟡 MinorMissing NVIDIA copyright header on modified file.
This file lacks a copyright header. As per coding guidelines, "Include NVIDIA copyright header on ALL new files and update year on modified files."
🤖 Fix all issues with AI agents
In `@tensorrt_llm/_torch/auto_deploy/compile/backends/torch_cudagraph.py`:
- Around line 510-524: The piecewise branch can be created without being
captured which leads to replaying a None graph; update the compile path in
torch_cudagraph.py (where PiecewiseCapturedGraph is created and
warmup_and_capture is conditionally called) to validate that
piecewise.cuda_graphs (or equivalent captured graph state on
PiecewiseCapturedGraph) is non-empty before returning DualModeCapturedGraph, and
if capture was skipped (get_mixed_args_kwargs_for_compile is None or
piecewise_num_tokens empty) either raise a clear error or set piecewise to None
so DualModeCapturedGraph is only constructed with a valid captured piecewise
graph; alternatively, implement a guard in
DualModeCapturedGraph.forward/ADPiecewiseRunner.forward to detect missing
captured graphs and fall back to eager execution when piecewise.cuda_graphs is
None/empty.
In `@tensorrt_llm/_torch/auto_deploy/compile/piecewise_runner.py`:
- Around line 238-240: The fallback currently silently proceeds to call
static_inp.copy_(new_inp, non_blocking=True) when shapes are incompatible;
instead, detect the mismatch and raise a clear RuntimeError (or ValueError) that
includes the shapes of static_inp and new_inp and the context (e.g., tensor
names or index) so callers see a helpful message; modify the branch in
piecewise_runner.py (where static_inp and new_inp are compared) to construct and
raise that explicit error (or log the shapes before raising) rather than
performing the blind copy.
- Around line 1-2: This file is missing the required NVIDIA copyright header
with the Apache License 2.0; add the standard NVIDIA copyright block (including
the year of latest meaningful modification) at the top of
tensorrt_llm/_torch/auto_deploy/compile/piecewise_runner.py before the module
docstring, ensuring the exact Apache-2.0 wording and license URL are included
and preserving the existing module docstring (ADPiecewiseRunner...) after the
header.
- Around line 182-196: Remove the unused entry parameter from the
_identify_dynamic_indices method signature and update its only caller to pass
just flat_args; specifically, change def _identify_dynamic_indices(self, entry:
SegmentEntry, flat_args: List[Any]) to def _identify_dynamic_indices(self,
flat_args: List[Any]) and update the call site that currently passes (entry,
flat_args) to call _identify_dynamic_indices(flat_args) so the method no longer
accepts or references entry and the argument list matches.
- Around line 338-347: Add defensive checks before replay: ensure
entry.cuda_graph is not None and entry.dynamic_indices is not None (used by
_prepare_replay_inputs) and raise a clear RuntimeError or fallback to a
non-cuda-graph path if either is missing. Specifically, in the code path around
_prepare_replay_inputs(entry, flat_args) and entry.cuda_graph.replay(), check
entry.dynamic_indices and entry.cuda_graph and if absent log/raise a message
like "cuda_graph capture missing for entry X; cannot replay" or call the
fallback execution path that computes and returns entry.static_output without
replay. Update callers to rely on the explicit error/fallback so failures are
debuggable.
In `@tensorrt_llm/_torch/auto_deploy/compile/piecewise_utils.py`:
- Around line 1-7: Add the standard NVIDIA Apache-2.0 copyright header to the
top of the new module tensorrt_llm._torch.auto_deploy.compile.piecewise_utils
(before the existing module docstring); include the year of latest meaningful
modification and the full Apache-2.0 boilerplate used across the repo so the
file has the required legal header. Ensure the header precedes the triple-quoted
module docstring and matches the format used in other files in the codebase.
- Around line 138-151: Add a unit test that builds a GraphModule whose first
non-placeholder node is a dynamic op to ensure partitioning/classification works
for that edge case; construct a small model or torch.fx graph where
is_dynamic_cached_op returns true for the first real node and then run the same
partitioning logic that populates node_to_partition, partition_counter,
dynamic_partitions and partition_ids_in_order (the mapping code in
piecewise_utils.py) and assert the first dynamic partition is recorded and
subsequent static ops get the correct partition ids. Target the partitioning
code paths that update partition_counter[0], node_to_partition,
dynamic_partitions and the partition_ids_in_order guard to validate behavior
when a dynamic op appears first.
In `@tensorrt_llm/_torch/auto_deploy/config/default.yaml`:
- Around line 206-207: Default piecewise settings are unsafe: change
piecewise_enabled to false and piecewise_num_tokens to null in default.yaml to
avoid silently ignored configs, and add explicit validation in the compiler
layer (e.g., CompilerBackend.compile or CompilerBackend.__init__) to raise or
log an error when piecewise_enabled=true is passed to incompatible backends
(e.g., TorchCompileCompiler, TorchSimpleCompiler or any backend whose name
contains "compile"/"simple"); also replace the hardcoded small buckets [8,16]
with null (auto-generate power-of-two buckets starting at 64) or
document/justify nonstandard buckets if you keep them.
In `@tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py`:
- Around line 549-555: _user-specified piecewise bucket sizes can exceed the
InputBuffer's max_num_tokens causing _shape_for_forward to slice out-of-bounds;
when processing user config (in compile_model.py) validate and clamp or reject
entries in piecewise_bucket_sizes against max_num_tokens before assigning to the
model so effective_num_tokens never exceeds InputBuffer size. Specifically, when
you parse piecewise_bucket_sizes in compile_model.py, iterate each bucket value
and ensure 0 < bucket <= max_num_tokens (either clamp to max_num_tokens or raise
a config error), then pass the validated list to the model (so attributes used
by _shape_for_forward and methods referencing _padded_num_tokens /
total_num_tokens remain safe).
In `@tensorrt_llm/_torch/auto_deploy/transform/library/compile_model.py`:
- Around line 127-136: The code currently only filters piecewise_num_tokens < 3
but not buckets larger than the model max causing out-of-bounds in
_shape_for_forward; update the logic that builds valid_buckets (and dropped) to
also remove values > cm.info.max_num_tokens (use cm.info.max_num_tokens to
compute the upper bound), log which buckets were dropped for being too large
(include max value context), and assign the filtered list to
config_overrides["piecewise_num_tokens"] so piecewise_num_tokens from
self.config cannot exceed cm.info.max_num_tokens.
🧹 Nitpick comments (8)
tensorrt_llm/_torch/auto_deploy/custom_ops/attention/flashinfer_attention.py (1)
456-470:torch.zerosallocation in branch 1 is not CUDA-graph-capture safe.Line 463 allocates a new tensor via
torch.zeros(...). If this path is ever reached during CUDA graph capture (or warm-up), the dynamic allocation will not be captured correctly, leading to silent failures or errors on replay.The
elifbranch (line 468) correctly uses in-placezero_(), which is safe. If the first branch is also expected to execute under CUDA graph capture, consider pre-allocating the padded output buffer (e.g., always allocateyat sizebsupfront) and filling it, rather than conditionally allocating a new tensor.If the invariant is that
y.shape[0] == bsalways holds during capture (because inputs are bucketed), a defensive assertion would make this contract explicit:Suggested assertion
bs = b * s + if torch.cuda.is_current_stream_capturing() or cuda_graph_state.in_warm_up(): + assert y.shape[0] >= bs, ( + f"During CG capture y.shape[0]={y.shape[0]} must be >= bs={bs}" + ) if y.shape[0] < bs:tensorrt_llm/_torch/auto_deploy/compile/piecewise_utils.py (1)
91-98: Substring matching is fragile — prefer exact match after suffix stripping.The comment on line 91 says "Strip the
.defaultsuffix if present for matching," but the code doesn't actually strip it. Instead it usesdyn_op in op_name, which is a loose substring check that could match unintended ops (e.g.,auto_deploy::torch_cached_ssmwould also match a hypotheticalauto_deploy::torch_cached_ssm_v2).Consider stripping
.default(or any overload suffix) and doing an exact match:Proposed fix
# Strip the ".default" suffix if present for matching dynamic_ops = _get_all_dynamic_op_names() - # Check with namespace::name format - for dyn_op in dynamic_ops: - if dyn_op in op_name: - return True - - return False + # Normalize: strip overload suffix (e.g. ".default") for exact matching + base_name = op_name.rsplit(".", 1)[0] if "." in op_name else op_name + return base_name in dynamic_opstensorrt_llm/_torch/auto_deploy/compile/piecewise_runner.py (2)
41-65:_warmup_data_ptrsinSegmentEntryand_track_warmup_ptrsare dead code.The module docstring (lines 21–26) explains that data-ptr change detection during warmup is unreliable and not used.
_identify_dynamic_indicescorrectly ignores_warmup_data_ptrsand instead classifies all non-weight tensors as dynamic. However,_warmup_data_ptrsis still defined inSegmentEntryand_track_warmup_ptrsis still called during warmup — both are dead code paths that serve no purpose and may confuse future readers.Consider removing
_warmup_data_ptrsfromSegmentEntryand_track_warmup_ptrsfromADPiecewiseRunner, or add a comment indicating these are retained intentionally for future use/debugging.
82-94: Mutable class-level attributes should useClassVarand phase should be validated with an enum/literal type.The static analysis tool (Ruff RUF012) correctly flags
_static_output_registryas a mutable class attribute withoutClassVar. Additionally,_current_num_tokensand_current_phaseare shared class-level state that should also be annotated withClassVarfor clarity.Proposed fix
+from typing import ClassVar ... - _current_num_tokens: Optional[int] = None - _current_phase: str = "replay" - _static_output_registry: Dict[Tuple[int, int], torch.Tensor] = {} + _current_num_tokens: ClassVar[Optional[int]] = None + _current_phase: ClassVar[str] = "replay" + _static_output_registry: ClassVar[Dict[Tuple[int, int], torch.Tensor]] = {}tensorrt_llm/_torch/auto_deploy/compile/backends/torch_cudagraph.py (4)
12-12: Remove unusednoqadirective.Ruff (RUF100) reports the
# noqa: I001directive onimport copyis unnecessary since theI001rule is not enabled.-import copy # noqa: I001 +import copy
28-29: Imports deviate from the project'sfrom package.subpackage import moduleguideline.The coding guidelines state: "Python imports must use
from package.subpackage import modulestyle; never usefrom module import Class." These lines import classes/functions directly from modules (ADPiecewiseRunnerfrompiecewise_runner,SplitInfoandsplit_graph_at_dynamic_opsfrompiecewise_utils). Strictly following the guideline would be:from .. import piecewise_runner from .. import piecewise_utilsand then reference
piecewise_runner.ADPiecewiseRunner,piecewise_utils.SplitInfo, etc. That said, the existing code in this file already usesfrom torch.cuda import CUDAGraph, so this is a pre-existing pattern. Flagging for awareness. As per coding guidelines: "Python imports must usefrom package.subpackage import modulestyle; never usefrom module import Class."
412-432:_is_decode_onlyfallback heuristic is fragile.The fallback (lines 426–429) checks
v.shape[1] == 1on 2D+ tensors to detect decode-only batches. This assumes a specific tensor layout ([batch, seq_len, ...]). If a model uses different shapes or the first matching kwarg doesn't follow this convention, misclassification could route traffic to the wrong execution path.Consider adding a debug-level log when the fallback heuristic is triggered, so misrouting issues are diagnosable:
Proposed fix
# Fallback heuristic: check if first batched input has sequence dim == 1 # (decode = 1 token per sequence) for name in self.batched_input_names: v = kwargs.get(name) if v is not None and isinstance(v, torch.Tensor) and v.ndim >= 2: + ad_logger.debug( + f"DualModeCapturedGraph: using fallback heuristic on '{name}' " + f"(shape={v.shape}) to determine decode-only status" + ) return v.shape[1] == 1
53-73:_submod_has_cuda_opsconservatively returnsTruefor all non-FX modules — consider logging.This is fine for correctness, but when debugging why a trivial submodule was wrapped, the conservative default could be confusing. A debug log when the fallback triggers would help:
if not isinstance(submod, GraphModule): + ad_logger.debug(f"_submod_has_cuda_ops: non-FX module assumed to have CUDA ops") return True # Conservative: non-FX modules assumed to have GPU ops
suyoggupta
reviewed
Feb 23, 2026
Collaborator
|
can we update the PR with TTFT numbers for GLM (h100 and b200) |
lucaslie
suggested changes
Feb 24, 2026
Contributor
There was a problem hiding this comment.
I don't think SequenceInfo is the right place to handle padding logic. Please remove the padding logic from there. Instead, I suggest the following approach:
Edit: [A: Improvement for current approach]
- Make use of padding wrapper we have already for _prepare_inputs
- In _prepare_inputs set batch_info by not counting the padding tokens. This way the attention operators won't process the padded tokens since all of them read batch_info_host to split the batch.
Edit: [B: Ideal architecture]
Ideally, we can also address the clean architecture with handling padding/truncation in the piecewise cg infrastructure. This removes all hacks for this PR across all attention backends, ad_executor, and SequenceInfo to 0. From a first principle, this PR should at most touch two files (torch_cudagraph compile backend and the compile_model transform)
Edit: per discussion with @nvchenghaoz, he will spend 1-2 days on B to see if it can be done as part of this PR or otherwise just incorporate the feedback A for the current approach
lucaslie
reviewed
Feb 26, 2026
Signed-off-by: Chenghao Zhang <211069071+nvchenghaoz@users.noreply.github.com>
Signed-off-by: Chenghao Zhang <211069071+nvchenghaoz@users.noreply.github.com>
Signed-off-by: Chenghao Zhang <211069071+nvchenghaoz@users.noreply.github.com>
Signed-off-by: Chenghao Zhang <211069071+nvchenghaoz@users.noreply.github.com>
Signed-off-by: Chenghao Zhang <211069071+nvchenghaoz@users.noreply.github.com>
Signed-off-by: Chenghao Zhang <211069071+nvchenghaoz@users.noreply.github.com>
Signed-off-by: Chenghao Zhang <211069071+nvchenghaoz@users.noreply.github.com>
Signed-off-by: Chenghao Zhang <211069071+nvchenghaoz@users.noreply.github.com>
Signed-off-by: nvchenghaoz <211069071+nvchenghaoz@users.noreply.github.com>
Signed-off-by: Chenghao Zhang <211069071+nvchenghaoz@users.noreply.github.com>
Signed-off-by: Chenghao Zhang <211069071+nvchenghaoz@users.noreply.github.com>
nvchenghaoz
force-pushed
the
chenghao/piecewise_cudagraph_0211
branch
from
February 26, 2026 04:54
70a5786 to
67ebed8
Compare
Collaborator
Author
|
/bot run |
nvchenghaoz
changed the title
Collaborator
|
PR_Github #36968 [ run ] triggered by Bot. Commit: |
Collaborator
Author
|
/bot run --extra-stage "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1" |
Collaborator
|
PR_Github #37746 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #37746 [ run ] completed with state
|
Collaborator
Author
|
/bot run --extra-stage "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1" |
Collaborator
|
PR_Github #37769 [ run ] triggered by Bot. Commit: |
Collaborator
Author
|
/bot run --extra-stage "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1" |
Collaborator
|
PR_Github #37780 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #37769 [ run ] completed with state
|
Collaborator
|
PR_Github #37780 [ run ] completed with state
|
Signed-off-by: Chenghao Zhang <211069071+nvchenghaoz@users.noreply.github.com>
Collaborator
Author
|
/bot run --extra-stage "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1" |
1 similar comment
Collaborator
Author
|
/bot run --extra-stage "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1" |
Collaborator
|
PR_Github #37894 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #37896 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #37894 [ run ] completed with state |
Collaborator
|
PR_Github #37896 [ run ] completed with state
|
Collaborator
Author
|
/bot run --extra-stage "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1" --disable-fail-fast |
Collaborator
|
PR_Github #37913 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #37913 [ run ] completed with state
|
Collaborator
Author
|
/bot run --extra-stage "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1" --disable-fail-fast |
Collaborator
|
PR_Github #37942 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #37942 [ run ] completed with state
|
Collaborator
Author
|
/bot run --extra-stage "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1" --disable-fail-fast |
Collaborator
|
PR_Github #37975 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #37975 [ run ] completed with state |
dominicshanshan
pushed a commit
to dominicshanshan/TensorRT-LLM
that referenced
this pull request
Mar 9, 2026
NVIDIA#11515) Signed-off-by: Chenghao Zhang <211069071+nvchenghaoz@users.noreply.github.com> Signed-off-by: nvchenghaoz <211069071+nvchenghaoz@users.noreply.github.com>
tianyuz-nv
pushed a commit
to wanqian-nv/TensorRT-LLM
that referenced
this pull request
Mar 19, 2026
NVIDIA#11515) Signed-off-by: Chenghao Zhang <211069071+nvchenghaoz@users.noreply.github.com> Signed-off-by: nvchenghaoz <211069071+nvchenghaoz@users.noreply.github.com>
limin2021
pushed a commit
to limin2021/TensorRT-LLM
that referenced
this pull request
Mar 19, 2026
NVIDIA#11515) Signed-off-by: Chenghao Zhang <211069071+nvchenghaoz@users.noreply.github.com> Signed-off-by: nvchenghaoz <211069071+nvchenghaoz@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
5 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
The PR is to enable the piecewise cudagraph support to improve the TTFT. Some perf numbers from my local test TTFT -
Qwen 3.5 35B 1k/1k TP2 NVFP4
Qwen 3.5 35B 1k/1k TP2 BF16
GLM 4.7 Flash 1k/1k BF16
GLM 4.7 Flash 1k/1k NVFP4
Nemotron Nano V3
Remaining Piecewise CG work:
Summary by CodeRabbit
New Features
Bug Fixes