Turn on breakable prefill cuda graph for dp attention by default (#31682)

python/sglang/srt/hardware_backend/npu/graph_runner/npu_cudagraph_backend.py

@@ -78,7 +78,7 @@ class NPUCudaGraphBackend(BaseCudaGraphBackend):
         self,
         shape_key: ShapeKey,
         forward_fn: Callable[[], Any],
-        dummies: Optional[Any] = None,
+        capture_inputs: Optional[Any] = None,
         post_warmup_hook: Optional[Callable[[], None]] = None,
     ) -> None:
         import torch_npu  # noqa: F401  (verifies NPU availability)
@@ -111,11 +111,14 @@ class NPUCudaGraphBackend(BaseCudaGraphBackend):
         else:
             graph_ctx = torch.npu.graph
 
-        with skip_guard_context, graph_ctx(
+        with (
-            graph,
+            skip_guard_context,
-            pool=self._pool,
+            graph_ctx(
-            stream=self._capture_stream,
+                graph,
-            auto_dispatch_capture=True,
+                pool=self._pool,
+                stream=self._capture_stream,
+                auto_dispatch_capture=True,
+            ),
         ):
             out = forward_fn()
 

python/sglang/srt/hardware_backend/xpu/graph_runner/xpu_full_graph_backend.py

@@ -58,7 +58,7 @@ class FullXPUGraphBackend(BaseCudaGraphBackend):
         self,
         shape_key: ShapeKey,
         forward_fn: Callable[[], Any],
-        dummies: Optional[Any] = None,
+        capture_inputs: Optional[Any] = None,
         post_warmup_hook: Optional[Callable[[], None]] = None,
     ) -> None:
         for _ in range(2):

python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py

@@ -1018,7 +1018,7 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
                 self.backend.capture_one(
                     shape_key,
                     run_once,
-                    dummies=None,
+                    capture_inputs=None,
                     post_warmup_hook=post_warmup_hook,
                 )
 

python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py

@@ -914,7 +914,9 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
         self.backend.capture_one(
             ShapeKey(size=num_tokens),
             run_once,
-            dummies=None,
+            # DP padding can install capture-only tensors on this dummy batch;
+            # BCG retains it so their recorded addresses remain valid.
+            capture_inputs=forward_batch,
             post_warmup_hook=post_warmup_hook,
         )
 

python/sglang/srt/model_executor/runner_backend/base_cuda_graph_backend.py

@@ -34,7 +34,7 @@ class BaseCudaGraphBackend(ABC):
       - capture_session(stream) — context wrapping the runner's outer
         capture loop; backends bind stream / pool and open per-backend
         capture flags here.
-      - capture_one(shape_key, forward_fn, dummies, post_warmup_hook)
+      - capture_one(shape_key, forward_fn, capture_inputs, post_warmup_hook)
         — record the replayable artifact for shape_key; one call per
         shape inside capture_session.
       - can_run(forward_batch, shape_key) — can this backend replay
@@ -49,6 +49,8 @@ class BaseCudaGraphBackend(ABC):
     Notes:
       - The outer capture loop is runner-specific; it lives on the
         runner, not here.
+      - capture_inputs optionally carries capture-time input owners that a
+        backend must retain when its graph records their tensor addresses.
     """
 
     @abstractmethod
@@ -59,7 +61,7 @@ class BaseCudaGraphBackend(ABC):
         self,
         shape_key: ShapeKey,
         forward_fn,
-        dummies: Optional[Any] = None,
+        capture_inputs: Optional[Any] = None,
         post_warmup_hook: Optional[Callable[[], None]] = None,
     ) -> None: ...
 

python/sglang/srt/model_executor/runner_backend/breakable_cuda_graph_backend.py

@@ -68,6 +68,7 @@ class BreakableCudaGraphBackend(DedupedCudaGraphMixin, BaseCudaGraphBackend):
         self._model_runner = cuda_graph_runner.model_runner
         self._graphs: Dict[Any, BreakableCUDAGraph] = {}
         self._outputs: Dict[Any, Any] = {}
+        self._capture_inputs: Dict[Any, Any] = {}
         self._pool = None
         self._device_module = cuda_graph_runner.device_module
         self._tp_group = cuda_graph_runner.model_runner.tp_group
@@ -107,7 +108,7 @@ class BreakableCudaGraphBackend(DedupedCudaGraphMixin, BaseCudaGraphBackend):
         self,
         shape_key: ShapeKey,
         forward_fn: Callable[[], Any],
-        dummies: Optional[Any] = None,
+        capture_inputs: Optional[Any] = None,
         post_warmup_hook: Optional[Callable[[], None]] = None,
     ) -> None:
         warmup_out = None
@@ -137,6 +138,8 @@ class BreakableCudaGraphBackend(DedupedCudaGraphMixin, BaseCudaGraphBackend):
         stored = self._slice_output(self._shared_output_buffer, out_rows)
         self._graphs[shape_key] = graph
         self._outputs[shape_key] = stored
+        # CUDA graphs retain tensor addresses, not Python tensor lifetimes.
+        self._capture_inputs[shape_key] = capture_inputs
 
     def _output_rows(self, output: Any, cap: int) -> int:
         """Leading-dim row count actually produced by the body, clamped to ``cap``.
