[TRTLLM-13579][feat] BREAKING: Support BCG in Prefill - #16609
[TRTLLM-13579][feat] BREAKING: Support BCG in Prefill#16609GuanhuaWang2001 wants to merge 16 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.
There was a problem hiding this comment.
The old config path implicitly sets a default value for capture_num_tokens when enable_piecewise_cuda_graph is enabled. The new config path, however, validates conflicts between capture_num_tokens and prefill_capture_num_tokens. During the transition period, mixing old and new config styles may therefore introduce unexpected conflicts.
For example, a user may set prefill_capture_num_tokens together with enable_piecewise_cuda_graph without explicitly setting capture_num_tokens, which can still trigger a conflict unexpectedly.
Please deprecate the old config option as soon as possible, or add compatibility handling to avoid this transitional issue.
| 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.
|
PR_Github #63383 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
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>
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>
ac7061e to
fae1d66
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #63694 [ run ] triggered by Bot. Commit: |
Signed-off-by: GuanhuaWang2001 <300454435+GuanhuaWang2001@users.noreply.github.com>
|
PR_Github #63716 [ run ] triggered by Bot. Commit: |
|
PR_Github #63694 [ run ] completed with state |
|
PR_Github #63716 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #63760 [ run ] triggered by Bot. Commit: |
|
PR_Github #63760 [ run ] completed with state
|
12 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: