From 3e0f7c3f30e246a0a7fdba7eb9b11e6383943ea4 Mon Sep 17 00:00:00 2001 From: Yuwei An Date: Fri, 31 Jul 2026 16:45:18 -0700 Subject: [PATCH] [BCG][3/N] Enable bcg on dsa & deepep a2a backend (#31987) Co-authored-by: Claude Fable 5 --- python/sglang/benchmark/one_batch.py | 1 + python/sglang/srt/layers/moe/ep_moe/layer.py | 67 +++++++- python/sglang/srt/managers/scheduler.py | 9 ++ .../managers/scheduler_components/dp_attn.py | 61 +++++--- .../runner/prefill_cuda_graph_runner.py | 146 ++++++++++++------ .../breakable_cuda_graph_backend.py | 1 + .../breakable_cuda_graph.py | 24 ++- python/sglang/srt/models/deepseek_v2.py | 8 + python/sglang/srt/server_args.py | 81 ++++++---- .../srt/speculative/eagle_worker_common.py | 10 +- .../sglang/srt/speculative/eagle_worker_v2.py | 20 +-- .../multi_layer_eagle_worker_v2.py | 6 +- .../runner/test_prefill_cuda_graph_padding.py | 1 + 13 files changed, 319 insertions(+), 116 deletions(-) diff --git a/python/sglang/benchmark/one_batch.py b/python/sglang/benchmark/one_batch.py index 8556becec..642f4d863 100644 --- a/python/sglang/benchmark/one_batch.py +++ b/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, diff --git a/python/sglang/srt/layers/moe/ep_moe/layer.py b/python/sglang/srt/layers/moe/ep_moe/layer.py index f656fe9ec..f6a445b42 100644 --- a/python/sglang/srt/layers/moe/ep_moe/layer.py +++ b/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 diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index afea8748c..8857ae8cb 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/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, diff --git a/python/sglang/srt/managers/scheduler_components/dp_attn.py b/python/sglang/srt/managers/scheduler_components/dp_attn.py index 4ea7f2355..5b2dfafeb 100644 --- a/python/sglang/srt/managers/scheduler_components/dp_attn.py +++ b/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, 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 84860b028..79b033f2b 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 @@ -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 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 6235da423..55f8eabe2 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 @@ -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) diff --git a/python/sglang/srt/model_executor/runner_backend_utils/breakable_cuda_graph/breakable_cuda_graph.py b/python/sglang/srt/model_executor/runner_backend_utils/breakable_cuda_graph/breakable_cuda_graph.py index f484e3a90..9c7601515 100644 --- a/python/sglang/srt/model_executor/runner_backend_utils/breakable_cuda_graph/breakable_cuda_graph.py +++ b/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 diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index e63313f85..095dad38d 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/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) diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index a591308f6..a179e7753 100644 --- a/python/sglang/srt/server_args.py +++ b/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) diff --git a/python/sglang/srt/speculative/eagle_worker_common.py b/python/sglang/srt/speculative/eagle_worker_common.py index efb3891a6..89a371f3e 100644 --- a/python/sglang/srt/speculative/eagle_worker_common.py +++ b/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( diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py index 0139e9e94..60f4dad89 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/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 ) 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 2b2e15578..3c8568fee 100644 --- a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py +++ b/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)): diff --git a/test/registered/unit/model_executor/runner/test_prefill_cuda_graph_padding.py b/test/registered/unit/model_executor/runner/test_prefill_cuda_graph_padding.py index 9351253ad..36183c9b2 100644 --- a/test/registered/unit/model_executor/runner/test_prefill_cuda_graph_padding.py +++ b/test/registered/unit/model_executor/runner/test_prefill_cuda_graph_padding.py @@ -39,6 +39,7 @@ class TestPrefillCudaGraphPadding(CustomTestCase): global_num_tokens_cpu=None, return_logprob=False, input_ids=list(range(num_tokens)), + extend_prefix_lens_cpu=[0], ) def test_rejects_more_than_two_x_token_padding(self):