[TRTLLM-13579][feat] BREAKING: Support BCG in Prefill - #16609
[TRTLLM-13579][feat] BREAKING: Support BCG in Prefill#16609GuanhuaWang2001 wants to merge 24 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.
| 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.
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>
97f0c96 to
a656945
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #65013 [ run ] triggered by Bot. Commit: |
|
PR_Github #65013 [ run ] completed with state
|
07b46ba to
d10bbda
Compare
Signed-off-by: GuanhuaWang2001 <300454435+GuanhuaWang2001@users.noreply.github.com>
d10bbda to
efa156a
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #65082 [ run ] triggered by Bot. Commit: |
|
PR_Github #65082 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #65182 [ run ] triggered by Bot. Commit: |
Signed-off-by: GaunhuaWang-nv <300454435+GuanhuaWang2001@users.noreply.github.com>
|
PR_Github #65182 [ run ] completed with state
|
| decoder_model = (self.model if isinstance( | ||
| self.model, DecoderModelForCausalLM) else getattr( | ||
| self.model, "llm", None)) | ||
| if not isinstance(decoder_model, DecoderModelForCausalLM): |
There was a problem hiding this comment.
The advertised speculative-decoding support still fails for chain-drafter configurations. py_executor_creator.py:607-648 wraps the draft model in LinearDraftingLoopWrapper or StaticTreeDraftingLoopWrapper and copies these BREAKABLE args into the draft engine, but those wrappers expose .draft_model, not .llm; this branch therefore raises during engine construction. Either disable BCG for the draft engine or unwrap .draft_model here, and add a chain-drafter construction test.
There was a problem hiding this comment.
This affects the deprecated two-model/chain-drafter path. MTP uses the one-model and does not instantiate these wrappers. Deepseek V3 / V4 with mtp can run safety.
|
|
||
| max_num_streams: PositiveInt = Field( | ||
| default=1, | ||
| default=3, |
There was a problem hiding this comment.
| default=3, | |
| default=1, |
This default also drives compile-only TorchCompileConfig() with prefill_cuda_graph_backend=disabled; Backend.__call__ enables multi-stream scheduling whenever max_num_streams > 1, so this silently changes existing compile-only execution. Keep the public default at 1 and opt piecewise configs into 3 explicitly if that is desired.
There was a problem hiding this comment.
This default is intentional and was confirmed with the PCG maintainer, Liji. We want TorchCompileConfig to default to three streams, including compile-only configurations, so I’ll keep it unchanged.
|
|
||
| captured_args = tuple(make_weak_ref(arg) for arg in args) | ||
| captured_kwargs = {key: make_weak_ref(value) for key, value in kwargs.items()} | ||
| captured_output = make_weak_ref(output) |
There was a problem hiding this comment.
This bridge output must retain owning storage. It is allocated after _end_current_segment(), so unlike captured_args produced by the preceding graph segment it is not pinned by a CUDAGraph private pool. make_weak_ref() retains only its raw pointer; once wrapper returns, the caching allocator may reuse the original allocation, and replay then copies into a stale/reassigned address before the next segment reads it. The cited SGLang implementation deliberately keeps this output strongly referenced for this reason. Please add an allocator-churn replay test as well.
| captured_output = make_weak_ref(output) | |
| captured_output = output |
There was a problem hiding this comment.
Agreed. The eager output is allocated outside the graph-private pool and must retain owning storage. I changed it to a strong reference and added a weakref/allocator-churn replay regression test.
Signed-off-by: GuanhuaWang2001 <300454435+GuanhuaWang2001@users.noreply.github.com>
|
/bot run --disable-fail-fast |
|
PR_Github #65296 [ run ] triggered by Bot. Commit: |
|
PR_Github #65296 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #65460 [ run ] triggered by Bot. Commit: |
|
PR_Github #65460 [ run ] completed with state
|
15 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: