diff --git a/python/sglang/srt/hardware_backend/npu/graph_runner/npu_cudagraph_backend.py b/python/sglang/srt/hardware_backend/npu/graph_runner/npu_cudagraph_backend.py index 919f46619..a0339d8fb 100644 --- a/python/sglang/srt/hardware_backend/npu/graph_runner/npu_cudagraph_backend.py +++ b/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( - 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() diff --git a/python/sglang/srt/hardware_backend/xpu/graph_runner/xpu_full_graph_backend.py b/python/sglang/srt/hardware_backend/xpu/graph_runner/xpu_full_graph_backend.py index 6399c9df3..9c59fccf7 100644 --- a/python/sglang/srt/hardware_backend/xpu/graph_runner/xpu_full_graph_backend.py +++ b/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): diff --git a/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py index ec1275999..60f4cab53 100644 --- a/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py +++ b/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, ) diff --git a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py index aa07a2754..a8852226f 100644 --- a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py +++ b/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, ) diff --git a/python/sglang/srt/model_executor/runner_backend/base_cuda_graph_backend.py b/python/sglang/srt/model_executor/runner_backend/base_cuda_graph_backend.py index 5b603da97..ce8e1e34c 100644 --- a/python/sglang/srt/model_executor/runner_backend/base_cuda_graph_backend.py +++ b/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: ... diff --git a/python/sglang/srt/model_executor/runner_backend/breakable_cuda_graph_backend.py b/python/sglang/srt/model_executor/runner_backend/breakable_cuda_graph_backend.py index 4864bbd67..6235da423 100644 --- a/python/sglang/srt/model_executor/runner_backend/breakable_cuda_graph_backend.py +++ b/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 diff --git a/python/sglang/srt/model_executor/runner_backend/full_cuda_graph_backend.py b/python/sglang/srt/model_executor/runner_backend/full_cuda_graph_backend.py index e475ed9ca..e4f13381e 100644 --- a/python/sglang/srt/model_executor/runner_backend/full_cuda_graph_backend.py +++ b/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. diff --git a/python/sglang/srt/model_executor/runner_backend/tc_piecewise_cuda_graph_backend.py b/python/sglang/srt/model_executor/runner_backend/tc_piecewise_cuda_graph_backend.py index c1e0d681a..a416d0c74 100644 --- a/python/sglang/srt/model_executor/runner_backend/tc_piecewise_cuda_graph_backend.py +++ b/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. diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index f20de4419..5c77b66d4 100644 --- a/python/sglang/srt/server_args.py +++ b/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", diff --git a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py b/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py index 5f1f46894..0e47b2546 100644 --- a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py +++ b/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, ) diff --git a/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py b/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py index 183997310..1eea9b969 100644 --- a/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py +++ b/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, ) diff --git a/python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py b/python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py index 05486883b..8eb3d18bc 100644 --- a/python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py +++ b/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: diff --git a/python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py b/python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py index 3cbc9516b..db21f9e06 100644 --- a/python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py +++ b/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 ), diff --git a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py b/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py index fb2e1de93..d451ac0ee 100644 --- a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py +++ b/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, diff --git a/test/registered/models_e2e/test_mimo_v2_flash.py b/test/registered/models_e2e/test_mimo_v2_flash.py index a5a7e4ba2..e431b0f7a 100644 --- a/test/registered/models_e2e/test_mimo_v2_flash.py +++ b/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",