Allow chunk-only prefill/decode interval protection

python/sglang/srt/environ.py

@@ -655,6 +655,10 @@ class Envs:
     # PP: skip output send/recv when the entire batch consists of non-final chunked prefill requests,
     # since process_batch_result_prefill discards next_token_ids for those anyway.
     SGLANG_PP_SKIP_PURE_CHUNKED_OUTPUT_COMM = EnvBool(False)
+    # Restrict prefill-decode interval protection to unfinished prompt chunks.
+    # DP attention retains its synchronized fixed cadence. Default preserves
+    # protection after every prefill batch.
+    SGLANG_PDI_CHUNKED_ONLY = EnvBool(False)
     # Run PP tensor communication on a dedicated stream so asynchronous sends
     # do not fence the next forward through the scheduler stream.
     SGLANG_PP_COMM_OVERLAP = EnvBool(False)

python/sglang/srt/managers/scheduler.py

@@ -1313,6 +1313,7 @@ def init_running_status(self):
     def init_chunked_prefill(self):
         self.chunked_prefill_size = get_schedule().chunked_prefill_size
         self.prefill_decode_interval = get_schedule().prefill_decode_interval or 0
+        self.pdi_chunked_only = envs.SGLANG_PDI_CHUNKED_ONLY.get()
         self._prefill_decode_interval_remaining = 0
         uses_transformers_backend = (
             get_resolved_model_impl(self.model_config) == ModelImpl.TRANSFORMERS
@@ -1376,7 +1377,17 @@ def _arm_prefill_decode_interval(self, batch: Optional[ScheduleBatch]) -> None:
             else batch.forward_mode.is_extend()
         )
         if is_extend:
-            self._prefill_decode_interval_remaining = self.prefill_decode_interval
+            # Protect decoding between unfinished prompt chunks. Completed
+            # prefills can admit the next request immediately in this mode.
+            # DP ranks retain the globally synchronized fixed cadence.
+            protect_decode = (
+                not self.pdi_chunked_only
+                or self.require_mlp_sync
+                or batch.chunked_req is not None
+            )
+            self._prefill_decode_interval_remaining = (
+                self.prefill_decode_interval if protect_decode else 0
+            )
 
     def init_metrics_reporter(self) -> None:
         # Override point for deployments that need a specialized reporter.

test/registered/unit/managers/test_scheduler_prefill_decode_interval.py

@@ -13,18 +13,22 @@
 register_cpu_ci(est_time=11, suite="base-a-test-cpu")
 
 
-def _make_scheduler(*, interval: int, require_mlp_sync: bool) -> Scheduler:
+def _make_scheduler(
+    *, interval: int, require_mlp_sync: bool, chunked_only: bool = False
+) -> Scheduler:
     scheduler = Scheduler.__new__(Scheduler)
     scheduler.prefill_decode_interval = interval
+    scheduler.pdi_chunked_only = chunked_only
     scheduler._prefill_decode_interval_remaining = 0
     scheduler.require_mlp_sync = require_mlp_sync
     return scheduler
 
 
-def _make_batch(*, local_extend: bool, global_extend: bool):
+def _make_batch(*, local_extend: bool, global_extend: bool, chunked_req=None):
     return SimpleNamespace(
         forward_mode=SimpleNamespace(is_extend=lambda: local_extend),
         is_extend_in_batch=global_extend,
+        chunked_req=chunked_req,
     )
 
 
@@ -70,6 +74,75 @@ def test_decode_batch_does_not_rearm_interval(self):
 
         self.assertEqual(scheduler._prefill_decode_interval_remaining, 1)
 
+    def test_completed_prefill_does_not_delay_next_admission(self):
+        scheduler = _make_scheduler(
+            interval=64, require_mlp_sync=False, chunked_only=True
+        )
+        scheduler._arm_prefill_decode_interval(
+            _make_batch(local_extend=True, global_extend=False)
+        )
+        self.assertFalse(scheduler._should_defer_prefill())
+
+    def test_unfinished_prompt_keeps_decode_protection(self):
+        scheduler = _make_scheduler(
+            interval=2, require_mlp_sync=False, chunked_only=True
+        )
+        scheduler._arm_prefill_decode_interval(
+            _make_batch(local_extend=True, global_extend=False, chunked_req=object())
+        )
+        self.assertTrue(scheduler._should_defer_prefill())
+        self.assertTrue(scheduler._should_defer_prefill())
+        self.assertFalse(scheduler._should_defer_prefill())
+
+    def test_finishing_chunk_clears_remaining_hold(self):
+        scheduler = _make_scheduler(
+            interval=64, require_mlp_sync=False, chunked_only=True
+        )
+        scheduler._prefill_decode_interval_remaining = 17
+        scheduler._arm_prefill_decode_interval(
+            _make_batch(local_extend=True, global_extend=False)
+        )
+        self.assertEqual(scheduler._prefill_decode_interval_remaining, 0)
+
+    def test_dp_ranks_keep_matching_cadence(self):
+        for local_extend, chunked_req in (
+            (True, object()),
+            (True, None),
+            (False, None),
+        ):
+            with self.subTest(
+                local_extend=local_extend, unfinished=chunked_req is not None
+            ):
+                scheduler = _make_scheduler(
+                    interval=2, require_mlp_sync=True, chunked_only=True
+                )
+                scheduler._arm_prefill_decode_interval(
+                    _make_batch(
+                        local_extend=local_extend,
+                        global_extend=True,
+                        chunked_req=chunked_req,
+                    )
+                )
+                self.assertEqual(scheduler._prefill_decode_interval_remaining, 2)
+
+    def test_decode_does_not_clear_chunk_protection(self):
+        scheduler = _make_scheduler(
+            interval=2, require_mlp_sync=False, chunked_only=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)
+
+    def test_none_batch_does_not_change_protection(self):
+        scheduler = _make_scheduler(
+            interval=2, require_mlp_sync=False, chunked_only=True
+        )
+        scheduler._prefill_decode_interval_remaining = 1
+        scheduler._arm_prefill_decode_interval(None)
+        self.assertEqual(scheduler._prefill_decode_interval_remaining, 1)
+
 
 if __name__ == "__main__":
     unittest.main()