From 6480cce9bc7bb83e18537be3995a094cf2426e74 Mon Sep 17 00:00:00 2001 From: YAMY <74099316+YAMY1234@users.noreply.github.com> Date: Thu, 27 Aug 2026 20:04:54 -0700 Subject: [PATCH] perf: skip redundant scheduler metadata gather for DP1 (#36568) --- .../managers/scheduler_components/dp_attn.py | 43 ++++++++-- .../scheduler_components/test_dp_attn.py | 79 +++++++++++++++++++ 2 files changed, 116 insertions(+), 6 deletions(-) create mode 100644 test/registered/unit/managers/scheduler_components/test_dp_attn.py diff --git a/python/sglang/srt/managers/scheduler_components/dp_attn.py b/python/sglang/srt/managers/scheduler_components/dp_attn.py index 8ac9da6f7..bffbfab0f 100644 --- a/python/sglang/srt/managers/scheduler_components/dp_attn.py +++ b/python/sglang/srt/managers/scheduler_components/dp_attn.py @@ -136,6 +136,16 @@ class MLPSyncBatchInfo: dtype=dtype, ) + def finalize_local(self): + """Populate gather-derived metadata from the sole attention-DP rank.""" + self.tp0_info_cpu = self._get_local_tensor(device="cpu").view(1, -1) + self.global_num_tokens = [self.num_tokens] + self.global_num_tokens_for_logprob = [self.num_tokens_for_logprob] + if _ENABLE_METRICS_DP_ATTENTION: + self.dp_cooperation_info = DPCooperationInfo.create( + self.tp0_info_cpu[:, 5].tolist() + ) + def all_gather( self, device, @@ -212,7 +222,7 @@ def _update_gather_batch( batch: ScheduleBatch, mlp_sync_info: MLPSyncBatchInfo, require_mlp_tp_gather: bool, - skip_all_gather=False, + skip_global_metadata=False, ): # TODO: handle the case when moe_dense_tp_size != 1 if not require_mlp_tp_gather: @@ -223,7 +233,7 @@ def _update_gather_batch( batch.global_num_tokens_for_logprob = ( mlp_sync_info.global_num_tokens_for_logprob ) - if not skip_all_gather: + if not skip_global_metadata: batch.is_extend_in_batch = mlp_sync_info.is_extend_in_batch batch.tbo_split_seq_index = mlp_sync_info.tbo_split_seq_index batch.global_forward_mode = mlp_sync_info.global_forward_mode @@ -233,6 +243,20 @@ def _update_gather_batch( batch.can_run_dp_prefill_cuda_graph = mlp_sync_info.can_run_prefill_cuda_graph +def should_skip_scheduler_all_gather(dp_size: int) -> bool: + """Return whether scheduler metadata is already local and rank-invariant. + + With one attention-DP rank there is no cross-DP state to reconcile. The + TP schedulers consume the same broadcast request stream, so gathering the + identical batch mode, graph eligibility, and token counts only adds a + device collective plus host synchronization. Preserve the environment + override for deployments that explicitly guarantee this invariant beyond + DP1. + """ + + return dp_size == 1 or envs.SGLANG_SCHEDULER_SKIP_ALL_GATHER.get() + + def _local_decode_cuda_graph_vote( *, local_batch: Optional[ScheduleBatch], @@ -358,7 +382,6 @@ def prepare_mlp_sync_batch_raw( or num_tokens_for_logprob == local_batch.batch_size() ) - skip_all_gather = envs.SGLANG_SCHEDULER_SKIP_ALL_GATHER.get() can_run_decode_cuda_graph = _local_decode_cuda_graph_vote( local_batch=local_batch, disable_cuda_graph=disable_cuda_graph ) @@ -408,6 +431,7 @@ def prepare_mlp_sync_batch_raw( local_num_tokens=num_tokens, local_forward_mode=local_forward_mode, ) + skip_all_gather = should_skip_scheduler_all_gather(dp_size) mlp_sync_info = MLPSyncBatchInfo( dp_size=dp_size, @@ -422,13 +446,17 @@ def prepare_mlp_sync_batch_raw( local_forward_mode=local_forward_mode, ) - if not skip_all_gather: + if dp_size == 1: + mlp_sync_info.finalize_local() + elif not skip_all_gather: mlp_sync_info.all_gather( device=device, group=group, use_all_reduce=use_world_group, ) + metadata_ready = mlp_sync_info.tp0_info_cpu is not None + if metadata_ready: mlp_sync_info.tbo_split_seq_index, mlp_sync_info.global_forward_mode = ( tbo_preparer.compute_output( mlp_sync_info.tp0_info_cpu[:, 4:6], @@ -452,12 +480,15 @@ def prepare_mlp_sync_batch_raw( if batch_to_gather is not None: _update_gather_batch( - batch_to_gather, mlp_sync_info, require_mlp_tp_gather, skip_all_gather + batch_to_gather, + mlp_sync_info, + require_mlp_tp_gather, + skip_global_metadata=not metadata_ready, ) # Set on `local_batch`, not `batch_to_gather`: for PREBUILT batches the # scheduler's `last_batch` is the prebuilt batch, not its inner idle batch. - if local_batch is not None and not skip_all_gather: + if local_batch is not None and metadata_ready: local_batch.recv_skipper_forward_mode = ( SchedulerRecvSkipper.derive_forward_mode( mlp_sync_info.tp0_info_cpu[:, 5].tolist() diff --git a/test/registered/unit/managers/scheduler_components/test_dp_attn.py b/test/registered/unit/managers/scheduler_components/test_dp_attn.py new file mode 100644 index 000000000..59e792266 --- /dev/null +++ b/test/registered/unit/managers/scheduler_components/test_dp_attn.py @@ -0,0 +1,79 @@ +import unittest +from types import SimpleNamespace +from unittest.mock import Mock, patch + +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase, maybe_stub_sgl_kernel + +maybe_stub_sgl_kernel() + +from sglang.srt.environ import envs # noqa: E402 +from sglang.srt.managers.scheduler_components import dp_attn # noqa: E402 +from sglang.srt.model_executor.forward_batch_info import ForwardMode # noqa: E402 +from sglang.srt.speculative.spec_info import SpeculativeAlgorithm # noqa: E402 + +register_cpu_ci(est_time=2, suite="base-a-test-cpu") + + +class TestDPAttnSchedulerMetadata(CustomTestCase): + def test_skip_all_gather_policy(self): + with envs.SGLANG_SCHEDULER_SKIP_ALL_GATHER.override(False): + self.assertTrue(dp_attn.should_skip_scheduler_all_gather(dp_size=1)) + self.assertFalse(dp_attn.should_skip_scheduler_all_gather(dp_size=2)) + with envs.SGLANG_SCHEDULER_SKIP_ALL_GATHER.override(True): + self.assertTrue(dp_attn.should_skip_scheduler_all_gather(dp_size=2)) + + def test_dp1_skip_preserves_local_tbo_metadata(self): + batch = SimpleNamespace( + forward_mode=ForwardMode.DECODE, + batch_size=lambda: 4, + ) + tbo_preparer = Mock() + tbo_preparer.prepare_all_gather.return_value = ( + True, + ForwardMode.DECODE.value, + ) + tbo_preparer.compute_output.return_value = (2, ForwardMode.DECODE) + + with ( + envs.SGLANG_SCHEDULER_SKIP_ALL_GATHER.override(False), + patch.object(dp_attn, "TboDPAttentionPreparer", return_value=tbo_preparer), + patch.object(dp_attn, "world_dp_gather_enabled", return_value=False), + patch.object(dp_attn, "check_cuda_graph_backend", return_value=False), + patch.object(dp_attn.MLPSyncBatchInfo, "all_gather") as all_gather, + ): + result = dp_attn.prepare_mlp_sync_batch_raw( + batch, + model_runner=SimpleNamespace( + prefill_cuda_graph_runner=None, + spec_algorithm=SpeculativeAlgorithm.NONE, + model_config=object(), + ), + dp_size=1, + attn_tp_size=4, + attn_cp_size=1, + tp_group=SimpleNamespace( + device_group=object(), device="cpu", cpu_group=object() + ), + get_idle_batch=Mock( + side_effect=AssertionError("DP1 must not emit idle batch") + ), + disable_cuda_graph=False, + require_mlp_tp_gather=False, + disable_overlap_schedule=True, + offload_tags=set(), + ) + + all_gather.assert_not_called() + self.assertEqual(result.global_num_tokens, [4]) + self.assertEqual(result.tbo_split_seq_index, 2) + self.assertEqual(result.global_forward_mode, ForwardMode.DECODE) + self.assertEqual(result.recv_skipper_forward_mode, ForwardMode.DECODE) + self.assertEqual( + tbo_preparer.compute_output.call_args.args[0].tolist(), + [[1, ForwardMode.DECODE.value]], + ) + + +if __name__ == "__main__": + unittest.main()