[TRTLLM-13579][feat] BREAKING: Support BCG in Prefill

docs/source/features/torch_compile_and_piecewise_cuda_graph.md

@@ -1,4 +1,4 @@
-# Torch Compile & Piecewise CUDA Graph
+# Torch Compile & Prefill CUDA Graph
 
 In this guide, we show how to enable torch.compile and Piecewise CUDA Graph in TensorRT LLM. TensorRT LLM uses torch.compile for lightweight vertical fusion and Piecewise CUDA Graph.
 
@@ -41,12 +41,29 @@ To enable torch.compile and Piecewise CUDA Graph, add the following configuratio
 
 ```yaml
 ... # Other extra config
+prefill_cuda_graph_backend: piecewise
+prefill_capture_num_tokens: '${capture_num_tokens}' # e.g. [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, ..., 3072]
 torch_compile_config:
-  capture_num_tokens: '${capture_num_tokens}' # List of num tokens to capture. e.g., [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, ..., 3072]
   enable_userbuffers: false
-  enable_piecewise_cuda_graph: true
 ```
 
+`TorchCompileConfig.enable_piecewise_cuda_graph` and
+`TorchCompileConfig.capture_num_tokens` are deprecated aliases for these
+prefill-specific options.
+
+The experimental breakable implementation can capture the model body without
+torch.compile:
+
+```yaml
+prefill_cuda_graph_backend: breakable
+prefill_capture_num_tokens: [128, 256, 512]
+```
+
+The first version of the breakable backend supports BF16 Qwen3.5 on one GPU for
+context-only, tensor/pipeline parallelism and mixed context/decode batches with KV cache. Speculative
+decoding, LoRA, multimodal inputs, and context
+logits fall back to eager execution or are rejected during initialization.
+
 ## Tips for Piecewise CUDA Graph
 
 ### Piecewise CUDA Graph & Generation Only CUDA Graph
@@ -59,9 +76,10 @@ cuda_graph_config:
   max_batch_size: 1024 # Specify max capture batch size for generation only cuda graph. By default, TensorRT LLM will generate a capture list based on it.
 
 torch_compile_config:
-  capture_num_tokens: '${capture_num_tokens}' # Specify capture_num_tokens for piecewise cuda graph
   enable_userbuffers: false
-  enable_piecewise_cuda_graph: true
+
+prefill_cuda_graph_backend: piecewise
+prefill_capture_num_tokens: '${capture_num_tokens}'
 ```
 
 ### Piecewise CUDA Graph Padding

tensorrt_llm/_torch/compilation/piecewise_optimizer.py

@@ -12,7 +12,7 @@
 from tensorrt_llm.llmapi.utils import enable_llm_debug
 
 from ..utils import (get_model_extra_attrs,
-                     get_per_request_piecewise_cuda_graph_flag,
+                     get_per_request_prefill_cuda_graph_flag,
                      get_piecewise_cuda_graph_flag, make_weak_ref,
                      set_piecewise_running)
 from .multi_stream.auto_multi_stream import multi_stream_schedule
@@ -202,7 +202,7 @@ def __call__(self, *args):
         if (runtime_num_of_token is None
                 or runtime_num_of_token not in self.entries
                 or not get_piecewise_cuda_graph_flag()
-                or not get_per_request_piecewise_cuda_graph_flag()):
+                or not get_per_request_prefill_cuda_graph_flag()):
             return self.default_callable(*args)
 
         if self.is_first_runner or self.is_last_runner:

tensorrt_llm/_torch/models/modeling_minimaxm3.py

@@ -62,6 +62,7 @@
 )
 from ..modules.multi_stream_utils import maybe_execute_in_parallel
 from ..modules.rms_norm import RMSNorm
+from ..pyexecutor.breakable_cuda_graph import eager_on_graph, is_in_breakable_cuda_graph
 from ..utils import (
     ActivationType,
     AuxStreamType,
@@ -660,6 +661,9 @@ def minimax_m3_attn_custom_op_inplace(
     )
 
 
+maybe_bcg_minimax_m3_attn_custom_op_inplace = eager_on_graph(minimax_m3_attn_custom_op_inplace)
+
+
 class MiniMaxM3Attention(Attention):
     """M3 attention: dense (layers 0-2) or sparse (layers 3-59).
 
@@ -1238,8 +1242,8 @@ def _forward_attention_core(
         attn_metadata: AttentionMetadata,
     ) -> torch.Tensor:
         output = q.new_empty((q.shape[0], self.num_heads * self.head_dim))
-        if self.register_to_config and is_torch_compiling():
-            minimax_m3_attn_custom_op_inplace(
+        if self.register_to_config and (is_torch_compiling() or is_in_breakable_cuda_graph()):
+            maybe_bcg_minimax_m3_attn_custom_op_inplace(
                 q,
                 k,
                 v,

tensorrt_llm/_torch/modules/attention.py

@@ -22,6 +22,8 @@
                            cp_allgather, reducescatter)
 from ..model_config import ModelConfig
 from ..peft.lora.layer import LoraLayer, LoraModuleType
+from ..pyexecutor.breakable_cuda_graph import (eager_on_graph,
+                                               is_in_breakable_cuda_graph)
 from ..utils import (Fp4QuantizedTensor, get_model_extra_attrs,
                      is_nvfp4_marlin_enabled, is_torch_compiling)
 from .linear import Linear, TensorParallelMode, WeightMode, WeightsLoadingConfig
@@ -115,6 +117,9 @@ def attn_custom_op_inplace(
     )
 
 
+maybe_bcg_attn_custom_op_inplace = eager_on_graph(attn_custom_op_inplace)
+
+
 def _helix_zero_kv_mask(
         attn_metadata: AttentionMetadata,
         num_tokens: int,
@@ -932,20 +937,19 @@ def forward_impl(
             if "mrope_position_deltas" in mrope_config:
                 mrope_position_deltas = mrope_config["mrope_position_deltas"]
 
-        # Currently only TRTLLM and FLASHINFER are torch compile compatible backends.
-        # Only enable custom inplace op when torch compiling.
-        use_custom_inplace_op = (self.register_to_config
-                                 and (self.attn_backend == "TRTLLM"
-                                      or self.attn_backend == "FLASHINFER")
-                                 and is_torch_compiling()
-                                 and not self.is_marlin_enabled)
+        # Currently only TRTLLM and FLASHINFER support the custom inplace op.
+        use_custom_inplace_op = (
+            self.register_to_config and
+            (self.attn_backend == "TRTLLM" or self.attn_backend == "FLASHINFER")
+            and (is_torch_compiling() or is_in_breakable_cuda_graph())
+            and not self.is_marlin_enabled)
 
         if use_custom_inplace_op:
             outputs = create_attn_outputs(q, attention_mask, self.layer_idx_str)
             assert len(outputs) == 1 or len(outputs) == 2
             output = outputs[0]
             output_sf = outputs[1] if len(outputs) == 2 else None
-            attn_custom_op_inplace(
+            maybe_bcg_attn_custom_op_inplace(
                 q,
                 k,
                 v,

tensorrt_llm/_torch/modules/mamba/gdn_mixer.py

@@ -31,6 +31,7 @@
 from ...attention_backend import AttentionMetadata
 from ...distributed import AllReduceParams
 from ...model_config import ModelConfig
+from ...pyexecutor.breakable_cuda_graph import eager_on_graph, is_in_breakable_cuda_graph
 from ...speculative import SpecMetadata
 from ...utils import EventType, get_model_extra_attrs, is_gdn_replay_enabled, is_torch_compiling
 from ..linear import FP8QDQLinearMethod, Linear, TensorParallelMode
@@ -174,6 +175,9 @@ def gdn_custom_op_inplace(
     )
 
 
+maybe_bcg_gdn_custom_op_inplace = eager_on_graph(gdn_custom_op_inplace)
+
+
 def ensure_divisibility(numerator, denominator):
     """Ensure that numerator is divisible by the denominator."""
     assert numerator % denominator == 0, "{} is not divisible by {}".format(numerator, denominator)
@@ -1053,11 +1057,12 @@ def forward(
     ):
         mixed_qkv, z, a, b = self._compute_tokenwise_inputs(hidden_states)
 
-        if self.register_to_config and is_torch_compiling():
+        use_breakable_cuda_graph = not is_torch_compiling() and is_in_breakable_cuda_graph()
+        if self.register_to_config and (is_torch_compiling() or use_breakable_cuda_graph):
             attn_out = mixed_qkv.new_empty(
                 (1, mixed_qkv.shape[0], self.num_v_heads_per_tp, self.head_v_dim)
             )
-            gdn_custom_op_inplace(mixed_qkv, a, b, self.layer_idx_str, attn_out)
+            maybe_bcg_gdn_custom_op_inplace(mixed_qkv, a, b, self.layer_idx_str, attn_out)
         else:
             attn_out = self.forward_core(
                 mixed_qkv,