[Scheduler] Add configurable decode interval after prefill (#35017)

test/registered/unit/managers/test_scheduler_chunked_req_gate.py

@@ -91,6 +91,8 @@ def _scheduler_for_get_next_batch(*, tree_cache, chunked_req) -> Scheduler:
     s.running_batch.is_prefill_only = False
     s.running_batch.batch_is_full = False
     s.running_batch.reqs = []
+    s.prefill_decode_interval = 0
+    s._prefill_decode_interval_remaining = 0
     s.get_new_batch_prefill = MagicMock(
         return_value=NextBatchPlan(batch_to_run=None, running_batch=s.running_batch)
     )

test/registered/unit/managers/test_scheduler_prefill_decode_interval.py

@@ -0,0 +1,75 @@
+"""Tests for scheduler prefill/decode interleaving."""
+
+import unittest
+from types import SimpleNamespace
+
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import maybe_stub_sgl_kernel
+
+maybe_stub_sgl_kernel()
+
+from sglang.srt.managers.scheduler import Scheduler
+
+register_cpu_ci(est_time=1, suite="base-a-test-cpu")
+
+
+def _make_scheduler(*, interval: int, require_mlp_sync: bool) -> Scheduler:
+    scheduler = Scheduler.__new__(Scheduler)
+    scheduler.prefill_decode_interval = interval
+    scheduler._prefill_decode_interval_remaining = 0
+    scheduler.require_mlp_sync = require_mlp_sync
+    return scheduler
+
+
+def _make_batch(*, local_extend: bool, global_extend: bool):
+    return SimpleNamespace(
+        forward_mode=SimpleNamespace(is_extend=lambda: local_extend),
+        is_extend_in_batch=global_extend,
+    )
+
+
+class TestPrefillDecodeInterval(unittest.TestCase):
+    def test_disabled_interval_does_not_arm(self):
+        scheduler = _make_scheduler(interval=0, require_mlp_sync=False)
+
+        scheduler._arm_prefill_decode_interval(
+            _make_batch(local_extend=True, global_extend=False)
+        )
+
+        self.assertFalse(scheduler._should_defer_prefill())
+
+    def test_non_dp_interval_uses_local_forward_mode(self):
+        scheduler = _make_scheduler(interval=2, require_mlp_sync=False)
+
+        scheduler._arm_prefill_decode_interval(
+            _make_batch(local_extend=True, global_extend=False)
+        )
+
+        self.assertTrue(scheduler._should_defer_prefill())
+        self.assertTrue(scheduler._should_defer_prefill())
+        self.assertFalse(scheduler._should_defer_prefill())
+
+    def test_dp_interval_uses_globally_synchronized_extend_flag(self):
+        scheduler = _make_scheduler(interval=2, require_mlp_sync=True)
+
+        # This rank is locally decoding, but another DP rank is prefilling.
+        scheduler._arm_prefill_decode_interval(
+            _make_batch(local_extend=False, global_extend=True)
+        )
+
+        self.assertEqual(scheduler._prefill_decode_interval_remaining, 2)
+        self.assertTrue(scheduler._should_defer_prefill())
+
+    def test_decode_batch_does_not_rearm_interval(self):
+        scheduler = _make_scheduler(interval=2, require_mlp_sync=True)
+        scheduler._prefill_decode_interval_remaining = 1
+
+        scheduler._arm_prefill_decode_interval(
+            _make_batch(local_extend=False, global_extend=False)
+        )
+
+        self.assertEqual(scheduler._prefill_decode_interval_remaining, 1)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/server_args/test_server_args.py

@@ -43,6 +43,15 @@ _mock_device.start()
 
 
 class TestPrepareServerArgs(CustomTestCase):
+    def test_prefill_decode_interval(self):
+        args = ServerArgs(model_path="dummy", prefill_decode_interval=16)
+        self.assertEqual(args.prefill_decode_interval, 16)
+
+        with self.assertRaisesRegex(
+            ValueError, "--prefill-decode-interval must be non-negative"
+        ):
+            ServerArgs(model_path="dummy", prefill_decode_interval=-1)
+
     def test_return_hidden_states_mode_configuration(self):
         disabled = ServerArgs(model_path="dummy")
         self.assertFalse(disabled.enable_return_hidden_states)