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

2026-08-27 20:04:54 -07:00
parent daf6317196
commit 6480cce9bc
2 changed files with 116 additions and 6 deletions
@@ -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()