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

2026-07-21 00:03:43 -07:00
parent 37a830b667
commit 429f6b6d15
15 changed files with 36 additions and 23 deletions
@@ -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(
-            graph,
-            pool=self._pool,
-            stream=self._capture_stream,
-            auto_dispatch_capture=True,
+        with (
+            skip_guard_context,
+            graph_ctx(
+                graph,
+                pool=self._pool,
+                stream=self._capture_stream,
+                auto_dispatch_capture=True,
+            ),
        ):
            out = forward_fn()

@@ -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):
@@ -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,
                )

@@ -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,
        )

@@ -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: ...

@@ -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
@@ -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.
@@ -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.
@@ -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",
@@ -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,
            )

@@ -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,
                )

@@ -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:
@@ -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
            ),
@@ -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,
@@ -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",