[BCG][3/N] Enable bcg on dsa & deepep a2a backend (#31987)

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

python/sglang/benchmark/one_batch.py

@@ -536,6 +536,7 @@ def _maybe_prepare_mlp_sync_batch(batch: ScheduleBatch, model_runner):
     if require_mlp_sync(model_runner.server_args):
         prepare_mlp_sync_batch_raw(
             batch,
+            model_runner=model_runner,
             dp_size=model_runner.server_args.dp_size,
             attn_tp_size=get_parallel().attn_tp_size,
             attn_cp_size=model_runner.ps.attn_cp_size,

python/sglang/srt/layers/moe/ep_moe/layer.py

@@ -8,6 +8,10 @@ import torch
 from sglang.kernels.ops.quantization.fp8_kernel import is_fp8_fnuz
 from sglang.srt.environ import envs
 from sglang.srt.layers import deep_gemm_wrapper
+from sglang.srt.layers.dp_attention import (
+    get_is_extend_in_batch,
+    set_is_extend_in_batch,
+)
 from sglang.srt.layers.moe import (
     get_deepep_mode,
     get_moe_a2a_backend,
@@ -21,10 +25,20 @@ from sglang.srt.layers.moe.token_dispatcher.deepep import (
     DeepEPLLCombineInput,
     DeepEPNormalCombineInput,
 )
-from sglang.srt.layers.moe.topk import TopKOutput, TopKOutputChecker
+from sglang.srt.layers.moe.topk import (
+    StandardTopKOutput,
+    TopKOutput,
+    TopKOutputChecker,
+)
 from sglang.srt.layers.quantization.base_config import QuantizationConfig
 from sglang.srt.layers.quantization.fp8 import Fp8Config
 from sglang.srt.layers.quantization.w4afp8 import W4AFp8Config, W4AFp8MoEMethod
+from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph import (
+    eager_on_graph,
+)
+from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.context import (
+    is_in_breakable_cuda_graph,
+)
 from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
     is_in_tc_piecewise_cuda_graph,
 )
@@ -155,11 +169,62 @@ class DeepEPMoE(FusedMoE):
                 deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM
             ), f"DeepEP {self.deepep_mode} mode requires deep_gemm"
 
+    def _a2a_forward_with_output_impl(
+        self,
+        hidden_states: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        router_logits: torch.Tensor,
+        output: torch.Tensor,
+    ) -> None:
+        # eager run under breakable cuda graph
+        saved_is_extend_in_batch = get_is_extend_in_batch()
+        set_is_extend_in_batch(True)
+        try:
+            output.copy_(
+                self.forward_impl(
+                    hidden_states,
+                    StandardTopKOutput(topk_weights, topk_ids, router_logits),
+                )
+            )
+        finally:
+            set_is_extend_in_batch(saved_is_extend_in_batch)
+
+    def _a2a_forward_capture_stub(
+        self,
+        hidden_states: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        router_logits: torch.Tensor,
+        output: torch.Tensor,
+    ) -> None:
+        # Capture pass only: record the buffer address, skip the
+        # rank-coupled a2a. Warmup and replay run the real body.
+        output.zero_()
+
+    a2a_forward_with_output = eager_on_graph(
+        True, capture_stub=_a2a_forward_capture_stub
+    )(_a2a_forward_with_output_impl)
+
     def forward(
         self,
         hidden_states: torch.Tensor,
         topk_output: TopKOutput,
     ):
+        # DeepEP NORMAL mode is not capturable; run it as an eager node.
+        if is_in_breakable_cuda_graph():
+            assert TopKOutputChecker.format_is_standard(
+                topk_output
+            ), "Only standard topk output is supported for breakable cuda graph"
+            output = torch.empty_like(hidden_states)
+            self.a2a_forward_with_output(
+                hidden_states,
+                topk_output.topk_weights,
+                topk_output.topk_ids,
+                topk_output.router_logits,
+                output,
+            )
+            return output
         if is_in_tc_piecewise_cuda_graph():
             assert TopKOutputChecker.format_is_standard(
                 topk_output

python/sglang/srt/managers/scheduler.py

@@ -269,6 +269,7 @@ from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo
 from sglang.srt.sampling.sampling_params import TOP_K_ALL
 from sglang.srt.server_args import PortArgs, ServerArgs
 from sglang.srt.session.session_controller import SessionController
+from sglang.srt.speculative.base_spec_worker import BaseSpecWorker
 from sglang.srt.speculative.dflash_utils import validate_dflash_request
 from sglang.srt.speculative.eagle_utils import get_draft_recurrent_hidden_state_spec
 from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
@@ -1882,7 +1883,15 @@ class Scheduler(
         )
 
     def init_dp_attn_adapter(self) -> None:
+        # Spec workers have no .model_runner of their own; the prefill graph
+        # runner that votes belongs to the target model.
+        target_worker = (
+            self.tp_worker.target_worker
+            if isinstance(self.tp_worker, BaseSpecWorker)
+            else self.tp_worker
+        )
         self.dp_attn_adapter = SchedulerDPAttnAdapter(
+            model_runner=target_worker.model_runner,
             tp_group=self.tp_group,
             req_to_token_pool=self.req_to_token_pool,
             token_to_kv_pool_allocator=self.token_to_kv_pool_allocator,

python/sglang/srt/managers/scheduler_components/dp_attn.py

@@ -33,6 +33,7 @@ from sglang.srt.utils.common import require_mlp_tp_gather
 
 if TYPE_CHECKING:
     from sglang.srt.distributed.parallel_state import GroupCoordinator
+    from sglang.srt.model_executor.model_runner import ModelRunner
 
 
 _ENABLE_METRICS_DP_ATTENTION = envs.SGLANG_ENABLE_METRICS_DP_ATTENTION.get()
@@ -83,11 +84,11 @@ class MLPSyncBatchInfo:
 
     num_tokens: int
     num_tokens_for_logprob: int
-    can_cuda_graph: bool
+    can_run_decode_cuda_graph: bool
+    can_run_prefill_cuda_graph: bool
     is_extend_in_batch: bool
     local_can_run_tbo: bool
     local_forward_mode: int
-    can_run_breakable_cuda_graph: bool
 
     # some gathered elements
     tp0_info: torch.Tensor = None
@@ -102,11 +103,11 @@ class MLPSyncBatchInfo:
             [
                 self.num_tokens,
                 self.num_tokens_for_logprob,
-                int(self.can_cuda_graph),
+                int(self.can_run_decode_cuda_graph),
                 int(self.is_extend_in_batch),
                 int(self.local_can_run_tbo),
                 self.local_forward_mode,
-                int(self.can_run_breakable_cuda_graph),
+                int(self.can_run_prefill_cuda_graph),
             ],
             device=device,
             dtype=dtype,
@@ -117,11 +118,11 @@ class MLPSyncBatchInfo:
             [
                 0,  # num_tokens
                 0,  # num_tokens_for_logprob
-                1,  # can_cuda_graph
+                1,  # can_run_decode_cuda_graph
                 0,  # is_extend_in_batch
                 1,  # local_can_run_tbo
                 ForwardMode.IDLE.value,  # local_forward_mode
-                0,  # can_run_breakable_cuda_graph
+                0,  # can_run_prefill_cuda_graph
             ],
             device=device,
             dtype=dtype,
@@ -184,9 +185,9 @@ class MLPSyncBatchInfo:
         cpu_data = tp0_info[:, :2].cpu()
         self.global_num_tokens = cpu_data[:, 0].tolist()
         self.global_num_tokens_for_logprob = cpu_data[:, 1].tolist()
-        self.can_cuda_graph = bool(tp0_info[:, 2].min().item())
+        self.can_run_decode_cuda_graph = bool(tp0_info[:, 2].min().item())
         self.is_extend_in_batch = bool(tp0_info[:, 3].max().item())
-        self.can_run_breakable_cuda_graph = bool(tp0_info[:, 6].min().item())
+        self.can_run_prefill_cuda_graph = bool(tp0_info[:, 6].min().item())
         if _ENABLE_METRICS_DP_ATTENTION:
             self.dp_cooperation_info = DPCooperationInfo.create(tp0_info[:, 5].tolist())
 
@@ -212,12 +213,13 @@ def _update_gather_batch(
         batch.global_forward_mode = mlp_sync_info.global_forward_mode
 
     # Check forward mode for cuda graph
-    batch.can_run_dp_cuda_graph = mlp_sync_info.can_cuda_graph
-    batch.can_run_dp_breakable_cuda_graph = mlp_sync_info.can_run_breakable_cuda_graph
+    batch.can_run_dp_cuda_graph = mlp_sync_info.can_run_decode_cuda_graph
+    batch.can_run_dp_breakable_cuda_graph = mlp_sync_info.can_run_prefill_cuda_graph
 
 
 def prepare_mlp_sync_batch_raw(
     local_batch: ScheduleBatch,
+    model_runner: ModelRunner,
     dp_size: int,
     attn_tp_size: int,
     attn_cp_size: int,
@@ -256,19 +258,38 @@ def prepare_mlp_sync_batch_raw(
         )
 
     skip_all_gather = envs.SGLANG_SCHEDULER_SKIP_ALL_GATHER.get()
-    can_cuda_graph = (
+    can_run_decode_cuda_graph = (
         local_batch is None
         or local_batch.forward_mode.is_decode_or_idle()
         or local_batch.forward_mode.is_prebuilt()
     ) and not disable_cuda_graph
-    # Idle/None ranks are permissive (like can_cuda_graph): the all-gather
-    # min()-reduces this across DP ranks, so a prefill batch with idle ranks
-    # still resolves to True (idle ranks become a padded dummy extend).
-    can_run_breakable_cuda_graph = (
+    breakable_prefill = check_cuda_graph_backend(Phase.PREFILL, Backend.BREAKABLE)
+    prefill_graph_runner = (
+        model_runner.prefill_cuda_graph_runner if breakable_prefill else None
+    )
+    can_run_prefill_cuda_graph = (
         local_batch is None
         or local_batch.forward_mode.is_idle()
-        or local_batch.forward_mode in (ForwardMode.EXTEND, ForwardMode.MIXED)
-    ) and check_cuda_graph_backend(Phase.PREFILL, Backend.BREAKABLE)
+        # Breakable Cuda Graph Backend Check.
+        or (
+            local_batch.forward_mode in (ForwardMode.EXTEND, ForwardMode.MIXED)
+            and (
+                prefill_graph_runner is None
+                or prefill_graph_runner.can_replay_locally(
+                    batch_size=local_batch.batch_size(),
+                    num_tokens=local_batch.extend_num_tokens,
+                    input_embeds=local_batch.input_embeds,
+                    replace_embeds=None,
+                    prefix_lens=local_batch.prefix_lens,
+                    is_target_verify=local_batch.forward_mode.is_target_verify(),
+                    capture_hidden_mode=None,
+                    return_logprob=local_batch.return_logprob,
+                    lora_ineligible=prefill_graph_runner.enable_lora,
+                )
+            )
+            and breakable_prefill
+        )
+    )
 
     is_extend_in_batch = local_batch.forward_mode.is_extend() if local_batch else False
     if local_batch is not None:
@@ -307,11 +328,11 @@ def prepare_mlp_sync_batch_raw(
         cp_size=attn_cp_size,
         num_tokens=num_tokens,
         num_tokens_for_logprob=num_tokens_for_logprob,
-        can_cuda_graph=can_cuda_graph,
+        can_run_decode_cuda_graph=can_run_decode_cuda_graph,
+        can_run_prefill_cuda_graph=can_run_prefill_cuda_graph,
         is_extend_in_batch=is_extend_in_batch,
         local_can_run_tbo=local_can_run_tbo,
         local_forward_mode=local_forward_mode,
-        can_run_breakable_cuda_graph=can_run_breakable_cuda_graph,
     )
 
     if not skip_all_gather:
@@ -364,6 +385,7 @@ def prepare_mlp_sync_batch_raw(
 
 @dataclass(kw_only=True, slots=True, frozen=True)
 class SchedulerDPAttnAdapter:
+    model_runner: ModelRunner
     tp_group: GroupCoordinator
     req_to_token_pool: ReqToTokenPool
     token_to_kv_pool_allocator: BaseTokenToKVPoolAllocator
@@ -379,6 +401,7 @@ class SchedulerDPAttnAdapter:
     def prepare_mlp_sync_batch(self, local_batch: ScheduleBatch):
         return prepare_mlp_sync_batch_raw(
             local_batch,
+            model_runner=self.model_runner,
             dp_size=self.server_args.dp_size,
             attn_tp_size=self.ps.attn_tp_size,
             attn_cp_size=self.ps.attn_cp_size,

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

@@ -50,6 +50,7 @@ import tqdm
 from sglang.kernels.ops.kvcache.kv_indices import (
     create_chunked_prefix_cache_kv_indices,
 )
+from sglang.srt.configs.model_config import is_deepseek_dsa
 from sglang.srt.distributed.parallel_state import graph_capture
 from sglang.srt.layers.attention.dsa.utils import is_dsa_enable_prefill_cp
 from sglang.srt.layers.dp_attention import (
@@ -236,6 +237,11 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
     buffer population, attention metadata init, and output slicing.
     """
 
+    # DSA forces use_mha=False in BCG capture/replay, so the sparse path
+    # serves any prefix and the MHA-prefix ban does not apply. Class
+    # default keeps __new__-built test instances on the ban.
+    dsa_sparse_prefill_forced: bool = False
+
     def __init__(self, model_runner: ModelRunner):
         super().__init__(model_runner)
         # --- model flags ----------------------------------------------
@@ -316,6 +322,10 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
             source=self.buffers,
         )
 
+        self.dsa_sparse_prefill_forced = is_deepseek_dsa(
+            self.model_runner.model_config.hf_config
+        )
+
         self.attention_layers = self.model_runner.attention_layers
         self.mha_companion_layers = self.model_runner.mha_companion_layers
         self.has_mha_companion_layers = any(
@@ -986,14 +996,6 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
             static_forward_batch=static_forward_batch,
         )
 
-    def _has_unsupported_mha_prefix(self, forward_batch: ForwardBatch) -> bool:
-        return (
-            self.prefill_backend_name == Backend.BREAKABLE
-            and self.has_mha_companion_layers
-            and forward_batch.extend_prefix_lens_cpu is not None
-            and any(forward_batch.extend_prefix_lens_cpu)
-        )
-
     @staticmethod
     def _restore_mha_capture_state(forward_batch: ForwardBatch) -> None:
         """Restore Python state omitted from breakable graph segments."""
@@ -1001,59 +1003,113 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
         forward_batch.mha_return_lse = False
         forward_batch.set_attn_attend_prefix_cache(False)
 
-    def can_run_graph(self, forward_batch: ForwardBatch) -> bool:
-        if self._is_full_backend and forward_batch.batch_size > self._capture_req_slots:
+    def can_replay_locally(
+        self,
+        *,
+        batch_size: int,
+        num_tokens: Optional[int],
+        input_embeds,
+        replace_embeds,
+        prefix_lens,
+        is_target_verify: bool,
+        capture_hidden_mode,
+        return_logprob: bool,
+        lora_ineligible: bool = False,
+        chunked_prefix_uncapturable: bool = False,
+    ) -> bool:
+        """Rank-local replay eligibility: the single source of truth for
+        ``can_run_graph`` (ForwardBatch, forward time) and the dp mlp-sync
+        vote (ScheduleBatch, schedule time) — all dp ranks must reach the
+        same replay-vs-eager decision or their collectives mismatch. Pass
+        ``capture_hidden_mode=None`` when unknown at the call site (it is
+        rank-uniform; forward-time-only checking cannot split the group).
+        """
+        if self._is_full_backend and batch_size > self._capture_req_slots:
             return False
-        # LoRA batches may only replay the graph when prepare_lora_batch put
-        # their metadata in the static buffers (same predicate); keyed off
-        # enable_lora, not lora_ids, which is non-None even without LoRA.
-        if self.enable_lora and not (
-            self._capture_lora
-            and self.model_runner.lora_manager.can_use_prefill_cuda_graph(forward_batch)
+        # LoRA replays need prepare_lora_batch's static metadata. lora_manager
+        # keeps LoRA prefill eager on every rank under dp attention, so the
+        # schedule-time vote derives this from enable_lora alone.
+        if lora_ineligible:
+            return False
+        if input_embeds is not None:
+            return False
+        if replace_embeds is not None:
+            return False
+        # A prefix forces the MHA companion path, whose captured state is
+        # frozen prefix-free; DSA models are exempt (capture/replay force
+        # the sparse path, which takes any prefix via device metadata).
+        if (
+            self.prefill_backend_name == Backend.BREAKABLE
+            and self.has_mha_companion_layers
+            and not self.dsa_sparse_prefill_forced
+            and prefix_lens is not None
+            and any(prefix_lens)
         ):
             return False
-        if forward_batch.input_embeds is not None:
-            return False
-        if forward_batch.replace_embeds is not None:
-            return False
-        if self._has_unsupported_mha_prefix(forward_batch):
+        # FullCG's chunked-prefix topology covers a bounded prefix. The flag
+        # gating it is FULL-backend-only, so this is inert for the breakable
+        # vote path.
+        if chunked_prefix_uncapturable:
             return False
         # tc_piecewise captures with ForwardMode.EXTEND and spec_info=None.
-        if forward_batch.forward_mode.is_target_verify():
+        if is_target_verify:
             return False
-        if forward_batch.capture_hidden_mode != self.capture_hidden_mode:
+        if (
+            capture_hidden_mode is not None
+            and capture_hidden_mode != self.capture_hidden_mode
+        ):
             return False
-        # BCG-with-captured-metadata under DP attention: every rank must
-        # have local tokens, and the batch must declare itself replayable.
-        # These gates are no-ops for non-DP / non-opt-in paths because
-        # global_num_tokens_cpu stays None.
-        if self._has_inactive_dp_rank(forward_batch):
+        if return_logprob and not self._uses_eager_prefill_tail():
             return False
+        if num_tokens is None:
+            return True
+        if num_tokens > self.max_num_tokens:
+            return False
+        # No exact-shape check: load_batch bucket-pads; only reject
+        # disproportionate padding waste.
+        padded_num_tokens = self._pad_to_bucket(num_tokens, self.capture_num_tokens)
+        if padded_num_tokens > num_tokens * _MAX_PREFILL_CUDA_GRAPH_PADDING_FACTOR:
+            return False
+        return True
+
+    def can_run_graph(self, forward_batch: ForwardBatch) -> bool:
+        # DP check: group verdict from the schedule-time all-gather
+        # (min-reduced votes; also requires every rank to hold tokens).
         if (
             forward_batch.global_num_tokens_cpu is not None
             and not forward_batch.can_run_dp_breakable_cuda_graph
         ):
             return False
-        num_tokens = len(forward_batch.input_ids)
-        if forward_batch.return_logprob and not self._uses_eager_prefill_tail():
+
+        # Every dp rank must hold tokens this forward (reads the synced
+        # table post dp-padding; idle ranks vote permissively upstream).
+        if self._has_inactive_dp_rank(forward_batch):
             return False
-        if num_tokens > self.max_num_tokens:
-            return False
-        padded_num_tokens = self._pad_to_bucket(num_tokens, self.capture_num_tokens)
-        if padded_num_tokens > num_tokens * _MAX_PREFILL_CUDA_GRAPH_PADDING_FACTOR:
-            return False
-        # Other backends and non-MLA FullCG keep using their normal graph with
-        # replay-refreshed metadata; only this extra topology has a prefix cap.
-        if (
-            self._capture_chunked_prefix
-            and self._has_prefix_hit(forward_batch)
-            and self._select_prefix_capture_chunks(forward_batch) is None
+
+        # Non-DP local check (sole decision for tp-only).
+        if not self.can_replay_locally(
+            batch_size=forward_batch.batch_size,
+            num_tokens=len(forward_batch.input_ids),
+            input_embeds=forward_batch.input_embeds,
+            replace_embeds=forward_batch.replace_embeds,
+            prefix_lens=forward_batch.extend_prefix_lens_cpu,
+            is_target_verify=forward_batch.forward_mode.is_target_verify(),
+            capture_hidden_mode=forward_batch.capture_hidden_mode,
+            return_logprob=forward_batch.return_logprob,
+            lora_ineligible=self.enable_lora
+            and not (
+                self._capture_lora
+                and self.model_runner.lora_manager.can_use_prefill_cuda_graph(
+                    forward_batch
+                )
+            ),
+            chunked_prefix_uncapturable=(
+                self._capture_chunked_prefix
+                and self._has_prefix_hit(forward_batch)
+                and self._select_prefix_capture_chunks(forward_batch) is None
+            ),
         ):
             return False
-        # load_batch bucket-pads to the nearest captured shape. The factor
-        # above rejects replays whose padded model work is disproportionate
-        # to the useful token count.
-        #
         # Multi-req replay is supported by body-capture backends via the
         # layer_model.forward monkey-patch in replay(): the captured graph runs
         # the transformer stack, then the outer model.forward runs

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

@@ -130,6 +130,7 @@ class BreakableCudaGraphBackend(DedupedCudaGraphMixin, BaseCudaGraphBackend):
             cuda_graph=graph,
             pool=self._pool,
             stream=self._capture_stream,
+            barrier_fn=self._tp_group.barrier,
         ):
             out = captured_fn()
             out_rows = self._output_rows(out, size)

python/sglang/srt/model_executor/runner_backend_utils/breakable_cuda_graph/breakable_cuda_graph.py

@@ -25,7 +25,7 @@ buffers to keep break-point tensors at stable addresses.
 import logging
 import threading
 from contextvars import ContextVar
-from typing import Any, Callable
+from typing import Any, Callable, Optional
 
 import torch
 
@@ -216,7 +216,7 @@ def _copy_output(dst: Any, src: Any) -> Any:
     return src
 
 
-def eager_on_graph(enable: bool):
+def eager_on_graph(enable: bool, capture_stub: Optional[Callable] = None):
     def decorator(inner: Callable):
         if not enable:
             return inner
@@ -231,9 +231,21 @@ def eager_on_graph(enable: bool):
             # End the segment that captured up to this break point.
             capture._end_current_segment()
 
-            # Run the eager function once so it allocates its outputs and
-            # writes real data into them.
-            output = inner(*args, **kwargs)
+            # Re-sync ranks after segment teardown (the slow, variable
+            # step) before break fns with rank-coupled collectives and hard
+            # timeouts (DeepEP NORMAL: 100s). Capture-only; replay bypasses
+            # this wrapper.
+            if capture._barrier_fn is not None:
+                capture._barrier_fn()
+
+            # Run the break once so its outputs are allocated and their
+            # addresses recorded. A capture_stub replaces the body during
+            # capture (contents are never consumed; warmup and replay run
+            # the real inner), letting rank-coupled bodies skip the work.
+            if capture_stub is not None:
+                output = capture_stub(*args, **kwargs)
+            else:
+                output = inner(*args, **kwargs)
 
             # Weak-ref captured inputs produced by graph segments. Their storage
             # is pinned by the segment CUDAGraphs' mempool use-count, so Python
@@ -308,6 +320,7 @@ class BreakableCUDAGraphCapture:
         pool=None,
         stream: torch.Stream | None = None,
         capture_error_mode: str = "global",
+        barrier_fn: Callable[[], None] | None = None,
     ):
         assert isinstance(
             cuda_graph, BreakableCUDAGraph
@@ -316,6 +329,7 @@ class BreakableCUDAGraphCapture:
         self._pool = pool if pool is not None else (0, 0)
         self._stream = stream
         self._capture_error_mode = capture_error_mode
+        self._barrier_fn = barrier_fn
         self._stream_ctx = None
         self._capture_token = None
         self._stream_token = None

python/sglang/srt/models/deepseek_v2.py

@@ -1270,6 +1270,13 @@ class DeepseekV2MoE(nn.Module):
                         shared_output = self._forward_shared_experts(hidden_states)
                         shared_output.record_stream(self.alt_stream)
                         shared_event = self.alt_stream.record_event()
+                    if is_in_breakable_cuda_graph():
+                        # The MoE call below is an eager break, so record
+                        # and wait must share one capture; joining here means
+                        # the shared experts overlap nothing. The alt stream
+                        # is kept for record_stream: without that marking the
+                        # allocator recycles shared_output across the break.
+                        torch.cuda.current_stream().wait_event(shared_event)
                 else:
                     shared_output = self._forward_shared_experts(hidden_states)
             topk_kwargs = (
@@ -1453,6 +1460,7 @@ class DeepseekV2MoE(nn.Module):
             and not sbo_enabled_flag
             and self.num_fused_shared_experts == 0
             and self.alt_stream is not None
+            and not is_in_breakable_cuda_graph()
         ):
             torch.cuda.current_stream().wait_event(shared_event)
 

python/sglang/srt/server_args.py

@@ -41,6 +41,7 @@ from sglang.srt.arg_groups.argparse_actions import (
     DeprecatedStoreTrueAction,
     LoRAPathAction,
 )
+from sglang.srt.arg_groups.overrides import resolved_view
 from sglang.srt.configs.embedding_model_spec import BCGPrefillPolicy
 from sglang.srt.configs.linear_attn_model_registry import get_linear_attn_spec_by_arch
 from sglang.srt.connector import ConnectorType
@@ -4207,6 +4208,7 @@ class ServerArgs:
     def _handle_cuda_graph_config(self):
         self._parse_cuda_graph_config()
         self._apply_cuda_graph_compatibility()
+        self._apply_deepep_adjustments()
         self._apply_cuda_graph_disaggregation_roles()
         self._validate_cuda_graph_config()
         # Warn on the final resolved config (not inside the compat cascade —
@@ -4218,6 +4220,30 @@ class ServerArgs:
                 "Use breakable or tc_piecewise for production workloads."
             )
 
+    def _apply_deepep_adjustments(self):
+        """Config adjustments required by the DeepEP a2a backend."""
+        if resolved_view(self).moe_a2a_backend != "deepep":
+            return
+
+        # Non-multiple-of-8 prefill buckets can hang DeepEP a2a capture under
+        # breakable CUDA graph
+        if self.cuda_graph_config.prefill.backend == Backend.BREAKABLE:
+            bs = self.cuda_graph_config.prefill.bs
+            if bs is None:
+                # 2048 = documented prefill default; max_bs unresolved here.
+                max_bs = self.cuda_graph_config.prefill.max_bs or 2048
+                bs = self._generate_prefill_cuda_graph_batch_sizes(max_bs)
+            aligned = sorted({((b + 7) // 8) * 8 for b in bs})
+            if aligned != sorted(bs):
+                logger.info(
+                    "Breakable prefill CUDA graph with DeepEP requires bucket "
+                    "sizes divisible by 8; aligning %s -> %s.",
+                    sorted(bs),
+                    aligned,
+                )
+                self.cuda_graph_config.prefill.bs = aligned
+                self.cuda_graph_config.prefill.max_bs = aligned[-1]
+
     def _parse_cuda_graph_config(self):
         """Resolve cuda_graph_config from explicit JSON, per-phase
         convenience flags, legacy global flags, and defaults.
@@ -4321,8 +4347,6 @@ class ServerArgs:
                 self.cuda_graph_config.prefill.backend = Backend.DISABLED
 
     def _disable_tc_piecewise_cudagraph_if_incompatible(self):
-        from sglang.srt.arg_groups.overrides import resolved_view as _resolved_view
-
         """TcPiecewise (torch.compile + piecewise) is incompatible with
         these configurations. Most are torch.compile / dynamo limitations.
         """
@@ -4346,7 +4370,7 @@ class ServerArgs:
             ),
             (
                 "MoE A2A backend",
-                lambda: _resolved_view(self).moe_a2a_backend != "none",
+                lambda: resolved_view(self).moe_a2a_backend != "none",
             ),
             # Dynamo blocks LoRA under tc_piecewise (per-batch LoRABatchInfo
             # rebinds break guards); breakable/full support LoRA.
@@ -4359,7 +4383,7 @@ class ServerArgs:
             (
                 "GGUF quantization",
                 lambda: self.load_format == "gguf"
-                or _resolved_view(self).quantization == "gguf"
+                or resolved_view(self).quantization == "gguf"
                 or check_gguf_file(self.model_path),
             ),
             ("DLLM (diffusion LLM)", lambda: self.dllm_algorithm is not None),
@@ -4398,17 +4422,21 @@ class ServerArgs:
                 self.cuda_graph_config.prefill.backend = Backend.DISABLED
 
     def _disable_breakable_cudagraph_if_incompatible(self):
-        from sglang.srt.arg_groups.overrides import resolved_view as _resolved_view
-
         """Breakable (segmented capture, no torch.compile). Breakable enforces
         memory-saver rejection in its own __init__; config-time rules can be
         added here as they're discovered.
         """
-        from sglang.srt.configs.model_config import is_deepseek_v4
+        from sglang.srt.configs.model_config import is_deepseek_dsa, is_deepseek_v4
 
         rules = [
-            # MLA prefill takes a different attn-forward path under BCG.
-            ("MLA attention", lambda: self.use_mla_backend()),
+            # MLA prefill under BCG takes forward_mha, which has no eager
+            # breaks. DSA is exempt: BCG forces the sparse path, whose
+            # indexer already splits eagerly.
+            (
+                "MLA attention (non-DSA)",
+                lambda: self.use_mla_backend()
+                and not is_deepseek_dsa(self.get_model_config().hf_config),
+            ),
             # DSV4 is BCG-compatible but introduces heavy memory pressure: the
             # c4 indexer scratch is pinned in the capture pool and OOMs. Disable.
             (
@@ -4425,10 +4453,15 @@ class ServerArgs:
                 "decode context parallel (dcp_size > 1)",
                 lambda: self.dcp_size > 1,
             ),
-            # BCG bucket sizes exceed FlashInfer MoE A2A's dispatch cap.
+            # TBO capture is unsupported.
             (
-                "MoE A2A backend",
-                lambda: _resolved_view(self).moe_a2a_backend != "none",
+                "two-batch overlap",
+                lambda: self.enable_two_batch_overlap,
+            ),
+            # Only DeepEP's a2a is validated under BCG.
+            (
+                "non-DeepEP a2a backend",
+                lambda: resolved_view(self).moe_a2a_backend not in ("none", "deepep"),
             ),
             # Multimodal prefill replay faults under BCG; allowlisted archs opt back in.
             (
@@ -4826,12 +4859,19 @@ class ServerArgs:
                 # MLA backend overhead is much higher than expected with fa3.
                 reserved_mem += 1.5 * 1024
 
+            if (
+                prefill_cuda_graph_config.backend == Backend.BREAKABLE
+                and resolved_view(self).moe_a2a_backend == "deepep"
+            ):
+                # Prefill-BCG DeepEP delta (bridge pool + NVL first-touch
+                # during capture); decode-side DeepEP is a baseline cost.
+                reserved_mem += 1 * 1024
+
         return reserved_mem
 
     def reserve_for_deepep_a2a_mb(self) -> float:
         # DeepEP all-to-all buffers captured in the decode graph are real extra
         # allocations, reserved on top of the floor.
-        from sglang.srt.arg_groups.overrides import resolved_view
 
         decode_cuda_graph_config = self.cuda_graph_config.decode
         if (
@@ -4991,7 +5031,6 @@ class ServerArgs:
         # flags tier.
         from sglang.srt.arg_groups.overrides import (
             collect_model_override_declarations,
-            resolved_view,
             validate_declarations,
         )
 
@@ -5524,7 +5563,6 @@ class ServerArgs:
         from sglang.srt.arg_groups.overrides import (
             _mamba_radix_cache_resolution,
             mamba_extra_buffer_of,
-            resolved_view,
             run_post_process_pass,
         )
 
@@ -5574,7 +5612,6 @@ class ServerArgs:
 
         if not use_mla_backend:
             # MHA architecture
-            from sglang.srt.arg_groups.overrides import resolved_view
 
             if is_hopper_with_cuda_12_3() and is_no_spec_infer_or_topk_one(
                 resolved_view(self)
@@ -5637,7 +5674,6 @@ class ServerArgs:
             _fa4_page_constraint,
             _intel_xpu_page_constraint,
             _mla_backend_page_constraints,
-            resolved_view,
             run_post_process_pass,
         )
 
@@ -5748,7 +5784,6 @@ class ServerArgs:
 
     def _handle_kv4_compatibility(self):
         """Check FP4 KV cache compatibility with the attention backend"""
-        from sglang.srt.arg_groups.overrides import resolved_view
 
         if self.kv_cache_dtype not in ("nvfp4", "fp4_mx_block16"):
             return
@@ -6014,7 +6049,6 @@ class ServerArgs:
                 )
             from sglang.srt.arg_groups.overrides import (
                 mamba_extra_buffer_of,
-                resolved_view,
             )
 
             if mamba_extra_buffer_of(resolved_view(self)):
@@ -6374,7 +6408,6 @@ class ServerArgs:
             _cutlass_moe_env_override,
             _moe_runner_backend_quant_constraints,
             _moe_runner_fusion_disable,
-            resolved_view,
             run_post_process_pass,
         )
 
@@ -6493,7 +6526,6 @@ class ServerArgs:
         """Fail fast if the FlashInfer A2A dispatcher workspace cannot cover the
         largest CuteDSL MoE forward. Runs after speculative decoding is resolved
         so cutedsl_moe_max_num_tokens() sees the final num_tokens_per_req."""
-        from sglang.srt.arg_groups.overrides import resolved_view
 
         view = resolved_view(self)
         if not (
@@ -6536,7 +6568,6 @@ class ServerArgs:
             _a2a_backend_overrides,
             _a2a_ep_size,
             _a2a_fusion_adjustments,
-            resolved_view,
             run_post_process_pass,
         )
 
@@ -7022,7 +7053,6 @@ class ServerArgs:
         still None, backends haven't settled yet and the resolved (prefill,
         decode) pair would be a stale (None, None).
         """
-        from sglang.srt.arg_groups.overrides import resolved_view
 
         if not self.prefill_only_disable_kv_cache:
             return
@@ -7595,7 +7625,6 @@ class ServerArgs:
             from sglang.srt.arg_groups.overrides import (
                 _deterministic_attention_backend,
                 _deterministic_sampling_backend,
-                resolved_view,
                 run_post_process_pass,
             )
 
@@ -7841,7 +7870,6 @@ class ServerArgs:
             )
 
     def _handle_other_validations(self):
-        from sglang.srt.arg_groups.overrides import resolved_view
 
         # Handle optimistic prefill validation
         if (
@@ -8276,7 +8304,6 @@ class ServerArgs:
     def _resolved(self):
         """Read-only view of the resolving configuration: declared fields
         resolve from the declaration stash."""
-        from sglang.srt.arg_groups.overrides import resolved_view
 
         return resolved_view(self)
 
@@ -8338,7 +8365,6 @@ class ServerArgs:
         view so declared fields resolve from the declaration stash."""
         from sglang.srt.arg_groups.overrides import (
             attention_backends_of,
-            resolved_view,
         )
 
         return attention_backends_of(resolved_view(self))
@@ -8406,7 +8432,6 @@ class ServerArgs:
         # (or mamba_chunk_size if it is defined in the model's config) and page_size.
         # It is used to determine the caching point in a sequence during prefill.
         if not hasattr(self, "_mamba_cache_chunk_size"):
-            from sglang.srt.arg_groups.overrides import resolved_view
 
             hf_config = self.get_model_config().hf_config
             chunk_size = getattr(hf_config, "mamba_chunk_size", FLA_CHUNK_SIZE)

python/sglang/srt/speculative/eagle_worker_common.py

@@ -191,10 +191,10 @@ def prepare_for_draft_extend(
         # Supply CPU mirror (extend_seq_lens are all num_window_tokens) so
         # backend max() reads from list without a per-iter D2H sync.
         forward_batch.extend_seq_lens_cpu = [num_window_tokens] * bs
-    can_cuda_graph = cuda_graph_runner and cuda_graph_runner.can_run_graph(
+    can_run_decode_cuda_graph = cuda_graph_runner and cuda_graph_runner.can_run_graph(
         forward_batch
     )
-    if not batch.forward_mode.is_idle() and not can_cuda_graph:
+    if not batch.forward_mode.is_idle() and not can_run_decode_cuda_graph:
         draft_model_runner.attn_backend.init_forward_metadata(forward_batch)
         # Planned pre-pad; do NOT opt into post-pad re-plan. DSA's indexer
         # cannot rebuild its deep_gemm schedule_meta on a DP-padded batch
@@ -204,7 +204,7 @@ def prepare_for_draft_extend(
         # On NPU with --disable-cuda-graph, block_table shape won't match
         # after prepare_mlp_sync_batch padding; defer re-init to
         # forward_extend (post-pad) instead.
-        if not is_npu() or can_cuda_graph:
+        if not is_npu() or can_run_decode_cuda_graph:
             forward_batch.mark_forward_metadata_ready()
     return forward_batch
 
@@ -307,10 +307,10 @@ def prepare_for_draft(
         capture_hidden_mode=capture_mode,
         return_hidden_states_before_norm=False,
     )
-    can_cuda_graph = cuda_graph_runner and cuda_graph_runner.can_run_graph(
+    can_run_decode_cuda_graph = cuda_graph_runner and cuda_graph_runner.can_run_graph(
         forward_batch
     )
-    return forward_batch, can_cuda_graph
+    return forward_batch, can_run_decode_cuda_graph
 
 
 def build_eagle_verify_input(

python/sglang/srt/speculative/eagle_worker_v2.py

@@ -468,7 +468,7 @@ class EagleDraftWorker(EagleDraftWorkerBase):
 
     def draft(self, batch: ScheduleBatch):
         draft_input: EagleDraftInput = batch.spec_info
-        forward_batch, can_cuda_graph = prepare_for_draft(
+        forward_batch, can_run_decode_cuda_graph = prepare_for_draft(
             draft_input,
             self.req_to_token_pool,
             batch,
@@ -478,12 +478,12 @@ class EagleDraftWorker(EagleDraftWorkerBase):
             self.speculative_num_steps,
         )
         if (
-            can_cuda_graph
+            can_run_decode_cuda_graph
             and not forward_batch.forward_mode.is_idle()
             and self.seed_dsa_topk_from_draft_extend
             and draft_input.dsa_topk_indices is None
         ):
-            can_cuda_graph = False
+            can_run_decode_cuda_graph = False
 
         n_inner = self.speculative_num_steps - 1
         canary_outside_ctx = (
@@ -497,7 +497,7 @@ class EagleDraftWorker(EagleDraftWorkerBase):
 
         with canary_outside_ctx:
             # Run draft
-            if can_cuda_graph:
+            if can_run_decode_cuda_graph:
                 parent_list, top_scores_index, draft_tokens, draft_probs = (
                     self.cuda_graph_runner.execute(forward_batch)
                 )
@@ -882,14 +882,14 @@ class EagleDraftWorker(EagleDraftWorkerBase):
             )
 
         # Run draft extend batch in the main compute stream
-        can_cuda_graph = (
+        can_run_decode_cuda_graph = (
             self.cuda_graph_runner_for_draft_extend
             and self.cuda_graph_runner_for_draft_extend.can_run_graph(forward_batch)
         )
 
         # Eager path publishes the indexer top-k into a worker buffer (the graph
         # path uses the runner's static buffer). Gathered at select_index below.
-        if self.seed_dsa_topk_from_draft_extend and not can_cuda_graph:
+        if self.seed_dsa_topk_from_draft_extend and not can_run_decode_cuda_graph:
             forward_batch.spec_info.dsa_seed_topk_capture = (
                 self._get_dsa_extend_topk_buf(forward_batch.input_ids.shape[0])
             )
@@ -906,7 +906,7 @@ class EagleDraftWorker(EagleDraftWorkerBase):
             else contextlib.nullcontext()
         )
         with canary_ctx:
-            if can_cuda_graph:
+            if can_run_decode_cuda_graph:
                 draft_logits_output = self.cuda_graph_runner_for_draft_extend.execute(
                     forward_batch
                 )
@@ -917,18 +917,18 @@ class EagleDraftWorker(EagleDraftWorkerBase):
 
         maybe_detect_nan(
             draft_logits_output.next_token_logits,
-            f"draft_extend_for_decode (cuda_graph={can_cuda_graph})",
+            f"draft_extend_for_decode (cuda_graph={can_run_decode_cuda_graph})",
         )
         maybe_detect_inf(
             draft_logits_output.next_token_logits,
-            f"draft_extend_for_decode (cuda_graph={can_cuda_graph})",
+            f"draft_extend_for_decode (cuda_graph={can_run_decode_cuda_graph})",
         )
 
         # Gather the per-request last-position indexer top-k as the next loop's
         # seed (select_index already picks the last accepted position per req).
         dsa_seed_topk_indices = None
         if self.seed_dsa_topk_from_draft_extend:
-            if can_cuda_graph:
+            if can_run_decode_cuda_graph:
                 dsa_extend_topk_capture = (
                     self.cuda_graph_runner_for_draft_extend.buffers.dsa_seed_topk_capture
                 )

python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py

@@ -406,7 +406,7 @@ class MultiLayerEagleDraftWorker(EagleDraftWorkerBase):
 
     def draft(self, batch: ScheduleBatch):
         draft_input: EagleDraftInput = batch.spec_info
-        forward_batch, can_cuda_graph = prepare_for_draft(
+        forward_batch, can_run_decode_cuda_graph = prepare_for_draft(
             draft_input,
             self.req_to_token_pool,
             batch,
@@ -732,7 +732,7 @@ class MultiLayerEagleDraftWorker(EagleDraftWorkerBase):
         forward_batch.spec_info.num_accept_tokens = batch_result.accept_lens
 
         # Run draft extend batch in the main compute stream
-        can_cuda_graph = (
+        can_run_decode_cuda_graph = (
             self.cuda_graph_runner_for_draft_extend
             and self.cuda_graph_runner_for_draft_extend.can_run_graph(forward_batch)
         )
@@ -742,7 +742,7 @@ class MultiLayerEagleDraftWorker(EagleDraftWorkerBase):
         ret_draft_probs = None
         next_token_ids_backup = batch_result.next_token_ids.clone()
 
-        if can_cuda_graph:
+        if can_run_decode_cuda_graph:
             # Graph replay bypasses ModelRunner.forward, which emits the
             # step[...] trace span for every other phase; emit it here.
             with profile_range(build_step_span_name(forward_batch)):