diff --git a/python/sglang/srt/layers/dp_attention.py b/python/sglang/srt/layers/dp_attention.py index 12bc3944a..62a981f3c 100644 --- a/python/sglang/srt/layers/dp_attention.py +++ b/python/sglang/srt/layers/dp_attention.py @@ -820,12 +820,19 @@ def _dp_gather_via_all_gatherv( get_tp_group().all_gatherv(local_real, sizes=sizes, output=global_tokens) +def _note_dp_gather_in_prefill_graph() -> None: + dp = get_flags().dp + if dp.capturing_prefill_graph: + dp.prefill_graph_has_dp_gather = True + + def _dp_gather( global_tokens: torch.Tensor, local_tokens: torch.Tensor, forward_batch: ForwardBatch, is_partial: bool, ): + _note_dp_gather_in_prefill_graph() if ( is_dp_gatherv_active() and forward_batch.dp_padding_mode is not None @@ -883,6 +890,7 @@ def dp_scatter( global_tokens: torch.Tensor, # input forward_batch: ForwardBatch, ): + _note_dp_gather_in_prefill_graph() # local_num_tokens is not necessarily the same as local_tokens.shape[0], # since local_tokens may be padded for cuda graph local_start_pos, local_num_tokens = get_dp_local_info(forward_batch) @@ -899,6 +907,7 @@ def dp_scatter( def dp_reduce_scatter_tensor(output: torch.Tensor, input: torch.Tensor): + _note_dp_gather_in_prefill_graph() if is_dp_gatherv_active(): # Variable-length combine matching all_gatherv dispatch: scatter the # global (sum_len) tensor back to per-rank token counts. Fall through to diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index 84c449422..5fff91290 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/python/sglang/srt/model_executor/forward_batch_info.py @@ -53,6 +53,7 @@ from sglang.srt.model_executor.forward_batch_deepseek_mha_mixin import ( ) from sglang.srt.runtime_context import ( get_exec, + get_flags, get_lora, get_parallel, mamba_cache_chunk_size, @@ -303,13 +304,19 @@ def compute_local_num_token_non_padded_cpu( def prefill_graph_tolerates_sum_len() -> bool: - """Whether MegaMoE may replay prefill graphs with local shapes.""" + """Whether MegaMoE may replay prefill graphs with per-rank SUM_LEN buckets. + + The graph body is captured with MAX_LEN geometry, so a graph that recorded + a DP gather/scatter only replays correctly when every rank uses one bucket. + """ from sglang.srt.layers.attention.dsa.utils import is_dsa_enable_prefill_cp from sglang.srt.layers.cp.utils import is_mla_cp_enabled from sglang.srt.layers.moe.utils import get_moe_a2a_backend if not get_moe_a2a_backend().is_megamoe(): return False + if get_flags().dp.prefill_graph_has_dp_gather: + return False return not (is_dsa_enable_prefill_cp() or is_mla_cp_enabled()) 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 4ef9f69ac..7448c0ea4 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 @@ -126,6 +126,7 @@ from sglang.srt.model_executor.runner_utils.pool import ( from sglang.srt.model_loader.utils import resolve_language_model from sglang.srt.runtime_context import ( get_exec, + get_flags, get_memory, get_parallel, get_schedule, @@ -844,9 +845,8 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner): # DSV4 DP attention / DeepEP collectives need every DP rank to enter # the same replay path. Sparse-DP batches (one or more ranks with # zero local tokens) fall back to eager to avoid hanging ranks. - # MegaMoE is exempt (prefill_graph_tolerates_sum_len): its idle ranks - # still execute MegaMoE with 0 tokens, so per-rank SUM_LEN buckets stay - # collective-safe and need no eager fallback. + # MegaMoE graphs without a captured DP gather tolerate per-rank buckets + # and an eager idle rank; graphs with one fall through to the check. global_num_tokens = forward_batch.global_num_tokens_cpu if global_num_tokens is None: return False @@ -1412,13 +1412,23 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner): # Warm up + autotune kernels once before capture (run-once across the # decode + prefill runners; see BaseRunner.warmup). self.warmup() - with freeze_gc(get_exec().graph.enable_cudagraph_gc): - with graph_capture( - stream=get_or_create_global_graph_capture_stream() - ) as graph_capture_context: - self.stream = graph_capture_context.stream - with self.backend.capture_session(self.stream): - self._capture_one_stream() + dp_flags = get_flags().dp + dp_flags.capturing_prefill_graph = True + try: + with freeze_gc(get_exec().graph.enable_cudagraph_gc): + with graph_capture( + stream=get_or_create_global_graph_capture_stream() + ) as graph_capture_context: + self.stream = graph_capture_context.stream + with self.backend.capture_session(self.stream): + self._capture_one_stream() + finally: + dp_flags.capturing_prefill_graph = False + if dp_flags.prefill_graph_has_dp_gather: + logger.info( + "Prefill CUDA graph captured a DP gather/scatter; " + "DP ranks will replay a shared MAX_LEN bucket." + ) def _capture_one_stream(self) -> None: avail_mem = get_available_gpu_memory( diff --git a/python/sglang/srt/runtime_context.py b/python/sglang/srt/runtime_context.py index e6687cf9f..8eff31689 100644 --- a/python/sglang/srt/runtime_context.py +++ b/python/sglang/srt/runtime_context.py @@ -565,6 +565,10 @@ class DpFlags(_FlagGroupBase): # Hybrid-SSM models materialize idle ranks via the MAX_LEN fabricated-row # conversion (set when hf_config has hybrid_override_pattern). max_len_with_idle: bool = False + # Set while the prefill CUDA graph runner captures; latched by the DP + # gather/scatter helpers, whose captured geometry needs one shared bucket. + capturing_prefill_graph: bool = False + prefill_graph_has_dp_gather: bool = False # DP gathered-buffer allocation metadata (model hidden size / dtype / # device), set by initialize_dp_attention alongside the flags above. buffer_hidden_size: Any = None 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 af402b707..ec33002f4 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 @@ -2,10 +2,13 @@ import unittest from types import SimpleNamespace from unittest import mock +from sglang.srt.layers.moe.utils import MoeA2ABackend +from sglang.srt.model_executor import forward_batch_info from sglang.srt.model_executor.cuda_graph_config import Backend from sglang.srt.model_executor.forward_batch_info import ( CaptureHiddenMode, ForwardMode, + prefill_graph_tolerates_sum_len, ) from sglang.srt.model_executor.runner.prefill_cuda_graph_runner import ( PrefillCudaGraphRunner, @@ -72,6 +75,68 @@ class TestPrefillCudaGraphPadding(CustomTestCase): forward_batch, num_qo_tokens=16 ) + def _megamoe_no_prefill_cp(self, graph_has_dp_gather=False): + return ( + mock.patch.object( + forward_batch_info, + "get_flags", + return_value=SimpleNamespace( + dp=SimpleNamespace(prefill_graph_has_dp_gather=graph_has_dp_gather) + ), + ), + mock.patch( + "sglang.srt.layers.moe.utils.get_moe_a2a_backend", + return_value=MoeA2ABackend.MEGAMOE, + ), + mock.patch( + "sglang.srt.layers.attention.dsa.utils.is_dsa_enable_prefill_cp", + return_value=False, + ), + mock.patch( + "sglang.srt.layers.cp.utils.is_mla_cp_enabled", + return_value=False, + ), + ) + + def test_megamoe_idle_rank_without_graph_gather_keeps_sum_len(self): + # Without a DP gather in the graph, an eager idle rank matches the + # replaying peers, so the sparse batch keeps per-rank buckets. + runner = self._make_runner() + flags, a2a, dsa_cp, mla_cp = self._megamoe_no_prefill_cp() + with flags, a2a, dsa_cp, mla_cp: + self.assertTrue(prefill_graph_tolerates_sum_len()) + self.assertFalse( + runner._has_inactive_dp_rank( + SimpleNamespace(global_num_tokens_cpu=[8, 0]) + ) + ) + + def test_megamoe_all_ranks_busy_keeps_per_rank_buckets(self): + runner = self._make_runner() + flags, a2a, dsa_cp, mla_cp = self._megamoe_no_prefill_cp() + with flags, a2a, dsa_cp, mla_cp: + self.assertTrue(prefill_graph_tolerates_sum_len()) + self.assertFalse( + runner._has_inactive_dp_rank( + SimpleNamespace(global_num_tokens_cpu=[8, 16]) + ) + ) + + def test_megamoe_graph_with_dp_gather_forces_shared_bucket(self): + # A DP gather captured in the graph has fixed MAX_LEN geometry; per-rank + # buckets or an eager idle rank would deadlock its all_gather. + runner = self._make_runner() + flags, a2a, dsa_cp, mla_cp = self._megamoe_no_prefill_cp( + graph_has_dp_gather=True + ) + with flags, a2a, dsa_cp, mla_cp: + self.assertFalse(prefill_graph_tolerates_sum_len()) + self.assertTrue( + runner._has_inactive_dp_rank( + SimpleNamespace(global_num_tokens_cpu=[8, 0]) + ) + ) + if __name__ == "__main__": unittest.main()