[TRTLLM-13579][feat] BREAKING: Support BCG in Prefill - #16609
[TRTLLM-13579][feat] BREAKING: Support BCG in Prefill#16609GuanhuaWang2001 wants to merge 13 commits into
Conversation
|
|
||
| if self.prefill_cuda_graph_backend == PrefillCudaGraphBackend.PIECEWISE: | ||
| if self.torch_compile_config is None: | ||
| self.torch_compile_config = TorchCompileConfig() |
There was a problem hiding this comment.
enable_piecewise_cuda_graph is default off. And capture_num_tokens should be setup.
Besides, we would need to set max_num_streams to maybe 3 to ensure the best performance.
| @pytest.mark.threadleak(enabled=False) | ||
| def test_bf16_breakable_prefill_cuda_graph(self): | ||
| model_path = f"{llm_models_root()}/Qwen3.5-4B" | ||
| prompts = [ |
There was a problem hiding this comment.
Maybe directly test with GSM8K and MMLU.
| def _weak_ref_if_tensor(value: Any) -> Any: | ||
| if torch.is_tensor(value): | ||
| return make_weak_ref(value) | ||
| if isinstance(value, tuple): |
There was a problem hiding this comment.
I think make_weak_ref already handle the tuple/list/...
| return hidden_states | ||
|
|
||
| self.layer_model.forward = capture_forward | ||
| self.logits_processor.forward = passthrough_forward |
There was a problem hiding this comment.
Wondering why we need to handle logits_processor differently. That might be a blocker when enabling spec dec support.
0d2ecb0 to
eeae18c
Compare
| o_lora.transpose(0, 1), | ||
| ) | ||
| return self.o_b_proj(o_lora.flatten(1)) | ||
| if attn_out_latent.shape[1:] == (self.n_local_groups, self.o_lora_rank): |
There was a problem hiding this comment.
Directly pass in enable_dsv4_epilogue_fusion is more clear.
| raise RuntimeError( | ||
| "DSv4 epilogue fusion requires a context-only or generation-only batch." | ||
| ) | ||
| enable_dsv4_epilogue_fusion = output.ndim == 3 |
There was a problem hiding this comment.
Same here, explicitly pass in enable_dsv4_epilogue_fusion
| return source | ||
|
|
||
|
|
||
| def eager_on_graph(enable: bool) -> Callable[[Callable], Callable]: |
There was a problem hiding this comment.
What is the meaning of eager_on_graph(False) comparing without using the decorator?
|
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:
WalkthroughThis PR adds prefill CUDA graph backends, segmented breakable capture and replay, PyTorch engine integration, DeepSeek-V4 MLA epilogue-fusion changes, breakable custom-op dispatch, and related unit, integration, and API-stability coverage. ChangesPrefill configuration and compatibility
Breakable CUDA graph runtime
Prefill engine integration
Custom operations and MLA epilogue fusion
Functionalization mutation validation
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant TorchLlmArgs
participant PyTorchModelEngine
participant BreakableCUDAGraphRunner
participant ModelBody
TorchLlmArgs->>PyTorchModelEngine: configure prefill backend and token buckets
PyTorchModelEngine->>BreakableCUDAGraphRunner: warm up and capture buckets
BreakableCUDAGraphRunner->>ModelBody: capture segmented prefill execution
PyTorchModelEngine->>BreakableCUDAGraphRunner: replay matching token bucket
BreakableCUDAGraphRunner->>ModelBody: execute captured segments and eager breaks
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 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.
Actionable comments posted: 9
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
docs/source/features/torch_compile_and_piecewise_cuda_graph.md (1)
9-9: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winStale TOC anchor after the heading rename.
Line 1 is now "Torch Compile & Prefill CUDA Graph", so the anchor
#torch-compile--piecewise-cuda-graphno longer resolves.📝 Proposed fix
-- [Torch Compile & Piecewise CUDA Graph](`#torch-compile--piecewise-cuda-graph`) +- [Torch Compile & Prefill CUDA Graph](`#torch-compile--prefill-cuda-graph`)🤖 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 `@docs/source/features/torch_compile_and_piecewise_cuda_graph.md` at line 9, Update the table-of-contents link in the “Torch Compile & Prefill CUDA Graph” documentation to use the heading’s current generated anchor, replacing the stale piecewise-cuda-graph fragment while preserving the existing link target.
🧹 Nitpick comments (9)
tests/unittest/_torch/compilation/test_remove_copy_pass.py (1)
25-87: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winCover the new malformed optional-output path.
Neither test sets a write argument’s
_<arg>_base_indextoNoneand then requests its getitem output, so the new assertion inremove_copy_for_mutates_args()is untested. Add a focusedpytest.raises(AssertionError, match="graph is malformed")case.Test coverage summary
- Added:
test_remove_copy_for_mutates_args_auto_functionalized_v2;test_remove_copy_for_mla_restores_final_output_mutation.- Test-list registration: no
tests/integration/test_lists/file was supplied; not assessable from this review context.- Verdict: insufficient — the newly added
None-base invariant lacks coverage.🤖 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_remove_copy_pass.py` around lines 25 - 87, Add a focused pytest.raises(AssertionError, match="graph is malformed") test alongside test_remove_copy_for_mutates_args_auto_functionalized_v2 that constructs an auto_functionalized_v2 graph with a write argument’s _<arg>_base_index set to None while requesting its getitem output, then calls remove_copy_pass.remove_copy_for_mutates_args(graph).Source: Path instructions
tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/cuda_utils.py (1)
17-25: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winType the CUDA result helper contract.
check_cuda_errors(result)has no parameter or return annotation despite being the central adapter for cuda-python result tuples. Add overloads or a generic tuple contract so callers cannot accidentally pass an incompatible result shape.As per coding guidelines, “Annotate every function” and avoid imprecise types.
🤖 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/breakable_cuda_graph/cuda_utils.py` around lines 17 - 25, Annotate the check_cuda_errors function with precise parameter and return types matching cuda-python result tuples, using overloads or a generic tuple contract to represent one-element, two-element, and multi-element results. Ensure the annotations reject incompatible result shapes while preserving the existing None, single-value, and tuple return behavior.Source: Coding guidelines
tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/breakable_cuda_graph.py (1)
146-185: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winPreserve decorated callable types.
eager_on_graph()exposes bareCallable, whilewrapper, the weak-ref helper, andreplay_fnlose argument and return types throughAny. UseParamSpecandTypeVarso decorated custom operations retain their callable contract.As per coding guidelines, “Annotate every function” and “use precise
Callabletypes.” Based on learnings, this repository supports Python 3.10+ typing features.🤖 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/breakable_cuda_graph/breakable_cuda_graph.py` around lines 146 - 185, Update eager_on_graph and its nested functions to preserve callable signatures using ParamSpec and TypeVar: type the decorator as accepting and returning Callable[P, R], annotate wrapper with P/R, and annotate replay_fn and make_weak_ref_with_str_none with precise argument and return types instead of bare Callable or Any. Preserve the existing runtime behavior while applying Python 3.10+ typing features throughout the added functions.Sources: Coding guidelines, Learnings
tensorrt_llm/_torch/modules/attention.py (1)
932-944: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winSplit the compile and BCG custom-op paths
tensorrt_llm/_torch/modules/attention.py#L932-L944,tensorrt_llm/_torch/models/modeling_minimaxm3.py#L1102-L1113, andtensorrt_llm/_torch/modules/mla.py#L3069-L3079all route through theeager_on_graph(True)wrapper for bothtorch.compileand BCG. Followtensorrt_llm/_torch/modules/mamba/gdn_mixer.py#L996-L1006and use the raw custom op on the compile path, reserving the wrapper for the breakable-CUDA-graph case.🤖 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/modules/attention.py` around lines 932 - 944, Split the custom-op handling for torch.compile and breakable CUDA graph paths, using the raw custom op during compilation and reserving the eager_on_graph(True) wrapper for BCG execution. Apply this consistently at tensorrt_llm/_torch/modules/attention.py:932-944, tensorrt_llm/_torch/models/modeling_minimaxm3.py:1102-1113, and tensorrt_llm/_torch/modules/mla.py:3069-3079, following the existing pattern in gdn_mixer.py.tensorrt_llm/_torch/pyexecutor/model_engine.py (2)
6194-6208: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueCapture path returns early, skipping post-forward hooks.
return breakable_runner.capture_model_body(forward_step)bypassesself.forward_pass_callable()and_execute_logit_post_processors(...)that every other branch runs. That is presumably intentional for warmup-only capture, but it is not obvious from the code — a short comment stating that capture never produces user-visible logits would prevent a future regression here.🤖 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 6194 - 6208, Add a concise comment immediately before the early return from breakable_runner.capture_model_body(forward_step) explaining that the capture path is warmup-only and produces no user-visible logits, so it intentionally bypasses self.forward_pass_callable() and _execute_logit_post_processors(...).
775-784: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRunner construction rejects wrapped models with an opaque message.
drafting_loop_wrapperreplacesself.modelwith aBaseDraftingLoopWrapper, soisinstance(decoder_model, DecoderModelForCausalLM)fails and the user sees "requires a decoder model body" rather than "speculative decoding is unsupported by the breakable prefill backend". Consider naming the unsupported configuration in the error.🤖 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 775 - 784, Update the validation in the breakable prefill runner initialization around decoder_model to recognize models wrapped by drafting_loop_wrapper and reject that unsupported speculative-decoding configuration with an explicit error message stating that speculative decoding is unsupported by the breakable prefill backend.docs/source/features/torch_compile_and_piecewise_cuda_graph.md (1)
44-66: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueClarify the breakable support matrix sentence.
"supports BF16 Qwen3.5 on one GPU for context-only, tensor/pipeline parallelism and mixed context/decode batches" reads self-contradictory (one GPU vs TP/PP). Consider splitting into an explicit supported/unsupported list.
Also applies to: 80-82
🤖 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 `@docs/source/features/torch_compile_and_piecewise_cuda_graph.md` around lines 44 - 66, Clarify the breakable backend support statement by separating supported and unsupported configurations, avoiding the contradiction between “one GPU” and tensor/pipeline parallelism. Explicitly state the supported BF16 Qwen3.5 context-only and mixed context/decode cases, identify the supported parallelism scope, and list unsupported configurations separately; update the corresponding repeated statement as well.tests/unittest/llmapi/test_llm_args.py (1)
1645-1685: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚖️ Poor tradeoffPadding tests rely on a hand-built
PyTorchModelEngineviaobject.__new__.This works today but silently breaks whenever
_get_padding_paramsstarts reading another attribute. Consider extracting the padding decision into a small free function (or a@staticmethodtaking the bucket list/backend) so the tests can call it without constructing a partially-initialized engine.🤖 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/llmapi/test_llm_args.py` around lines 1645 - 1685, Extract the padding decision logic used by PyTorchModelEngine._get_padding_params into a standalone helper or static method that accepts the padding bucket list and backend explicitly. Update _get_padding_params and the tests test_piecewise_and_breakable_use_identical_padding and test_attention_dp_prefill_graph_uses_all_rank_decision to use this helper, removing reliance on object.__new__(PyTorchModelEngine) for padding-only behavior.tests/integration/defs/accuracy/test_llm_api_pytorch.py (1)
3975-4084: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winTrim the benchmark machinery from this correctness test.
The only assertions are the three determinism/equality checks at the end, yet the test builds two 8-GPU engines, runs 12 generation passes, and computes latency/throughput medians and "p90".
statistics.quantiles(latencies, n=10, method="inclusive")[8]over five samples is not a meaningful p90, and the JSON dump toTRTLLM_DSV4_BCG_AB_RESULT_PATHis ad-hoc reporting rather than test output. Dropping the timing/stats/JSON block and reducing the round count would cut CI cost substantially without weakening the assertions; keep the perf comparison in a dedicated benchmark if it is still needed.🤖 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/integration/defs/accuracy/test_llm_api_pytorch.py` around lines 3975 - 4084, Simplify run_variant and the surrounding result/reporting flow to make this a correctness-only test: remove timing, throughput, statistics, perf fields, console/result-file JSON reporting, and retain only generated token IDs for determinism and cross-variant equality checks. Reduce the repeated generation rounds to the minimum needed to validate repeatability, while preserving disabled_repeatable, fusion_repeatable, and token_ids_match assertions.
🤖 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.
Inline comments:
In `@tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph_runner.py`:
- Around line 106-111: Update the exception cleanup in the capture flow to clear
`_shared_output` whenever `created_memory_pool` is true and the newly created
pool is discarded because `self._graphs` is empty. Keep the existing graph reset
and pool cleanup behavior unchanged, ensuring a later `capture()` allocates a
fresh shared output.
In `@tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/breakable_cuda_graph.py`:
- Around line 248-256: Make BreakableCUDAGraphCapture.__enter__ transactional:
wrap stream/context setup, ContextVar token installation, and _begin_new_segment
(which invokes capture_begin) in an except BaseException cleanup path. On
failure, undo every resource acquired so far—including resetting ContextVar
tokens, exiting the stream context, and removing the _install_wait_stream_hook
monkeypatch—then re-raise the original exception; keep normal successful entry
behavior unchanged.
In `@tensorrt_llm/_torch/pyexecutor/model_engine.py`:
- Around line 1997-2012: Guard the prefill capture block around
breakable_cuda_graph_runner.capture and the torch.compile warmup loop with
_assert_all_tp_ranks_have_warmup_batch(...). Invoke the assertion before
rank-local capture decisions so all TP ranks consistently skip or perform
capture when batch is unavailable, preventing missing BCG entries and replay
KeyErrors.
In `@tensorrt_llm/_torch/utils.py`:
- Around line 383-390: Update the stale flag call in
_apply_steady_gen_fast_prepare to use
set_per_request_prefill_cuda_graph_flag(...) instead of
set_per_request_piecewise_cuda_graph_flag(...). Preserve the existing argument
and fast-path behavior, ensuring the called setter matches the imported symbol
and avoids the NameError.
In `@tensorrt_llm/llmapi/llm_args.py`:
- Around line 5212-5226: Update the legacy bucket handling in the TorchLlmArgs
initialization flow to read capture_num_tokens only when that field was
explicitly set by the user, rather than when
TorchCompileConfig.set_default_capture_num_tokens auto-populated it. Preserve
the conflict validation for explicitly provided legacy buckets and the migration
warning/assignment for non-explicit prefill_capture_num_tokens.
In `@tests/integration/defs/accuracy/test_llm_api_pytorch.py`:
- Around line 6306-6342: Add the new Qwen test to the appropriate single-GPU
test-db list, and register the DeepSeekV4Flash epilogue-fusion A/B test in the
corresponding multi-GPU or QA list. Use the exact pytest node for
TestQwen3_5_4B::test_bf16_breakable_prefill_cuda_graph and the matching node for
the DeepSeekV4Flash test; update the relevant files under
tests/integration/test_lists/test-db/ or tests/integration/test_lists/qa/.
In `@tests/unittest/_torch/executor/test_breakable_cuda_graph.py`:
- Around line 5-6: Remove the unused gc and weakref imports from
test_breakable_cuda_graph.py, leaving the remaining imports and test behavior
unchanged.
In `@tests/unittest/_torch/modules/test_mla_registry.py`:
- Around line 87-112: Update
test_dsv4_epilogue_fusion_returns_only_final_output_inside_breakable_graph so it
does not patch is_in_breakable_cuda_graph, since create_mla_outputs_impl does
not consult it. Remove the unused context patch and rename the test to describe
only the behavior it actually verifies, or otherwise add a meaningful comparison
that distinguishes breakable and non-breakable contexts.
In `@tests/unittest/llmapi/test_llm_args.py`:
- Around line 1632-1643: Update the remaining later test cases in
test_prefill_filter_sorts_dedupes_and_drops_nonpositive and related tests to
import and call _filter_prefill_capture_num_tokens instead of
_filter_piecewise_capture_num_tokens. Use the defined model_engine symbol
consistently and do not add a compatibility alias.
---
Outside diff comments:
In `@docs/source/features/torch_compile_and_piecewise_cuda_graph.md`:
- Line 9: Update the table-of-contents link in the “Torch Compile & Prefill CUDA
Graph” documentation to use the heading’s current generated anchor, replacing
the stale piecewise-cuda-graph fragment while preserving the existing link
target.
---
Nitpick comments:
In `@docs/source/features/torch_compile_and_piecewise_cuda_graph.md`:
- Around line 44-66: Clarify the breakable backend support statement by
separating supported and unsupported configurations, avoiding the contradiction
between “one GPU” and tensor/pipeline parallelism. Explicitly state the
supported BF16 Qwen3.5 context-only and mixed context/decode cases, identify the
supported parallelism scope, and list unsupported configurations separately;
update the corresponding repeated statement as well.
In `@tensorrt_llm/_torch/modules/attention.py`:
- Around line 932-944: Split the custom-op handling for torch.compile and
breakable CUDA graph paths, using the raw custom op during compilation and
reserving the eager_on_graph(True) wrapper for BCG execution. Apply this
consistently at tensorrt_llm/_torch/modules/attention.py:932-944,
tensorrt_llm/_torch/models/modeling_minimaxm3.py:1102-1113, and
tensorrt_llm/_torch/modules/mla.py:3069-3079, following the existing pattern in
gdn_mixer.py.
In `@tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/breakable_cuda_graph.py`:
- Around line 146-185: Update eager_on_graph and its nested functions to
preserve callable signatures using ParamSpec and TypeVar: type the decorator as
accepting and returning Callable[P, R], annotate wrapper with P/R, and annotate
replay_fn and make_weak_ref_with_str_none with precise argument and return types
instead of bare Callable or Any. Preserve the existing runtime behavior while
applying Python 3.10+ typing features throughout the added functions.
In `@tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/cuda_utils.py`:
- Around line 17-25: Annotate the check_cuda_errors function with precise
parameter and return types matching cuda-python result tuples, using overloads
or a generic tuple contract to represent one-element, two-element, and
multi-element results. Ensure the annotations reject incompatible result shapes
while preserving the existing None, single-value, and tuple return behavior.
In `@tensorrt_llm/_torch/pyexecutor/model_engine.py`:
- Around line 6194-6208: Add a concise comment immediately before the early
return from breakable_runner.capture_model_body(forward_step) explaining that
the capture path is warmup-only and produces no user-visible logits, so it
intentionally bypasses self.forward_pass_callable() and
_execute_logit_post_processors(...).
- Around line 775-784: Update the validation in the breakable prefill runner
initialization around decoder_model to recognize models wrapped by
drafting_loop_wrapper and reject that unsupported speculative-decoding
configuration with an explicit error message stating that speculative decoding
is unsupported by the breakable prefill backend.
In `@tests/integration/defs/accuracy/test_llm_api_pytorch.py`:
- Around line 3975-4084: Simplify run_variant and the surrounding
result/reporting flow to make this a correctness-only test: remove timing,
throughput, statistics, perf fields, console/result-file JSON reporting, and
retain only generated token IDs for determinism and cross-variant equality
checks. Reduce the repeated generation rounds to the minimum needed to validate
repeatability, while preserving disabled_repeatable, fusion_repeatable, and
token_ids_match assertions.
In `@tests/unittest/_torch/compilation/test_remove_copy_pass.py`:
- Around line 25-87: Add a focused pytest.raises(AssertionError, match="graph is
malformed") test alongside
test_remove_copy_for_mutates_args_auto_functionalized_v2 that constructs an
auto_functionalized_v2 graph with a write argument’s _<arg>_base_index set to
None while requesting its getitem output, then calls
remove_copy_pass.remove_copy_for_mutates_args(graph).
In `@tests/unittest/llmapi/test_llm_args.py`:
- Around line 1645-1685: Extract the padding decision logic used by
PyTorchModelEngine._get_padding_params into a standalone helper or static method
that accepts the padding bucket list and backend explicitly. Update
_get_padding_params and the tests
test_piecewise_and_breakable_use_identical_padding and
test_attention_dp_prefill_graph_uses_all_rank_decision to use this helper,
removing reliance on object.__new__(PyTorchModelEngine) for padding-only
behavior.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 93727c2d-fa02-4e5c-bb81-bebfd79a85ef
📒 Files selected for processing (25)
docs/source/features/torch_compile_and_piecewise_cuda_graph.mdtensorrt_llm/_torch/compilation/piecewise_optimizer.pytensorrt_llm/_torch/compilation/remove_copy_pass.pytensorrt_llm/_torch/models/modeling_minimaxm3.pytensorrt_llm/_torch/modules/attention.pytensorrt_llm/_torch/modules/mamba/gdn_mixer.pytensorrt_llm/_torch/modules/mla.pytensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/__init__.pytensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/breakable_cuda_graph.pytensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/context.pytensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/cuda_utils.pytensorrt_llm/_torch/pyexecutor/breakable_cuda_graph_runner.pytensorrt_llm/_torch/pyexecutor/model_engine.pytensorrt_llm/_torch/utils.pytensorrt_llm/llmapi/__init__.pytensorrt_llm/llmapi/llm_args.pytensorrt_llm/usage/llm_args_golden_manifest.jsontests/integration/defs/accuracy/test_llm_api_pytorch.pytests/integration/test_lists/waives.txttests/unittest/_torch/compilation/test_remove_copy_pass.pytests/unittest/_torch/executor/test_breakable_cuda_graph.pytests/unittest/_torch/modules/test_mla_registry.pytests/unittest/api_stability/api_stability_core.pytests/unittest/api_stability/references/llm.yamltests/unittest/llmapi/test_llm_args.py
💤 Files with no reviewable changes (1)
- tests/integration/test_lists/waives.txt
| except Exception: | ||
| if graph is not None: | ||
| graph.reset() | ||
| if created_memory_pool and not self._graphs: | ||
| self._memory_pool = None | ||
| raise |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
Reset _shared_output when the freshly created pool is discarded.
On a failed first capture the pool handle is dropped but _shared_output still references a tensor allocated inside that (now discarded) pool. A later capture() would then route bucket outputs into that stale buffer instead of allocating a new shared output.
🛠️ Proposed fix
if created_memory_pool and not self._graphs:
self._memory_pool = None
+ self._shared_output = None
raise📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| except Exception: | |
| if graph is not None: | |
| graph.reset() | |
| if created_memory_pool and not self._graphs: | |
| self._memory_pool = None | |
| raise | |
| except Exception: | |
| if graph is not None: | |
| graph.reset() | |
| if created_memory_pool and not self._graphs: | |
| self._memory_pool = None | |
| self._shared_output = None | |
| raise |
🤖 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/breakable_cuda_graph_runner.py` around lines
106 - 111, Update the exception cleanup in the capture flow to clear
`_shared_output` whenever `created_memory_pool` is true and the newly created
pool is discarded because `self._graphs` is empty. Keep the existing graph reset
and pool cleanup behavior unchanged, ensuring a later `capture()` allocates a
fresh shared output.
| def __enter__(self) -> "BreakableCUDAGraphCapture": | ||
| _install_wait_stream_hook() | ||
| if self._stream is not None: | ||
| self._stream_context = torch.cuda.stream(self._stream) | ||
| self._stream_context.__enter__() | ||
| self._capture_token = _current_capture.set(self) | ||
| self._stream_token = _current_stream.set(self._stream or torch.cuda.current_stream()) | ||
| self._forked_token = _forked_streams.set(set()) | ||
| self._begin_new_segment() |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Make capture entry transactional.
If stream setup or capture_begin() on Line 274 raises, __exit__() is never called. The global wait_stream monkeypatch, ContextVar tokens, and possibly the stream context remain installed, corrupting later captures. Clean up each successfully acquired resource in an except BaseException path before re-raising.
🤖 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/breakable_cuda_graph/breakable_cuda_graph.py`
around lines 248 - 256, Make BreakableCUDAGraphCapture.__enter__ transactional:
wrap stream/context setup, ContextVar token installation, and _begin_new_segment
(which invokes capture_begin) in an except BaseException cleanup path. On
failure, undo every resource acquired so far—including resetting ContextVar
tokens, exiting the stream context, and removing the _install_wait_stream_hook
monkeypatch—then re-raise the original exception; keep normal successful entry
behavior unchanged.
| if self.breakable_cuda_graph_runner is not None: | ||
| self.breakable_cuda_graph_runner.capture( | ||
| num_tokens, lambda: self.forward( | ||
| batch, | ||
| new_tensors_device=None, | ||
| resource_manager=resource_manager)) | ||
| else: | ||
| # Run a few times to ensure torch.compile capture. | ||
| for _ in range(4): | ||
| self.forward(batch, | ||
| new_tensors_device=None, | ||
| resource_manager=resource_manager) | ||
|
|
||
| self.forward(batch, | ||
| new_tensors_device=None, | ||
| resource_manager=resource_manager) | ||
| torch.cuda.synchronize() | ||
| gc.collect() | ||
| torch.cuda.empty_cache() |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Check whether prefill capture reconciles per-rank capture success.
ast-grep outline tensorrt_llm/_torch/pyexecutor/model_engine.py --match '_capture_prefill_cuda_graphs' --view expanded
rg -nP -C3 '_assert_all_tp_ranks_have_warmup_batch' tensorrt_llm/_torch/pyexecutor/model_engine.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 2300
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Map the relevant areas of model_engine.py and search for capture/bucket logic.
ast-grep outline tensorrt_llm/_torch/pyexecutor/model_engine.py --view expanded | sed -n '1,260p'
printf '\n--- search hits ---\n'
rg -n -C 3 'breakable_cuda_graph_runner|_prefill_cuda_graph_num_tokens|_get_padding_params|execute\(|capture\(|warmup|bcg|bucket|KeyError' tensorrt_llm/_torch/pyexecutor/model_engine.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 50375
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the relevant slices if the outline output is too large.
sed -n '1880,2075p' tensorrt_llm/_torch/pyexecutor/model_engine.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 11004
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Read the local region around the prefill CUDA graph capture code.
nl -ba tensorrt_llm/_torch/pyexecutor/model_engine.py | sed -n '1940,2055p'Repository: NVIDIA/TensorRT-LLM
Length of output: 197
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the exact prefill capture loop and nearby control flow.
nl -ba tensorrt_llm/_torch/pyexecutor/model_engine.py | sed -n '1970,2035p'
printf '\n--- surrounding prefill symbols ---\n'
rg -n -C 4 'prefill|breakable_cuda_graph_runner|_prefill_cuda_graph_num_tokens|_assert_all_tp_ranks_have_warmup_batch' tensorrt_llm/_torch/pyexecutor/model_engine.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 197
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
p = Path('tensorrt_llm/_torch/pyexecutor/model_engine.py')
lines = p.read_text().splitlines()
for start, end, title in [
(2938, 3038, "_get_padding_params and nearby runtime bucket logic"),
(2078, 2325, "_create_warmup_request and _create_cuda_graph_warmup_request"),
(2828, 2928, "rank/token coordination helpers"),
]:
print(f"\n=== {title} ({start}-{end}) ===")
for i in range(start, min(end, len(lines)) + 1):
print(f"{i:5d}: {lines[i-1]}")
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 25062
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
p = Path('tensorrt_llm/_torch/pyexecutor/model_engine.py')
lines = p.read_text().splitlines()
for start, end, title in [
(2878, 3018, "_get_all_rank_num_tokens, _set_spec_metadata_all_rank_num_tokens, and _get_padding_params"),
(2080, 2325, "_create_warmup_request and _create_cuda_graph_warmup_request"),
]:
print(f"\n=== {title} ({start}-{end}) ===")
for i in range(start, min(end, len(lines)) + 1):
print(f"{i:5d}: {lines[i-1]}")
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 21522
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
paths = [
Path('tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py'),
Path('tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph_runner.py'),
Path('tensorrt_llm/_torch/pyexecutor/model_engine.py'),
]
for p in paths:
print(f"\n### {p}")
if not p.exists():
print("missing")
continue
lines = p.read_text().splitlines()
needles = [
'maybe_get_cuda_graph', 'execute(', 'KeyError', 'capture_num_tokens',
'padding_dummy_requests', 'fallback', 'graph_cache', 'capture(',
'warmup(', 'cuda graph'
]
for needle in needles:
hits = [i + 1 for i, line in enumerate(lines) if needle in line]
if hits:
print(f"{needle}: {hits[:20]}")
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 978
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
for rel in [
'tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py',
'tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph_runner.py',
]:
p = Path(rel)
print(f"\n=== {rel} ===")
lines = p.read_text().splitlines()
for start, end in [(1, 260), (260, 520), (520, 820)]:
if start > len(lines):
continue
print(f"\n--- lines {start}-{min(end, len(lines))} ---")
for i in range(start, min(end, len(lines)) + 1):
line = lines[i-1]
if any(k in line for k in [
'def capture', 'def warmup', 'def execute', 'maybe_get_cuda_graph',
'KeyError', 'capture_num_tokens', 'padding_dummy_requests',
'fallback', 'graph_cache'
]):
print(f"{i:5d}: {line}")
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 1845
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
# Focus on the runtime lookup path around graph execution in the runner(s).
targets = [
('tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py',
['def execute', 'def maybe_get_cuda_graph', 'graph_cache', 'KeyError', 'capture_num_tokens']),
('tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph_runner.py',
['def capture', 'def warmup', 'execute', 'KeyError', 'fallback']),
]
for rel, needles in targets:
p = Path(rel)
print(f"\n=== {rel} ===")
lines = p.read_text().splitlines()
for needle in needles:
matches = [i + 1 for i, line in enumerate(lines) if needle in line]
print(f"{needle}: {matches[:20]}")
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 508
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
for rel in [
'tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py',
'tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph_runner.py',
]:
p = Path(rel)
print(f"\n=== {rel} ===")
lines = p.read_text().splitlines()
for needle in [
'def capture', 'def warmup', 'def execute', 'maybe_get_cuda_graph',
'KeyError', 'capture_num_tokens', 'padding_dummy_requests',
'fallback', 'graph_cache'
]:
hits = [i + 1 for i, line in enumerate(lines) if needle in line]
if hits:
print(f"{needle}: {hits[:20]}")
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 525
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
for rel, ranges in {
'tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py': [(260, 340), (360, 430), (780, 930)],
'tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph_runner.py': [(1, 210)],
}.items():
p = Path(rel)
lines = p.read_text().splitlines()
print(f"\n=== {rel} ===")
for start, end in ranges:
print(f"\n--- lines {start}-{min(end, len(lines))} ---")
for i in range(start, min(end, len(lines)) + 1):
print(f"{i:5d}: {lines[i-1]}")
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 25252
Guard prefill capture with _assert_all_tp_ranks_have_warmup_batch(...).
The prefill capture loops can skip batch is None independently per TP rank, but runtime only allgathers the can-run decision in _get_padding_params. If one rank misses a bucket during warmup, that rank can later miss the BCG entry and fall back to eager or hit KeyError while peers replay.
🤖 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 1997 - 2012,
Guard the prefill capture block around breakable_cuda_graph_runner.capture and
the torch.compile warmup loop with _assert_all_tp_ranks_have_warmup_batch(...).
Invoke the assertion before rank-local capture decisions so all TP ranks
consistently skip or perform capture when batch is unavailable, preventing
missing BCG entries and replay KeyErrors.
| legacy_buckets = (compile_config.capture_num_tokens | ||
| if compile_config is not None else None) | ||
| if legacy_buckets is not None: | ||
| if (buckets_are_explicit | ||
| and self.prefill_capture_num_tokens is not None | ||
| and sorted(set(legacy_buckets)) != sorted( | ||
| set(self.prefill_capture_num_tokens))): | ||
| raise ValueError( | ||
| "torch_compile_config.capture_num_tokens conflicts with " | ||
| "prefill_capture_num_tokens") | ||
| if not buckets_are_explicit: | ||
| logger.warning( | ||
| "TorchCompileConfig.capture_num_tokens is deprecated; use " | ||
| "prefill_capture_num_tokens instead.") | ||
| self.prefill_capture_num_tokens = list(legacy_buckets) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Auto-populated legacy buckets can trigger a spurious "conflicts" error.
TorchCompileConfig.set_default_capture_num_tokens (Line 4685) fills capture_num_tokens with the default list whenever enable_piecewise_cuda_graph=True, so legacy_buckets is non-None even when the user never set it. A user who writes
TorchLlmArgs(model=...,
prefill_capture_num_tokens=[128],
torch_compile_config=TorchCompileConfig(enable_piecewise_cuda_graph=True))is rejected with "torch_compile_config.capture_num_tokens conflicts with prefill_capture_num_tokens" despite specifying only one bucket list. Gate the legacy read on the field actually being set.
🐛 Proposed fix
- legacy_buckets = (compile_config.capture_num_tokens
- if compile_config is not None else None)
+ legacy_buckets = (compile_config.capture_num_tokens if
+ (compile_config is not None
+ and "capture_num_tokens"
+ in compile_config.model_fields_set) else None)🤖 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/llmapi/llm_args.py` around lines 5212 - 5226, Update the legacy
bucket handling in the TorchLlmArgs initialization flow to read
capture_num_tokens only when that field was explicitly set by the user, rather
than when TorchCompileConfig.set_default_capture_num_tokens auto-populated it.
Preserve the conflict validation for explicitly provided legacy buckets and the
migration warning/assignment for non-explicit prefill_capture_num_tokens.
| def test_dsv4_epilogue_fusion_returns_only_final_output_inside_breakable_graph() -> None: | ||
| metadata = SimpleNamespace(num_contexts=1, num_generations=1, num_tokens=5) | ||
| mla_layer = Mock(spec=MLA) | ||
| mla_layer._should_use_dsv4_epilogue_fusion.return_value = True | ||
| mla_layer.create_output.return_value = torch.empty(8, 4, 2) | ||
| hidden_states = torch.empty(8, 8) | ||
|
|
||
| with ( | ||
| patch( | ||
| "tensorrt_llm._torch.modules.mla._extract_mla_extra_attrs", | ||
| return_value=(metadata, mla_layer), | ||
| ), | ||
| patch( | ||
| "tensorrt_llm._torch.modules.mla.is_in_breakable_cuda_graph", | ||
| return_value=True, | ||
| ), | ||
| ): | ||
| outputs = create_mla_outputs_impl(hidden_states, "0") | ||
|
|
||
| assert outputs == [mla_layer.create_output.return_value] | ||
| mla_layer.create_output.assert_called_once_with( | ||
| hidden_states, | ||
| 1, | ||
| enable_dsv4_epilogue_fusion=True, | ||
| ) | ||
| mla_layer._create_dsv4_epilogue_buffers.assert_not_called() |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
The is_in_breakable_cuda_graph patch has no effect on create_mla_outputs_impl.
create_mla_outputs_impl (mla.py lines 137-148) only calls _extract_mla_extra_attrs, _should_use_dsv4_epilogue_fusion, and create_output; it never consults is_in_breakable_cuda_graph. The test therefore asserts the same thing inside and outside a breakable region, so the "inside_breakable_graph" guarantee in the name is not actually exercised. Either drop the patch (and rename) or add an assertion that distinguishes the two contexts.
🤖 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/modules/test_mla_registry.py` around lines 87 - 112,
Update
test_dsv4_epilogue_fusion_returns_only_final_output_inside_breakable_graph so it
does not patch is_in_breakable_cuda_graph, since create_mla_outputs_impl does
not consult it. Remove the unused context patch and rename the test to describe
only the behavior it actually verifies, or otherwise add a meaningful comparison
that distinguishes breakable and non-breakable contexts.
|
Hi all, this PR adds Breakable CUDA Graph (BCG) support for prefill. Here is a short review guide by ownership:
Thanks! |
There was a problem hiding this comment.
Actionable comments posted: 3
♻️ Duplicate comments (1)
tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/breakable_cuda_graph.py (1)
238-247: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
__enter__is still not transactional.If
_stream_context.__enter__(), any of theContextVar.set()calls, or_begin_new_segment()'scapture_begin()raises,__exit__is never invoked (Python only calls__exit__after a successful__enter__). The_install_wait_stream_hook()monkeypatch, the stream context, and the ContextVar tokens are left installed, corrupting subsequent captures. This was already flagged on a prior revision and remains unresolved.🔧 Proposed fix
def __enter__(self) -> "BreakableCUDAGraphCapture": _install_wait_stream_hook() - if self._stream is not None: - self._stream_context = torch.cuda.stream(self._stream) - self._stream_context.__enter__() - self._capture_token = _current_capture.set(self) - self._stream_token = _current_stream.set(self._stream or torch.cuda.current_stream()) - self._forked_token = _forked_streams.set(set()) - self._begin_new_segment() - return self + try: + if self._stream is not None: + self._stream_context = torch.cuda.stream(self._stream) + self._stream_context.__enter__() + self._capture_token = _current_capture.set(self) + self._stream_token = _current_stream.set(self._stream or torch.cuda.current_stream()) + self._forked_token = _forked_streams.set(set()) + self._begin_new_segment() + return self + except BaseException: + if self._forked_token is not None: + _forked_streams.reset(self._forked_token) + if self._stream_token is not None: + _current_stream.reset(self._stream_token) + if self._capture_token is not None: + _current_capture.reset(self._capture_token) + if self._stream_context is not None: + self._stream_context.__exit__(None, None, None) + self._stream_context = None + _uninstall_wait_stream_hook() + raise🤖 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/breakable_cuda_graph/breakable_cuda_graph.py` around lines 238 - 247, Make BreakableCUDAGraphCapture.__enter__ transactional: if stream context entry, any ContextVar.set, or _begin_new_segment fails, immediately roll back the wait-stream hook, entered stream context, and all tokens already created before re-raising the original exception. Reuse the existing cleanup behavior from __exit__ or a shared rollback helper, ensuring successful entry remains unchanged.
🧹 Nitpick comments (1)
tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/breakable_cuda_graph.py (1)
10-10: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valuePrefer
X | NoneandLiteralper repo style.
typing.Optional[...]is used throughout instead of theX | Nonesyntax, andcapture_error_mode: strcould beLiteral["global", "thread_local", "relaxed"]to match the accepted values documented forCUDAGraph.capture_begin.As per coding guidelines: "Annotate every function, use
Nonefor non-returning functions, avoidAnyand unnecessary type ignores, prefer built-in generic types and|... useLiteral,overload,TypeVar, orProtocolwhen appropriate."Also applies to: 220-236
🤖 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/breakable_cuda_graph/breakable_cuda_graph.py` at line 10, Update the type annotations in the affected breakable CUDA graph functions to replace Optional[X] with X | None and use built-in generic syntax where applicable. Import and apply Literal for capture_error_mode, restricting it to "global", "thread_local", or "relaxed" in the relevant signatures and related code around the noted functions. Remove Any where a more specific type is available and ensure every affected function has an explicit return annotation, using None for non-returning functions.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.
Inline comments:
In `@tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/breakable_cuda_graph.py`:
- Around line 146-175: Update eager_on_graph’s wrapper to restore a valid
capture segment when inner raises after capture._end_current_segment(). Ensure
capture._begin_new_segment() executes on both success and exception paths, while
preserving the original exception and only recording replay metadata after
successful execution.
In `@tensorrt_llm/_torch/utils.py`:
- Around line 159-184: Complete annotations for make_weak_ref by typing its
input and return with a recursive container-preserving alias or overloads,
without using Any, while covering tensors, scalars, tuples, lists, dictionaries,
and preserved unsupported values. Add the explicit None return annotation to
set_per_request_prefill_cuda_graph_flag. Keep the existing recursive behavior
unchanged.
In `@tests/integration/defs/accuracy/test_llm_api_pytorch.py`:
- Around line 3984-4062: Strengthen BCG coverage by exposing a supported
capture/replay metric and asserting it: in
tests/integration/defs/accuracy/test_llm_api_pytorch.py:3984-4062, update
test_mixed_breakable_cuda_graph to verify replay during the BREAKABLE run; in
tests/integration/defs/accuracy/test_llm_api_pytorch.py:3564-3623, assert replay
for baseline and explicit eager fallback for mtp3_fp8kv_chunked; in
tests/integration/defs/accuracy/test_llm_api_pytorch.py:6295-6333, use async
streaming with second-request admission during decoding and assert BCG replay;
in tests/integration/defs/accuracy/test_disaggregated_serving.py:2483-2543,
expose the same context-worker metric and assert capture/replay instead of
relying on the smoke response.
---
Duplicate comments:
In `@tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/breakable_cuda_graph.py`:
- Around line 238-247: Make BreakableCUDAGraphCapture.__enter__ transactional:
if stream context entry, any ContextVar.set, or _begin_new_segment fails,
immediately roll back the wait-stream hook, entered stream context, and all
tokens already created before re-raising the original exception. Reuse the
existing cleanup behavior from __exit__ or a shared rollback helper, ensuring
successful entry remains unchanged.
---
Nitpick comments:
In `@tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/breakable_cuda_graph.py`:
- Line 10: Update the type annotations in the affected breakable CUDA graph
functions to replace Optional[X] with X | None and use built-in generic syntax
where applicable. Import and apply Literal for capture_error_mode, restricting
it to "global", "thread_local", or "relaxed" in the relevant signatures and
related code around the noted functions. Remove Any where a more specific type
is available and ensure every affected function has an explicit return
annotation, using None for non-returning functions.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 3c8b504c-980c-4c16-8bc3-9bb84fd35b8a
📒 Files selected for processing (27)
docs/source/features/torch_compile_and_piecewise_cuda_graph.mdtensorrt_llm/_torch/compilation/piecewise_optimizer.pytensorrt_llm/_torch/models/modeling_minimaxm3.pytensorrt_llm/_torch/modules/attention.pytensorrt_llm/_torch/modules/mamba/gdn_mixer.pytensorrt_llm/_torch/modules/mla.pytensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/__init__.pytensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/breakable_cuda_graph.pytensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/context.pytensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/cuda_utils.pytensorrt_llm/_torch/pyexecutor/breakable_cuda_graph_runner.pytensorrt_llm/_torch/pyexecutor/model_engine.pytensorrt_llm/_torch/utils.pytensorrt_llm/llmapi/__init__.pytensorrt_llm/llmapi/llm_args.pytensorrt_llm/usage/llm_args_golden_manifest.jsontests/integration/defs/accuracy/test_disaggregated_serving.pytests/integration/defs/accuracy/test_llm_api_pytorch.pytests/integration/test_lists/test-db/l0_b200.ymltests/integration/test_lists/test-db/l0_dgx_b200.ymltests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_o_proj.pytests/unittest/_torch/compilation/test_remove_copy_pass.pytests/unittest/_torch/executor/test_breakable_cuda_graph.pytests/unittest/_torch/modules/test_mla_registry.pytests/unittest/api_stability/api_stability_core.pytests/unittest/api_stability/references/llm.yamltests/unittest/llmapi/test_llm_args.py
🚧 Files skipped from review as they are similar to previous changes (18)
- tests/unittest/api_stability/references/llm.yaml
- tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/context.py
- tensorrt_llm/llmapi/init.py
- tensorrt_llm/usage/llm_args_golden_manifest.json
- tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/init.py
- tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/cuda_utils.py
- tests/unittest/api_stability/api_stability_core.py
- tensorrt_llm/_torch/modules/attention.py
- tensorrt_llm/_torch/compilation/piecewise_optimizer.py
- tensorrt_llm/_torch/models/modeling_minimaxm3.py
- tensorrt_llm/_torch/modules/mamba/gdn_mixer.py
- tests/unittest/_torch/modules/test_mla_registry.py
- docs/source/features/torch_compile_and_piecewise_cuda_graph.md
- tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph_runner.py
- tensorrt_llm/llmapi/llm_args.py
- tests/unittest/llmapi/test_llm_args.py
- tensorrt_llm/_torch/pyexecutor/model_engine.py
- tensorrt_llm/_torch/modules/mla.py
| def eager_on_graph(inner: Callable) -> Callable: | ||
| """Run a callable eagerly between captured CUDA graph segments.""" | ||
|
|
||
| @functools.wraps(inner) | ||
| def wrapper(*args, **kwargs): | ||
| capture = _current_capture.get() | ||
| if capture is None: | ||
| return inner(*args, **kwargs) | ||
|
|
||
| logger.debug( | ||
| "Break CUDA graph for function %s", getattr(inner, "__name__", type(inner).__name__) | ||
| ) | ||
| capture._end_current_segment() | ||
| output = inner(*args, **kwargs) | ||
|
|
||
| captured_args = tuple(make_weak_ref(arg, preserve_unsupported=True) for arg in args) | ||
| captured_kwargs = { | ||
| key: make_weak_ref(value, preserve_unsupported=True) for key, value in kwargs.items() | ||
| } | ||
| captured_output = make_weak_ref(output, preserve_unsupported=True) | ||
|
|
||
| def replay_fn() -> Any: | ||
| new_output = inner(*captured_args, **captured_kwargs) | ||
| return _copy_output(captured_output, new_output) | ||
|
|
||
| capture.cuda_graph._break_functions.append(replay_fn) | ||
| capture._begin_new_segment() | ||
| return output | ||
|
|
||
| return wrapper |
There was a problem hiding this comment.
🩺 Stability & Availability | 🔴 Critical | ⚡ Quick win
Exception inside an eager break corrupts capture state.
_end_current_segment() runs at Line 158 (which calls capture_end()), but if inner(*args, **kwargs) raises at Line 159, _begin_new_segment() at Line 172 never executes. When the exception unwinds through the enclosing with BreakableCUDAGraphCapture(...) block, __exit__ calls self._end_current_segment() again on a segment that was already ended, producing a confusing secondary failure that masks the original exception.
🐛 Proposed fix
capture._end_current_segment()
- output = inner(*args, **kwargs)
+ try:
+ output = inner(*args, **kwargs)
+ except BaseException:
+ capture._begin_new_segment()
+ raise📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| def eager_on_graph(inner: Callable) -> Callable: | |
| """Run a callable eagerly between captured CUDA graph segments.""" | |
| @functools.wraps(inner) | |
| def wrapper(*args, **kwargs): | |
| capture = _current_capture.get() | |
| if capture is None: | |
| return inner(*args, **kwargs) | |
| logger.debug( | |
| "Break CUDA graph for function %s", getattr(inner, "__name__", type(inner).__name__) | |
| ) | |
| capture._end_current_segment() | |
| output = inner(*args, **kwargs) | |
| captured_args = tuple(make_weak_ref(arg, preserve_unsupported=True) for arg in args) | |
| captured_kwargs = { | |
| key: make_weak_ref(value, preserve_unsupported=True) for key, value in kwargs.items() | |
| } | |
| captured_output = make_weak_ref(output, preserve_unsupported=True) | |
| def replay_fn() -> Any: | |
| new_output = inner(*captured_args, **captured_kwargs) | |
| return _copy_output(captured_output, new_output) | |
| capture.cuda_graph._break_functions.append(replay_fn) | |
| capture._begin_new_segment() | |
| return output | |
| return wrapper | |
| def eager_on_graph(inner: Callable) -> Callable: | |
| """Run a callable eagerly between captured CUDA graph segments.""" | |
| `@functools.wraps`(inner) | |
| def wrapper(*args, **kwargs): | |
| capture = _current_capture.get() | |
| if capture is None: | |
| return inner(*args, **kwargs) | |
| logger.debug( | |
| "Break CUDA graph for function %s", getattr(inner, "__name__", type(inner).__name__) | |
| ) | |
| capture._end_current_segment() | |
| try: | |
| output = inner(*args, **kwargs) | |
| except BaseException: | |
| capture._begin_new_segment() | |
| raise | |
| captured_args = tuple(make_weak_ref(arg, preserve_unsupported=True) for arg in args) | |
| captured_kwargs = { | |
| key: make_weak_ref(value, preserve_unsupported=True) for key, value in kwargs.items() | |
| } | |
| captured_output = make_weak_ref(output, preserve_unsupported=True) | |
| def replay_fn() -> Any: | |
| new_output = inner(*captured_args, **captured_kwargs) | |
| return _copy_output(captured_output, new_output) | |
| capture.cuda_graph._break_functions.append(replay_fn) | |
| capture._begin_new_segment() | |
| return output | |
| return wrapper |
🤖 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/breakable_cuda_graph/breakable_cuda_graph.py`
around lines 146 - 175, Update eager_on_graph’s wrapper to restore a valid
capture segment when inner raises after capture._end_current_segment(). Ensure
capture._begin_new_segment() executes on both success and exception paths, while
preserving the original exception and only recording replay metadata after
successful execution.
| def make_weak_ref(x, preserve_unsupported: bool = False): | ||
|
|
||
| if isinstance(x, torch.Tensor): | ||
| return convert_to_torch_tensor( | ||
| TensorWrapper(x.data_ptr(), x.dtype, x.shape, | ||
| x.stride())) if x.is_cuda else x | ||
| elif isinstance(x, tuple): | ||
| return tuple(make_weak_ref(i) for i in x) | ||
| return tuple( | ||
| make_weak_ref(i, preserve_unsupported=preserve_unsupported) | ||
| for i in x) | ||
| elif isinstance(x, list): | ||
| return [make_weak_ref(i) for i in x] | ||
| return [ | ||
| make_weak_ref(i, preserve_unsupported=preserve_unsupported) | ||
| for i in x | ||
| ] | ||
| elif isinstance(x, dict): | ||
| return {k: make_weak_ref(v) for k, v in x.items()} | ||
| return { | ||
| k: make_weak_ref(v, preserve_unsupported=preserve_unsupported) | ||
| for k, v in x.items() | ||
| } | ||
| elif isinstance(x, (int, float, bool)): | ||
| return x | ||
| elif preserve_unsupported: | ||
| return x | ||
| else: | ||
| raise TypeError(f"Invalid type {type(x)} to make weak ref") |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Add complete annotations to the changed helpers.
make_weak_ref still leaves x and its return untyped, and set_per_request_prefill_cuda_graph_flag needs -> None. Use a recursive alias or overloads rather than Any for the container-preserving helper. As per coding guidelines, “Annotate every function” and “avoid Any”.
Also applies to: 407-414
🤖 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/utils.py` around lines 159 - 184, Complete annotations
for make_weak_ref by typing its input and return with a recursive
container-preserving alias or overloads, without using Any, while covering
tensors, scalars, tuples, lists, dictionaries, and preserved unsupported values.
Add the explicit None return annotation to
set_per_request_prefill_cuda_graph_flag. Keep the existing recursive behavior
unchanged.
Source: Coding guidelines
| @pytest.mark.skip_less_mpi_world_size(8) | ||
| @pytest.mark.threadleak(enabled=False) | ||
| def test_mixed_breakable_cuda_graph(self): | ||
| from transformers import AutoTokenizer | ||
|
|
||
| tokenizer = AutoTokenizer.from_pretrained(self.MODEL_PATH) | ||
| base_prompt_ids = tokenizer.encode( | ||
| "TensorRT-LLM accelerates reliable large language model inference " | ||
| "with efficient attention, parallelism, and CUDA graphs. ", | ||
| add_special_tokens=False, | ||
| ) | ||
| assert base_prompt_ids | ||
|
|
||
| def make_prompt(prompt_length): | ||
| return (base_prompt_ids * | ||
| ((prompt_length + len(base_prompt_ids) - 1) // | ||
| len(base_prompt_ids)))[:prompt_length] | ||
|
|
||
| generation_prompt = make_prompt(64) | ||
| context_prompt = make_prompt(129) | ||
| sampling_params = SamplingParams( | ||
| max_tokens=8, | ||
| min_tokens=8, | ||
| seed=42, | ||
| temperature=0, | ||
| ignore_eos=True, | ||
| detokenize=False, | ||
| add_special_tokens=False, | ||
| ) | ||
| common_llm_kwargs = dict( | ||
| tensor_parallel_size=8, | ||
| moe_expert_parallel_size=8, | ||
| moe_config=MoeConfig(backend="TRTLLM"), | ||
| enable_attention_dp=True, | ||
| max_batch_size=8, | ||
| max_num_tokens=1024, | ||
| max_seq_len=2048, | ||
| kv_cache_config=KvCacheConfig( | ||
| enable_block_reuse=False, | ||
| dtype="fp8", | ||
| free_gpu_memory_fraction=0.6, | ||
| ), | ||
| cuda_graph_config=CudaGraphConfig( | ||
| batch_sizes=[1, 2, 4, 6, 8], | ||
| enable_padding=True, | ||
| ), | ||
| ) | ||
|
|
||
| def run(backend): | ||
| with LLM( | ||
| self.MODEL_PATH, | ||
| **common_llm_kwargs, | ||
| prefill_cuda_graph_backend=backend, | ||
| prefill_capture_num_tokens=[128, 256, 512, 1024], | ||
| ) as llm: | ||
| generation_request = llm.generate_async( | ||
| generation_prompt, | ||
| sampling_params=sampling_params, | ||
| streaming=True, | ||
| ) | ||
| next(generation_request) | ||
| assert not generation_request.finished | ||
|
|
||
| # Admit a context request while the first request is decoding. | ||
| context_request = llm.generate_async( | ||
| context_prompt, | ||
| sampling_params=sampling_params, | ||
| streaming=False, | ||
| ) | ||
| generation_output = generation_request.result() | ||
| context_output = context_request.result() | ||
| return [ | ||
| generation_output.outputs[0].token_ids, | ||
| context_output.outputs[0].token_ids, | ||
| ] | ||
|
|
||
| eager_token_ids = run(PrefillCudaGraphBackend.DISABLED) | ||
| breakable_token_ids = run(PrefillCudaGraphBackend.BREAKABLE) | ||
| assert breakable_token_ids == eager_token_ids |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy lift
Coverage summary — insufficient: the tests cannot prove BCG capture or replay occurred. Output equality and accuracy can pass when all requests fall back to eager execution; the MTP parameter is expected to fall back, but that expectation is not asserted either.
tests/integration/defs/accuracy/test_llm_api_pytorch.py#L3984-L4062: expose a supported capture/replay counter or test-only metric and assert replay for the BREAKABLE run.tests/integration/defs/accuracy/test_llm_api_pytorch.py#L3564-L3623: assert replay forbaselineand explicit eager fallback formtp3_fp8kv_chunked.tests/integration/defs/accuracy/test_llm_api_pytorch.py#L6295-L6333: use async streaming plus admission of the second request while the first decodes, then assert BCG replay; the current synchronous loop only submits prefill batches.tests/integration/defs/accuracy/test_disaggregated_serving.py#L2483-L2543: surface the same context-worker metric and assert capture/replay rather than only a one-sample smoke response.
Changed tests: DeepSeek V32 NVFP4 BCG, DeepSeek V4 Flash mixed BCG, Qwen3.5-4B BF16 BCG, and DeepSeek V4 Flash disaggregated BCG. Registration is present in tests/integration/test_lists/test-db/l0_b200.yml and tests/integration/test_lists/test-db/l0_dgx_b200.yml; no QA-list sync is needed.
📍 Affects 2 files
tests/integration/defs/accuracy/test_llm_api_pytorch.py#L3984-L4062(this comment)tests/integration/defs/accuracy/test_llm_api_pytorch.py#L3564-L3623tests/integration/defs/accuracy/test_llm_api_pytorch.py#L6295-L6333tests/integration/defs/accuracy/test_disaggregated_serving.py#L2483-L2543
🤖 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/integration/defs/accuracy/test_llm_api_pytorch.py` around lines 3984 -
4062, Strengthen BCG coverage by exposing a supported capture/replay metric and
asserting it: in
tests/integration/defs/accuracy/test_llm_api_pytorch.py:3984-4062, update
test_mixed_breakable_cuda_graph to verify replay during the BREAKABLE run; in
tests/integration/defs/accuracy/test_llm_api_pytorch.py:3564-3623, assert replay
for baseline and explicit eager fallback for mtp3_fp8kv_chunked; in
tests/integration/defs/accuracy/test_llm_api_pytorch.py:6295-6333, use async
streaming with second-request admission during decoding and assert BCG replay;
in tests/integration/defs/accuracy/test_disaggregated_serving.py:2483-2543,
expose the same context-worker metric and assert capture/replay instead of
relying on the smoke response.
Source: Path instructions
|
/bot run |
|
PR_Github #62468 [ run ] triggered by Bot. Commit: |
|
PR_Github #62468 [ run ] completed with state
|
|
/bot run |
|
PR_Github #62546 [ run ] triggered by Bot. Commit: |
|
PR_Github #62546 [ run ] completed with state
|
Signed-off-by: GuanhuaWang2001 <300454435+GuanhuaWang2001@users.noreply.github.com>
Signed-off-by: GuanhuaWang2001 <300454435+GuanhuaWang2001@users.noreply.github.com>
Signed-off-by: GuanhuaWang2001 <300454435+GuanhuaWang2001@users.noreply.github.com>
Signed-off-by: GuanhuaWang2001 <300454435+GuanhuaWang2001@users.noreply.github.com>
Signed-off-by: GuanhuaWang2001 <300454435+GuanhuaWang2001@users.noreply.github.com>
Signed-off-by: GuanhuaWang2001 <300454435+GuanhuaWang2001@users.noreply.github.com>
Signed-off-by: GuanhuaWang2001 <300454435+GuanhuaWang2001@users.noreply.github.com>
Signed-off-by: GuanhuaWang2001 <300454435+GuanhuaWang2001@users.noreply.github.com>
Signed-off-by: GuanhuaWang2001 <300454435+GuanhuaWang2001@users.noreply.github.com>
Signed-off-by: GuanhuaWang2001 <300454435+GuanhuaWang2001@users.noreply.github.com>
5ab1b82 to
4df0f72
Compare
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 (2)
tests/unittest/_torch/compilation/test_remove_copy_pass.py (1)
86-87: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winAssert that the functionalized wrapper is removed.
The test checks output rewiring but not cleanup of
auto_functionalized_v2. Add the same removal assertion used by the preceding test; otherwise dead functionalization nodes could remain undetected.Proposed assertion
assert clone.args[0] is output + assert all(node.target != auto_functionalized_v2 for node in graph.nodes) graph.lint()🤖 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_remove_copy_pass.py` around lines 86 - 87, Extend the test assertion block around clone.args[0] and graph.lint() to verify that the auto_functionalized_v2 functionalized wrapper has been removed, matching the removal assertion used by the preceding test. Preserve the existing output-rewiring assertion and lint check.tensorrt_llm/_torch/models/modeling_minimaxm3.py (1)
1481-1505: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
_load_index_qk_proj_weightsignoresallow_partial_loading.It is invoked unconditionally before
super().load_weights(...), and the comment relies onLinear.load_weightsasserting when a source is missing. Withallow_partial_loading=True(dummy/partial weight loads), every sparse layer now hard-fails instead of skipping. Gate the fusion on the flag, or skip a module when both filtered dicts are empty.🛠️ Suggested guard
-def _load_index_qk_proj_weights(model: nn.Module, weights) -> None: +def _load_index_qk_proj_weights( + model: nn.Module, weights, allow_partial_loading: bool = False +) -> None: @@ q_weights = filter_weights(f"{parent}.index_q_proj", weights) k_weights = filter_weights(f"{parent}.index_k_proj", weights) + if allow_partial_loading and not (q_weights and k_weights): + continueand at the call site:
- _load_index_qk_proj_weights(self, weights) + _load_index_qk_proj_weights(self, weights, allow_partial_loading)Also applies to: 1530-1532
🤖 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/models/modeling_minimaxm3.py` around lines 1481 - 1505, Update _load_index_qk_proj_weights and its unconditional call site to honor allow_partial_loading: when partial loading is enabled, skip fusion for index_qk_proj modules whose filtered index_q_proj and index_k_proj sources are absent instead of invoking module.load_weights. Preserve the current strict assertion behavior for complete loads, and ensure the call-site flag is propagated to the helper.
♻️ Duplicate comments (2)
tensorrt_llm/_torch/pyexecutor/model_engine.py (1)
2049-2104: 🩺 Stability & Availability | 🔴 Critical | ⚡ Quick winPrefill CUDA-graph capture loops still lack the TP-symmetry guard.
Both the capture loop and the "most requests" warmup loop in
_capture_prefill_cuda_graphsskip a bucket with a bareif batch is None: continue, independently per TP rank — the same hazard already flagged on this code path in a prior review round (_assert_all_tp_ranks_have_warmup_batchis called in every other warmup loop in this file:_general_warmup_impl,_run_attention_warmup,_run_autotuner_warmup,_run_mamba_hybrid_warmup, but not here).Under attention-DP, per-rank KV cache availability can differ, so one rank may capture a bucket while a peer skips it. At replay time,
_get_padding_params's all-rank agreement only checksmax(attn_all_rank_num_tokens) <= max_captured_num_tokens— a static, config-derived ceiling identical on every rank — never whether the bucket was actually captured everywhere.forward()then decides per-rank whether to callbreakable_runner.execute()(replays captured NCCL collectives) or fall back to eagerforward_step(), based onbreakable_runner.has_graph(num_tokens). If ranks diverge on that, they run different collective-op sequences on the same iteration — a collective mismatch/deadlock, not just a silent-fallback correctness gap.🛡️ Suggested fix
with self._release_batch_context(warmup_request, resource_manager) as batch: - if batch is None: - continue + if batch is None and self.mapping.tp_size <= 1: + continue + self._assert_all_tp_ranks_have_warmup_batch( + batch, num_tokens) + if batch is None: + continueApply the equivalent change to the second ("most requests") loop below.
🤖 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 2049 - 2104, Update both warmup loops in _capture_prefill_cuda_graphs to call _assert_all_tp_ranks_have_warmup_batch before handling a missing batch, replacing the rank-local bare batch-is-None skip with the established TP-symmetry guard. Apply the same behavior to both the capture loop and the “most requests” warmup loop, preserving continuation only after the guard confirms the batch is unavailable consistently across ranks.tensorrt_llm/_torch/utils.py (1)
159-184: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winMissing annotations still unresolved.
make_weak_ref'sxparameter and return type remain untyped, andset_per_request_prefill_cuda_graph_flag/get_per_request_prefill_cuda_graph_flagstill lack explicit-> None/-> booltreatment consistency (setter has none). This was flagged in a prior review round on this same file and appears not yet addressed.As per coding guidelines, "Annotate every function, use
Nonefor non-returning functions, avoidAny".🩹 Suggested fix
-def make_weak_ref(x, preserve_unsupported: bool = False): +def make_weak_ref(x: Any, preserve_unsupported: bool = False) -> Any:-def set_per_request_prefill_cuda_graph_flag(enable: bool): +def set_per_request_prefill_cuda_graph_flag(enable: bool) -> None:Also applies to: 407-414
🤖 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/utils.py` around lines 159 - 184, Annotate make_weak_ref with an explicit parameter type and precise return type without using Any, covering its supported recursive values and preserved unsupported values. Also update set_per_request_prefill_cuda_graph_flag to explicitly return None and get_per_request_prefill_cuda_graph_flag to explicitly return bool, keeping their existing behavior unchanged.Source: Coding guidelines
🧹 Nitpick comments (5)
tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/cuda_utils.py (1)
17-25: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAnnotate
check_cuda_errors.The repo requires every function to be annotated; this one has neither a parameter nor a return type. A
Sequenceinput with anobject(or overloaded) return keeps it precise withoutAny.♻️ Suggested signature
-def check_cuda_errors(result): +def check_cuda_errors(result: Sequence[object]) -> object:As per coding guidelines: "Annotate every function, use
Nonefor non-returning functions, avoidAny".🤖 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/breakable_cuda_graph/cuda_utils.py` around lines 17 - 25, Annotate the check_cuda_errors function’s result parameter as a Sequence type and add an appropriate return annotation covering None, the second element, or the sliced remainder without using Any. Preserve the existing CUDA error validation and return behavior.Source: Coding guidelines
tests/unittest/_torch/executor/test_breakable_cuda_graph.py (1)
223-251: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winConsider covering the two capture-time guard paths.
BreakableCUDAGraphRunner.capture()raisesTypeErrorfor a non-tensor body output, andcapture_output()raisesValueErrorwhen buckets are captured in ascending order. Both are cheap to assert and lock in the descending-bucket contract that_shared_outputslicing depends on.🤖 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/executor/test_breakable_cuda_graph.py` around lines 223 - 251, Extend the CUDA graph runner tests around BreakableCUDAGraphRunner.capture and capture_output to cover both guard paths: assert capture raises TypeError when the body returns a non-tensor, and assert capture_output raises ValueError when buckets are captured in ascending order. Reuse the existing runner and capture setup where appropriate, and preserve the descending-bucket ordering contract.tensorrt_llm/_torch/models/modeling_minimaxm3.py (1)
735-753: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚖️ Poor tradeoffThe fused
index_qk_projrefactor looks orthogonal to this PR's BCG scope.Merging
index_q_proj/index_k_projinto one GEMM plus its bespoke checkpoint loader is a MiniMax-M3 performance change, not part of breaking-prefill CUDA graphs. Splitting it into its own PR would keep the BCG change reviewable and make the weight-loading change independently bisectable.As per coding guidelines: "Keep each pull request focused on one concern; split unrelated changes into separate pull requests."
Also applies to: 1202-1203
🤖 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/models/modeling_minimaxm3.py` around lines 735 - 753, Remove the fused index_qk_proj refactor from the MiniMax-M3 changes, restoring separate index_q_proj and index_k_proj Linear modules and their corresponding checkpoint-loading logic. Keep the breaking-prefill CUDA graph changes intact, and ensure weight loading continues to target the separate projection symbols.Source: Coding guidelines
tensorrt_llm/_torch/modules/mamba/gdn_mixer.py (1)
1060-1070: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDispatch shape differs from the attention/MLA call sites.
eager_on_graph's wrapper already returnsinner(...)whentorch.compiler.is_compiling(), so the explicituse_breakable_cuda_graphbranch and the two-op selection are redundant here.attention.pyandmla.pysimply always call the wrapped variant; matching that keeps the three dispatch sites uniform.♻️ Suggested simplification
- use_breakable_cuda_graph = not is_torch_compiling() and is_in_breakable_cuda_graph() - if self.register_to_config and (is_torch_compiling() or use_breakable_cuda_graph): + if self.register_to_config and (is_torch_compiling() or is_in_breakable_cuda_graph()): attn_out = mixed_qkv.new_empty( (1, mixed_qkv.shape[0], self.num_v_heads_per_tp, self.head_v_dim) ) - custom_op = ( - breakable_gdn_custom_op_inplace - if use_breakable_cuda_graph - else gdn_custom_op_inplace - ) - custom_op(mixed_qkv, a, b, self.layer_idx_str, attn_out) + breakable_gdn_custom_op_inplace(mixed_qkv, a, b, self.layer_idx_str, attn_out)🤖 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/modules/mamba/gdn_mixer.py` around lines 1060 - 1070, Update the dispatch block in the relevant GDN mixer forward path to always call the wrapped breakable-GDN operation, removing the separate use_breakable_cuda_graph condition and custom_op selection. Preserve the existing attn_out allocation and arguments, matching the uniform dispatch pattern used by the attention and MLA call sites.tests/unittest/llmapi/test_llm_args.py (1)
1843-1858: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd a regression test for the "auto-populated legacy buckets" false-conflict scenario.
The tests here cover explicit-legacy-vs-explicit-new conflicts, but not the specific scenario a prior review flagged:
TorchCompileConfig.set_default_capture_num_tokensauto-populatescapture_num_tokenswheneverenable_piecewise_cuda_graph=True, even if the user never set it. A regression test asserting thatTorchLlmArgs(prefill_capture_num_tokens=[128], torch_compile_config=TorchCompileConfig(enable_piecewise_cuda_graph=True))does not raise would lock in that 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/llmapi/test_llm_args.py` around lines 1843 - 1858, Add a regression case to test_explicit_legacy_and_new_config_conflicts that constructs TorchLlmArgs with prefill_capture_num_tokens=[128] and TorchCompileConfig(enable_piecewise_cuda_graph=True), asserting construction succeeds without raising; keep the existing explicit-conflict assertions unchanged.
🤖 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/models/modeling_minimaxm3.py`:
- Around line 1481-1505: Update _load_index_qk_proj_weights and its
unconditional call site to honor allow_partial_loading: when partial loading is
enabled, skip fusion for index_qk_proj modules whose filtered index_q_proj and
index_k_proj sources are absent instead of invoking module.load_weights.
Preserve the current strict assertion behavior for complete loads, and ensure
the call-site flag is propagated to the helper.
In `@tests/unittest/_torch/compilation/test_remove_copy_pass.py`:
- Around line 86-87: Extend the test assertion block around clone.args[0] and
graph.lint() to verify that the auto_functionalized_v2 functionalized wrapper
has been removed, matching the removal assertion used by the preceding test.
Preserve the existing output-rewiring assertion and lint check.
---
Duplicate comments:
In `@tensorrt_llm/_torch/pyexecutor/model_engine.py`:
- Around line 2049-2104: Update both warmup loops in
_capture_prefill_cuda_graphs to call _assert_all_tp_ranks_have_warmup_batch
before handling a missing batch, replacing the rank-local bare batch-is-None
skip with the established TP-symmetry guard. Apply the same behavior to both the
capture loop and the “most requests” warmup loop, preserving continuation only
after the guard confirms the batch is unavailable consistently across ranks.
In `@tensorrt_llm/_torch/utils.py`:
- Around line 159-184: Annotate make_weak_ref with an explicit parameter type
and precise return type without using Any, covering its supported recursive
values and preserved unsupported values. Also update
set_per_request_prefill_cuda_graph_flag to explicitly return None and
get_per_request_prefill_cuda_graph_flag to explicitly return bool, keeping their
existing behavior unchanged.
---
Nitpick comments:
In `@tensorrt_llm/_torch/models/modeling_minimaxm3.py`:
- Around line 735-753: Remove the fused index_qk_proj refactor from the
MiniMax-M3 changes, restoring separate index_q_proj and index_k_proj Linear
modules and their corresponding checkpoint-loading logic. Keep the
breaking-prefill CUDA graph changes intact, and ensure weight loading continues
to target the separate projection symbols.
In `@tensorrt_llm/_torch/modules/mamba/gdn_mixer.py`:
- Around line 1060-1070: Update the dispatch block in the relevant GDN mixer
forward path to always call the wrapped breakable-GDN operation, removing the
separate use_breakable_cuda_graph condition and custom_op selection. Preserve
the existing attn_out allocation and arguments, matching the uniform dispatch
pattern used by the attention and MLA call sites.
In `@tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/cuda_utils.py`:
- Around line 17-25: Annotate the check_cuda_errors function’s result parameter
as a Sequence type and add an appropriate return annotation covering None, the
second element, or the sliced remainder without using Any. Preserve the existing
CUDA error validation and return behavior.
In `@tests/unittest/_torch/executor/test_breakable_cuda_graph.py`:
- Around line 223-251: Extend the CUDA graph runner tests around
BreakableCUDAGraphRunner.capture and capture_output to cover both guard paths:
assert capture raises TypeError when the body returns a non-tensor, and assert
capture_output raises ValueError when buckets are captured in ascending order.
Reuse the existing runner and capture setup where appropriate, and preserve the
descending-bucket ordering contract.
In `@tests/unittest/llmapi/test_llm_args.py`:
- Around line 1843-1858: Add a regression case to
test_explicit_legacy_and_new_config_conflicts that constructs TorchLlmArgs with
prefill_capture_num_tokens=[128] and
TorchCompileConfig(enable_piecewise_cuda_graph=True), asserting construction
succeeds without raising; keep the existing explicit-conflict assertions
unchanged.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 47d618f7-7b1c-4d71-bed6-3e4816593861
📒 Files selected for processing (27)
docs/source/features/torch_compile_and_piecewise_cuda_graph.mdtensorrt_llm/_torch/compilation/piecewise_optimizer.pytensorrt_llm/_torch/models/modeling_minimaxm3.pytensorrt_llm/_torch/modules/attention.pytensorrt_llm/_torch/modules/mamba/gdn_mixer.pytensorrt_llm/_torch/modules/mla.pytensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/__init__.pytensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/breakable_cuda_graph.pytensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/context.pytensorrt_llm/_torch/pyexecutor/breakable_cuda_graph/cuda_utils.pytensorrt_llm/_torch/pyexecutor/breakable_cuda_graph_runner.pytensorrt_llm/_torch/pyexecutor/model_engine.pytensorrt_llm/_torch/utils.pytensorrt_llm/llmapi/__init__.pytensorrt_llm/llmapi/llm_args.pytensorrt_llm/usage/llm_args_golden_manifest.jsontests/integration/defs/accuracy/test_disaggregated_serving.pytests/integration/defs/accuracy/test_llm_api_pytorch.pytests/integration/test_lists/test-db/l0_b200.ymltests/integration/test_lists/test-db/l0_dgx_b200.ymltests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_o_proj.pytests/unittest/_torch/compilation/test_remove_copy_pass.pytests/unittest/_torch/executor/test_breakable_cuda_graph.pytests/unittest/_torch/modules/test_mla_registry.pytests/unittest/api_stability/api_stability_core.pytests/unittest/api_stability/references/llm.yamltests/unittest/llmapi/test_llm_args.py
🚧 Files skipped from review as they are similar to previous changes (1)
- docs/source/features/torch_compile_and_piecewise_cuda_graph.md
|
/bot run |
|
PR_Github #62703 [ run ] triggered by Bot. Commit: |
|
PR_Github #62703 [ run ] completed with state
|
| - accuracy/test_llm_api_pytorch.py::TestGLM52::test_nvfp4[tp_size=8-ep_size=8] TIMEOUT (60) | ||
| - accuracy/test_llm_api_pytorch.py::TestDeepSeekV4Pro::test_gsm8k_full_accuracy TIMEOUT (240) | ||
| - accuracy/test_llm_api_pytorch.py::TestDeepSeekV4Flash::test_mixed_breakable_cuda_graph TIMEOUT (120) ISOLATION | ||
| - accuracy/test_disaggregated_serving.py::TestDeepSeekV4Flash::test_prefill_breakable_cuda_graph TIMEOUT (120) ISOLATION |
There was a problem hiding this comment.
BCG only affects the prefill computation path and does not change KV cache layout, write timing, or the transceiver's send path from the context server to the generation server. The existing test_auto_dtype already covers the disaggregated path for this model, and test_mixed_breakable_cuda_graph covers BCG accuracy. Would you be open to removing this test (and the corresponding yaml file)? After doing this, we will approve this PR. Thanks for the effort.
There was a problem hiding this comment.
thanks for review, i removed those test.
Signed-off-by: GuanhuaWang2001 <300454435+GuanhuaWang2001@users.noreply.github.com>
|
/bot run |
|
PR_Github #62724 [ run ] triggered by Bot. Commit: |
Signed-off-by: GaunhuaWang-nv <300454435+GuanhuaWang2001@users.noreply.github.com>
|
PR_Github #62724 [ run ] completed with state
|
6 participants
Summary
Adds breakable prefill CUDA graphs (BCG) support to the PyTorch backend by capturing decoder-body segments into CUDA-graph segments while executing dynamic Attention/GDN operations eagerly at explicit breakpoints. Introduces prefill-specific CUDA-graph configuration (DISABLED/PIECEWISE/BREAKABLE) and migrates legacy piecewise options into these new prefill fields (with deprecation warnings). Runtime routing uses per-request prefill gating so only supported prefill requests use BCG capture/replay; unsupported cases fall back to eager execution.
Integrates
BreakableCUDAGraphRunnerintoModelEnginefor token-bucket capture/replay, including prefill padding feasibility calculation, capture warmup/capture/execute routing, and robust capture teardown. Attention/GDN/minimaxm3 custom ops are wrapped with BCG-aware execution viaeager_on_graphandis_in_breakable_cuda_graph()to allow correct custom-op behavior inside breakable graph regions.Also updates DeepSeek-V4 “DSv4 epilogue fusion” MLA plumbing to a simplified single-tensor custom-op schema and routes MLA custom-op call sites through BCG-wrapped variants when breakable CUDA-graph execution is active. Includes docs, API/export updates, and expanded unit/integration test coverage for the breakable capture lifecycle and config migration.
Dev Engineer Review
BCG implementation & routing correctness
BreakableCUDAGraph,BreakableCUDAGraphCapture,eager_on_graph, andbreak_graph, including ContextVar-based capture state, stream capture segmentation, side-stream synchronization, and_copy_outputwriteback for structured outputs.BreakableCUDAGraphRunnerstate machine (IDLE/WARMUP/CAPTURE/REPLAY) that captures only the decoder-body by temporarily patchinglayer_model.forward, then replays for requestednum_tokensbuckets with validated output copying.ModelEngineviaPrefillCudaGraphBackend, including:_capture_prefill_cuda_graphsand_run_cuda_graph_warmup_get_padding_params(with TP-rank-uniform gating)set/get_per_request_prefill_cuda_graph_flag(...), with eager fallback when graphs are missing, warming up, capture is in progress, or the per-request flag is disabled.Correctness of custom-op behavior under breakable graphs
attn_custom_op_inplacewithin breakable regions by introducingmaybe_bcg_attn_custom_op_inplace = eager_on_graph(...)and expanding the dispatch condition to includeis_in_breakable_cuda_graph().gdn_custom_op_inplaceselection under breakable mode (while preserving non-marlin constraints), plus refactored GDN state zeroing into a helper.is_in_breakable_cuda_graph()is true.outputtensor; updated routing to call BCG-wrapped custom-op variants when inside breakable CUDA-graph execution.API/config consistency
PrefillCudaGraphBackendand newTorchLlmArgsfields:prefill_cuda_graph_backendprefill_capture_num_tokensnormalize_prefill_cuda_graph_config()migration:torch_compile_config.enable_piecewise_cuda_graph/torch_compile_config.capture_num_tokens→ prefill fields (with warnings)Nonetorch_compile_config, andencode_onlyconflicts)prefill_capture_num_tokensand PIECEWISE behavior.tensorrt_llm/_torch/utils.py.Performance/memory/execution boundaries
make_weak_ref(..., preserve_unsupported=...)) to safely replay complex outputs and preserve unsupported objects where needed.Config/docs updates
prefill_cuda_graph_backend: breakableexample.piecewise_optimizer.pyto use per-request prefill flags.CI / follow-ups
f8fad38, followed by retriggers where downstreamL0_MergeRequest_PRpipelines reportedFAILURE; additional follow-up was requested. Given current test-list entry TIMEOUTs (below), CI stability and runtime performance investigation likely needs follow-up.QA Engineer Review
Test-list / test-db changes (integration CI selection)
Modified:
tests/integration/test_lists/test-db/l0_b200.ymlaccuracy/test_llm_api_pytorch.py::TestQwen3_5_4B::test_bf16_breakable_prefill_cuda_graphtests/integration/test_lists/test-db/l0_dgx_b200.ymlaccuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_nvfp4_multi_gpus_breakable_cuda_graph[baseline](TIMEOUT)accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_nvfp4_multi_gpus_breakable_cuda_graph[mtp3_fp8kv_chunked](TIMEOUT)accuracy/test_llm_api_pytorch.py::TestDeepSeekV4Flash::test_mixed_breakable_cuda_graph(TIMEOUT, ISOLATION)accuracy/test_disaggregated_serving.py::TestDeepSeekV4Flash::test_prefill_breakable_cuda_graph(TIMEOUT, ISOLATION)Verdict (test-list coverage): needs follow-up (entries currently show TIMEOUTs).
Test code changes (files outside test-db)
Integration tests
tests/integration/defs/accuracy/test_llm_api_pytorch.pyTestQwen3_5_4B.test_bf16_breakable_prefill_cuda_graphTestDeepSeekV32.test_nvfp4_multi_gpus_breakable_cuda_graph(...)TestDeepSeekV4Flash.test_mixed_breakable_cuda_graphtest-db/:l0_b200.yml(Qwen3.5-4B bf16 breakable prefill)l0_dgx_b200.yml(DeepSeekV32 + DeepSeekV4Flash breakable cases; TIMEOUT observed)tests/integration/defs/accuracy/test_disaggregated_serving.pyTestDeepSeekV4Flash.test_prefill_breakable_cuda_graphtest-db/:l0_dgx_b200.yml(TIMEOUT observed)Unit tests (CPU/GPU dependent; executed in unit CI lanes)
tests/unittest/_torch/executor/test_breakable_cuda_graph.pytest_runner_warmup_capture_execute_and_shared_outputtest-db/: not indicated; should run in unit CI (no test-db mapping provided)tests/unittest/_torch/compilation/test_remove_copy_pass.pytests/unittest/_torch/modules/test_mla_registry.pytests/unittest/llmapi/test_llm_args.pyprefill_cuda_graph_backend/prefill_capture_num_tokensmigration/validation and padding equivalence behaviortests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_o_proj.pyenable_dsv4_epilogue_fusion=FalseVerdict (test code coverage): insufficient / needs follow-up (because the primary new integration test-db entries are timing out in
l0_dgx_b200.yml, and prior CI failures were reported in the objectives).Summary
Adds breakable prefill CUDA graphs (BCG) to the PyTorch backend. BCG captures the decoder body as segmented CUDA graphs while executing dynamic Attention/GDN operations eagerly at graph breakpoints.
Design
BreakableCUDAGraphRunnerintegrated directly intoModelEngine.Validation
Tested with a Release SM100 build on B200 GPUs:
All latency values are end-to-end
LLM.generate()P50 in milliseconds with one generated token, overlap scheduling disabled, and KV-cache block reuse disabled.Qwen3.5-4B BF16, 1×B200
Parameters:
Backend overrides:
Qwen3.5-397B-A17B NVFP4, 4×B200
Common parameters:
Attention-DP: