perf: skip redundant scheduler metadata gather for DP1 (#36568)

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()

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()