[TRTLLM-15936][fix] Enable breakable prefill CUDA graphs (BCG) for Nemotron-H hybrid models

tensorrt_llm/_torch/compilation/piecewise_optimizer.py

@@ -26,6 +26,7 @@ def _piecewise_boundary_ops():
         "mla_custom_op_inplace",
         "mla_dsa_attn_inplace",
         "gdn_custom_op_inplace",
+        "mamba2_custom_op_inplace",
         "minimax_m3_attn_custom_op_inplace",
     ]
     return [

tensorrt_llm/_torch/compilation/utils.py

@@ -201,6 +201,21 @@ def inplace_info():
         "gdn_custom_op_inplace": {
             1: "output"
         },
+        # Registered lazily: the op only exists once mamba2_mixer is imported
+        # (Mamba2/NemotronH family). Void boundary op mutating ssm_out:
+        # auto_functionalized returns (None, ssm_out), hence index 1.
+        "mamba2_custom_op_inplace": {
+            1: "ssm_out"
+        },
+        # Registered lazily: the op only exists once mamba2_mixer is
+        # imported (Mamba2/NemotronH family). Void op mutating (state, out):
+        # auto_functionalized returns (None, state, out), hence indices 1/2.
+        # Without this entry the pass leaves the functionalization clone of
+        # the full per-layer SSM state cache in every decode graph.
+        "flashinfer_selective_state_update": {
+            1: "state",
+            2: "out"
+        },
         "minimax_m3_attn_custom_op_inplace": {
             1: "output"
         },

tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py

@@ -15,6 +15,8 @@
 
 import functools
 import os
+import weakref
+from typing import Optional
 
 import torch
 from einops import rearrange, repeat
@@ -30,7 +32,10 @@
 from ...attention.backends import AttentionMetadata
 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 ...speculative import SpecMetadata
+from ...utils import get_model_extra_attrs, is_torch_compiling
 from ..linear import Linear, TensorParallelMode
 from .causal_conv1d import causal_conv1d_fn, causal_conv1d_update
 from .causal_conv1d_triton import \
@@ -46,6 +51,110 @@
 from .ssd_combined import mamba_chunk_scan_combined
 
 
+def _extract_mamba2_extra_attrs(layer_idx: str):
+    extra_attrs = get_model_extra_attrs()
+    assert extra_attrs is not None, "Model extra attrs is not set"
+
+    metadata_ref = extra_attrs.get("attention_metadata", None)
+    assert metadata_ref is not None, "Attention metadata is not set"
+    metadata = metadata_ref()
+    assert isinstance(metadata, AttentionMetadata)
+
+    mamba2_layers = extra_attrs.get("mamba2_layers", None)
+    assert mamba2_layers is not None, "Mamba2 layer is not registered"
+    layer_ref = mamba2_layers.get(layer_idx, None)
+    assert layer_ref is not None, \
+        f"Cannot find Mamba2 layer for layer {layer_idx}"
+    mamba_layer = layer_ref()
+    assert isinstance(mamba_layer, Mamba2Mixer)
+
+    return metadata, mamba_layer, extra_attrs.get("spec_metadata", None)
+
+
+@torch.library.custom_op("trtllm::mamba2_custom_op_inplace",
+                         mutates_args=("ssm_out", ))
+def mamba2_custom_op_inplace(zxbcdt: torch.Tensor, layer_idx: str,
+                             ssm_out: torch.Tensor) -> None:
+    # Piecewise boundary op (mirror of trtllm::gdn_custom_op_inplace): the
+    # whole conv+SSM core runs eagerly inside one opaque node, so the traced
+    # graph never reads batch-composition ints (num_ctx_tokens/num_decodes
+    # would otherwise specialize a dynamo variant per composition, which
+    # exhausts the recompile limit and hard-fails under fullgraph=True) and
+    # every surrounding piece stays uniform in the num_tokens dim.
+    attn_metadata, mamba_layer, spec_metadata = _extract_mamba2_extra_attrs(
+        layer_idx)
+    mamba_layer.forward_core(zxbcdt, attn_metadata,
+                             attn_metadata.mamba_metadata, spec_metadata,
+                             ssm_out)
+
+
+maybe_bcg_mamba2_custom_op_inplace = eager_on_graph(mamba2_custom_op_inplace)
+
+
+@torch.library.custom_op("trtllm::flashinfer_selective_state_update",
+                         mutates_args=("state", "out"),
+                         device_types="cuda")
+def _flashinfer_selective_state_update_op(
+        state: torch.Tensor,
+        x: torch.Tensor,
+        dt: torch.Tensor,
+        A: torch.Tensor,
+        B: torch.Tensor,
+        C: torch.Tensor,
+        D: torch.Tensor,
+        out: torch.Tensor,
+        dt_bias: Optional[torch.Tensor] = None,
+        dt_softplus: bool = False,
+        state_batch_indices: Optional[torch.Tensor] = None,
+        rand_seed: Optional[torch.Tensor] = None,
+        philox_rounds: int = 10) -> None:
+    # Opaque wrapper for the torch.compile path: flashinfer resolves its JIT
+    # module inside the Python call (torch.cuda.device_count() and friends),
+    # which dynamo cannot trace; the resulting graph break sits inside the
+    # decoder-layer loop, so dynamo skips the whole forward frame and the
+    # piecewise backend then fails on child frames that carry no
+    # input_ids/inputs_embeds placeholder.
+    # Contract notes: the op is void (remove_copy_for_mutates_args indexes
+    # every getitem user into the inplace_info map, so a real return value
+    # would KeyError at index 0), and `out` must be a plain intermediate,
+    # not a view of the caller's preallocated buffer (mutable view args are
+    # not handled by the rewrite).
+    kwargs = {}
+    if rand_seed is not None:
+        kwargs["rand_seed"] = rand_seed
+        kwargs["philox_rounds"] = philox_rounds
+    selective_state_update_fi(state,
+                              x,
+                              dt,
+                              A,
+                              B,
+                              C,
+                              D,
+                              z=None,
+                              dt_bias=dt_bias,
+                              dt_softplus=dt_softplus,
+                              state_batch_indices=state_batch_indices,
+                              out=out,
+                              **kwargs)
+
+
+@_flashinfer_selective_state_update_op.register_fake
+def _(state,
+      x,
+      dt,
+      A,
+      B,
+      C,
+      D,
+      out,
+      dt_bias=None,
+      dt_softplus=False,
+      state_batch_indices=None,
+      rand_seed=None,
+      philox_rounds=10) -> None:
+    return None
+
+
 class Mamba2Mixer(nn.Module):
 
     def __init__(
@@ -72,6 +181,22 @@ def __init__(
     ):
         super().__init__()
 
+        # Register into the model's extra attrs so the mamba2 boundary custom
+        # op can recover this module (and the live metadata) from just a
+        # layer-idx string while the surrounding forward is traced.
+        self.layer_idx_str = str(layer_idx)
+        self.register_to_config = False
+        if config is not None:
+            if "mamba2_layers" not in config.extra_attrs:
+                config.extra_attrs["mamba2_layers"] = {}
+            suffix = 0
+            while self.layer_idx_str in config.extra_attrs["mamba2_layers"]:
+                self.layer_idx_str = str(layer_idx) + f"_{suffix}"
+                suffix += 1
+            config.extra_attrs["mamba2_layers"][self.layer_idx_str] = \
+                weakref.ref(self)
+            self.register_to_config = True
+
         config = config or ModelConfig()
 
         if config.mapping.enable_attention_dp:
@@ -312,13 +437,67 @@ def forward(
         **kwargs,
     ) -> torch.Tensor:
 
+        # in_proj (LoRA is applied internally by Linear layer)
+        zxbcdt = self.in_proj(hidden_states,
+                              lora_params=lora_params,
+                              layer_idx=self.layer_idx)
+        z = zxbcdt[:, :self.tp_d_inner]
+
+        # Preallocated output buffer shared by the prefill and decode
+        # segments; forward_core fills the real rows and zeroes the pad tail.
+        # Allocated here, not in forward_core, on every path: the boundary op
+        # is void and mutates this buffer in place (mutates_args + inplace_info),
+        # which is what makes it a piecewise / BCG boundary. The eager path is
+        # unified on purpose; its only extra cost is zero_() of an empty pad
+        # slice (a no-op). Do not move the allocation into forward_core.
+        preallocated_ssm_out = torch.empty(
+            [zxbcdt.shape[0], (self.tp_nheads * self.head_dim)],
+            dtype=zxbcdt.dtype,
+            device=zxbcdt.device,
+        )
+
+        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):
+            # Route the conv+SSM core through the opaque boundary op (see its
+            # comment): under torch.compile the traced graph stays free of
+            # batch-composition ints; under breakable CUDA graph capture the
+            # op is the eager bridge between captured segments.
+            maybe_bcg_mamba2_custom_op_inplace(zxbcdt, self.layer_idx_str,
+                                               preallocated_ssm_out)
+        else:
+            self.forward_core(zxbcdt, attn_metadata, mamba_metadata,
+                              spec_metadata, preallocated_ssm_out)
+
+        # norm
+        # Full padded length through norm/out_proj so the residual stream
+        # keeps a consistent row count (pad rows are zeros; row-wise norm
+        # keeps them finite). The caller trims real tokens via gather_ids.
+        hidden_states = self.norm(preallocated_ssm_out, z)
+
+        # out_proj
+        out = self.out_proj(hidden_states,
+                            lora_params=lora_params,
+                            layer_idx=self.layer_idx)
+
+        return out
+
+    def forward_core(
+        self,
+        zxbcdt: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+        mamba_metadata: Mamba2Metadata,
+        spec_metadata: SpecMetadata | None,
+        preallocated_ssm_out: torch.Tensor,
+    ) -> None:
+
         # calculate split size
         num_prefills = attn_metadata.num_contexts
         num_decodes = attn_metadata.seq_lens.shape[0] - num_prefills
         num_prefill_tokens = attn_metadata.num_ctx_tokens
         num_decode_tokens = attn_metadata.num_tokens - num_prefill_tokens
         num_actual_tokens = attn_metadata.num_tokens
-        seqlen_split_size = [num_prefill_tokens, num_decode_tokens]
         batch_split_size = [num_prefills, num_decodes]
 
         state_indices = mamba_metadata.state_indices[:num_prefills +
@@ -331,32 +510,26 @@ def forward(
         state_indices_p, state_indices_d = torch.split(state_indices,
                                                        batch_split_size)
 
-        # in_proj (LoRA is applied internally by Linear layer)
-        zxbcdt = self.in_proj(hidden_states,
-                              lora_params=lora_params,
-                              layer_idx=self.layer_idx)
-
-        # Split z and dt with views.
-        z = zxbcdt[:, :self.tp_d_inner]
+        # Split dt with views (z is consumed by the gated norm in forward).
         dt = zxbcdt[:, self.tp_d_inner + self.tp_conv_dim:]
-        dt_p, dt_d = torch.split(dt, seqlen_split_size, dim=0)
+        # Slice instead of exact-sum split: under piecewise CUDA graphs the
+        # token dim is padded to the capture bucket, so hidden_states can
+        # carry more rows than num_actual_tokens; the pad tail belongs to
+        # neither the prefill nor the decode segment.
+        dt_p = dt[:num_prefill_tokens]
+        dt_d = dt[num_prefill_tokens:num_actual_tokens]
 
         # Decode path uses regular view since no transpose is needed.
         xbc_d = zxbcdt[num_prefill_tokens:num_actual_tokens,
                        self.tp_d_inner:self.tp_d_inner + self.tp_conv_dim]
 
-        # Preallocate output tensor to avoid memcpy cost for merging prefill
-        # and decode outputs
-        preallocated_ssm_out = torch.empty(
-            [zxbcdt.shape[0], (self.tp_nheads * self.head_dim)],
-            dtype=zxbcdt.dtype,
-            device=zxbcdt.device,
-        )
-        preallocated_ssm_out_p, preallocated_ssm_out_d = torch.split(
-            preallocated_ssm_out,
-            [num_prefill_tokens, num_decode_tokens],
-            dim=0,
-        )
+        # Zero the pad tail (the buffer comes from torch.empty in forward) so
+        # the full-length gated norm sees defined values in pad rows; no-op
+        # slice when not padded.
+        preallocated_ssm_out[num_actual_tokens:].zero_()
+        preallocated_ssm_out_p = preallocated_ssm_out[:num_prefill_tokens]
+        preallocated_ssm_out_d = preallocated_ssm_out[
+            num_prefill_tokens:num_actual_tokens]
 
         if num_prefills > 0:
 
@@ -752,26 +925,43 @@ def convert_dt():
                     ssu_kwargs['rand_seed'] = rand_seed[:1]
                     ssu_kwargs['philox_rounds'] = self._philox_rounds
 
-                self.selective_state_update_func(
-                    ssm_states,
-                    x_d,
-                    dt_d,
-                    A,
-                    B_d,
-                    C_d,
-                    D,
-                    **ssu_kwargs,
-                )
-
-        # norm
-        hidden_states = self.norm(preallocated_ssm_out, z[:num_actual_tokens])
-
-        # out_proj
-        out = self.out_proj(hidden_states,
-                            lora_params=lora_params,
-                            layer_idx=self.layer_idx)
-
-        return out[:num_actual_tokens]
+                # Only while dynamo itself traces forward_core, i.e. a mixer
+                # built without a model config. Registered layers run this
+                # body eagerly inside the boundary op, where the plain call is
+                # the right one; the engine-level is_torch_compiling() flag
+                # stays set for the whole run and is not the test here.
+                if self._use_flashinfer and torch.compiler.is_compiling():
+                    # Route through the opaque custom op (see its comment);
+                    # `out` is a fresh intermediate copied into the
+                    # preallocated view afterwards so the op itself has no
+                    # mutable view args.
+                    ssu_out = torch.empty_like(x_d)
+                    torch.ops.trtllm.flashinfer_selective_state_update(
+                        ssm_states,
+                        x_d,
+                        dt_d,
+                        A,
+                        B_d,
+                        C_d,
+                        D,
+                        ssu_out,
+                        dt_bias=dt_bias,
+                        dt_softplus=self.delta_softplus,
+                        state_batch_indices=state_indices_d,
+                        rand_seed=ssu_kwargs.get('rand_seed'),
+                        philox_rounds=self._philox_rounds)
+                    ssu_kwargs['out'].copy_(ssu_out)
+                else:
+                    self.selective_state_update_func(
+                        ssm_states,
+                        x_d,
+                        dt_d,
+                        A,
+                        B_d,
+                        C_d,
+                        D,
+                        **ssu_kwargs,
+                    )
 
 
 # We want to cache the largest indexing vector we'd ever need and mask it, vs

tests/integration/test_lists/test-db/l0_dgx_h100.yml

@@ -38,6 +38,7 @@ l0_dgx_h100:
   - disaggregated/test_disaggregated.py::test_disaggregated_cancel_large_context_requests[DeepSeek-V3-Lite-bf16]
   # ------------- Model specific tests ---------------
   - accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B::test_w4a8[tp2_ep1]
+  - unittest/_torch/modeling/test_modeling_nemotron_h.py::test_nemotron_h_breakable_prefill_cuda_graph[tp2] TIMEOUT (90)
 - condition:
     ranges:
       system_gpu_count:
@@ -267,6 +268,7 @@ l0_dgx_h100:
     - unittest/_torch/ray_orchestrator/multi_gpu/test_ops.py -m "part3"
     - unittest/_torch/ray_orchestrator/multi_gpu/test_ops.py -m "part4"
     - unittest/llmapi/test_llm_multi_gpu_pytorch.py -m "gpu2"
+    - unittest/_torch/modeling/test_modeling_nemotron_h.py::test_nemotron_h_breakable_prefill_cuda_graph -k "tp2" TIMEOUT (90)
     - unittest/llmapi/test_async_llm.py -m "gpu2"
     - examples/test_ray.py::test_llm_inference_distributed_ray[tp2]
     - examples/test_ray.py::test_llm_inference_distributed_ray[pp2]

tests/integration/test_lists/test-db/l0_h100.yml

@@ -83,6 +83,8 @@ l0_h100:
   # reuse dense-8B greedy/logprob equality; accuracy stays on GSM8K/MMLU).
   - unittest/_torch/modeling/test_modeling_nemotron_h.py::test_nemotron_h_cuda_graph_overlap_scheduler
   - unittest/_torch/modeling/test_modeling_nemotron_h.py::test_nemotron_h_chunked_prefill
+  # Breakable prefill CUDA graphs (BCG) on the hybrid Mamba path; mpi_ray_parity (Ray id runs in the ray stage below).
+  - unittest/_torch/modeling/test_modeling_nemotron_h.py::test_nemotron_h_breakable_prefill_cuda_graph[tp1] TIMEOUT (90)
   - unittest/_torch/modeling/test_multimodal_encoder_graph.py
   # Qwen3.5-MoE-VL is hybrid (Mamba SSM + attention); FlashInfer's
   # chunk_gated_delta_rule GDN prefill kernel is sm90+ only, so this
@@ -245,6 +247,7 @@ l0_h100:
     - disaggregated/test_disaggregated_single_gpu.py::test_disaggregated_logprobs[False-TinyLlama-1.1B-Chat-v1.0]
     - disaggregated/test_disaggregated_single_gpu.py::test_disaggregated_logprobs[True-TinyLlama-1.1B-Chat-v1.0]
     - unittest/_torch/executor/test_overlap_scheduler.py
+    - unittest/_torch/modeling/test_modeling_nemotron_h.py::test_nemotron_h_breakable_prefill_cuda_graph -k "tp1" TIMEOUT (90)
     - unittest/executor/test_shim_ray.py
     - unittest/_torch/ray_orchestrator/single_gpu/test_llm_sleep.py
     - unittest/_torch/ray_orchestrator/single_gpu/test_llm_update_weights.py -m "part0"