@@ -248,5 +251,6 @@ class BreakableCudaGraphBackend(DedupedCudaGraphMixin, BaseCudaGraphBackend):
         self.close()
         self._graphs.clear()
         self._outputs.clear()
+        self._capture_inputs.clear()
         self._pool = None
         self._shared_output_buffer = None

python/sglang/srt/model_executor/runner_backend/full_cuda_graph_backend.py

@@ -81,7 +81,7 @@ class FullCudaGraphBackend(BaseCudaGraphBackend):
         self,
         shape_key: ShapeKey,
         forward_fn: Callable[[], Any],
-        dummies: Optional[Any] = None,
+        capture_inputs: Optional[Any] = None,
         post_warmup_hook: Optional[Callable[[], None]] = None,
     ) -> None:
         # Two warmups so kernels are loaded and one-time setup is paid before capture.

python/sglang/srt/model_executor/runner_backend/tc_piecewise_cuda_graph_backend.py

@@ -231,7 +231,7 @@ class TcPiecewiseCudaGraphBackend(BaseCudaGraphBackend):
         self,
         shape_key: ShapeKey,
         forward_fn: Callable[[], Any],
-        dummies: Optional[Any] = None,
+        capture_inputs: Optional[Any] = None,
         post_warmup_hook: Optional[Callable[[], None]] = None,
     ) -> None:
         # Call 1 warms FX state; call 2 captures the cuda graph inside capture_session.

python/sglang/srt/server_args.py

@@ -3776,8 +3776,6 @@ class ServerArgs:
                 "MoE A2A backend",
                 lambda: _resolved_view(self).moe_a2a_backend != "none",
             ),
-            # DP-attn × BCG capture/replay not yet validated.
-            ("DP attention", lambda: self._resolved().enable_dp_attention),
             # Multimodal prefill replay faults under BCG; allowlisted archs opt back in.
             (
                 "multimodal model",

python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py

@@ -484,7 +484,7 @@ class EAGLEDraftCudaGraphRunner(DecodeCudaGraphRunner):
             self.backend.capture_one(
                 shape_key,
                 run_once,
-                dummies=None,
+                capture_inputs=None,
                 post_warmup_hook=post_warmup_hook,
             )
 

python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py

@@ -478,7 +478,7 @@ class EAGLEDraftExtendCudaGraphRunner(DecodeCudaGraphRunner):
                 self.backend.capture_one(
                     shape_key,
                     run_once,
-                    dummies=None,
+                    capture_inputs=None,
                     post_warmup_hook=post_warmup_hook,
                 )
 

python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py

@@ -365,7 +365,7 @@ class FrozenKVMTPCudaGraphRunner(DecodeCudaGraphRunner):
                 self.backend.capture_one(
                     shape_key,
                     run_once,
-                    dummies=None,
+                    capture_inputs=None,
                     post_warmup_hook=post_warmup_hook,
                 )
         finally:

python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py

@@ -447,7 +447,7 @@ class MultiLayerEagleDraftExtendCudaGraphRunner(DecodeCudaGraphRunner):
             self.backend.capture_one(
                 shape_key,
                 run_once,
-                dummies=None,
+                capture_inputs=None,
                 post_warmup_hook=post_warmup_hook,
             )
 
@@ -943,7 +943,7 @@ class OneGraphMultiLayerEagleMultiStepDraftExtendCudaGraphRunner(
         first.backend.capture_one(
             shape_key,
             multi_step_fn,
-            dummies=None,
+            capture_inputs=None,
             post_warmup_hook=getattr(
                 first.attn_backend, "on_after_cuda_graph_warmup", None
             ),

python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py

@@ -608,8 +608,11 @@ class MultiLayerEagleDraftWorker(EagleDraftWorkerBase):
             forward_batch.token_to_kv_pool = self.draft_runner_list[
                 step
             ].token_to_kv_pool
+            # DP/MLP-sync padding mutates ForwardBatch fields in place. Keep
+            # those per-runner mutations from leaking into the next MTP step.
+            step_forward_batch = replace(forward_batch)
             output: ModelRunnerOutput = self.draft_runner_list[step].forward(
-                forward_batch
+                step_forward_batch
             )
             maybe_detect_nan(
                 output.logits_output.next_token_logits,

test/registered/models_e2e/test_mimo_v2_flash.py

@@ -1,6 +1,7 @@
 import unittest
 
 from sglang.srt.environ import envs
+from sglang.srt.utils import is_blackwell
 from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.kits.eval_accuracy_kit import GSM8KMixin
 from sglang.test.kits.spec_decoding_kit import SpecDecodingMixin
@@ -23,7 +24,7 @@ class TestMiMoV2Flash(GSM8KMixin, SpecDecodingMixin, DefaultServerBase):
         "--enable-dp-attention",
         "--trust-remote-code",
         "--attention-backend",
-        "fa3",
+        "fa4" if is_blackwell() else "fa3",
         "--max-running-requests",
         "128",
         "--cuda-graph-max-bs-decode",