Refactor and fix prefill delayer (scheduler enhancer) (#16269)

python/sglang/srt/environ.py

@@ -223,6 +223,7 @@ class Envs:
     SGLANG_DYNAMIC_CHUNKING_SMOOTH_FACTOR = EnvFloat(0.75)
     SGLANG_SCHEDULER_SKIP_ALL_GATHER = EnvBool(False)
     SGLANG_SCHEDULER_DECREASE_PREFILL_IDLE = EnvBool(False)
+    SGLANG_PREFILL_DELAYER_MAX_DELAY_PASSES = EnvInt(30)
     SGLANG_DATA_PARALLEL_BUDGET_INTERVAL = EnvInt(1)
 
     # Test: pd-disaggregation

python/sglang/srt/managers/prefill_delayer.py

@@ -0,0 +1,72 @@
+import logging
+
+import torch
+
+from sglang.srt.environ import envs
+from sglang.srt.utils import get_bool_env_var
+
+_DEBUG_LOG = get_bool_env_var("SGLANG_PREFILL_DELAYER_DEBUG_LOG")
+
+logger = logging.getLogger(__name__)
+
+
+class PrefillDelayer:
+    def __init__(self, dp_size, attn_tp_size, tp_worker, server_args):
+        self.global_info = torch.empty(
+            (dp_size, attn_tp_size, 1),
+            dtype=torch.int64,
+            device="cpu",
+        )
+        self.cpu_group = tp_worker.get_tp_group().cpu_group
+
+        self.curr_delayed_count = 0
+        self.max_delay_passes = envs.SGLANG_PREFILL_DELAYER_MAX_DELAY_PASSES.get()
+
+        assert (
+            server_args.schedule_policy == "fcfs"
+        ), f"To use PrefillDelayer, schedule_policy must be 'fcfs'. '{server_args.schedule_policy}' is not supported."
+        assert (
+            server_args.enable_dp_attention
+        ), "To use PrefillDelayer, enable_dp_attention must be enabled."
+        assert (
+            server_args.disaggregation_mode == "null"
+        ), "To use PrefillDelayer, disaggregation_mode must be null."
+        assert (
+            not server_args.disable_overlap_schedule
+        ), "To use PrefillDelayer, disable_overlap_schedule must be False."
+
+    def _gather_info(self, local_prefillable: int):
+        local_info = torch.tensor(
+            [local_prefillable],
+            device="cpu",
+            dtype=torch.int64,
+        )
+        torch.distributed.all_gather_into_tensor(
+            self.global_info.flatten(),
+            local_info,
+            group=self.cpu_group,
+        )
+        tp0_info = self.global_info[:, 0, :]
+        return tp0_info
+
+    def should_allow_prefill(self, local_prefillable: int) -> bool:
+        tp0_info = self._gather_info(local_prefillable=local_prefillable)
+        global_prefillable = tp0_info[:, 0]
+        global_exists_not_prefillable = global_prefillable.min().item() == 0
+        global_exists_prefillable = global_prefillable.max().item() > 0
+        global_mixed_prefillable = (
+            global_exists_not_prefillable and global_exists_prefillable
+        )
+
+        if global_mixed_prefillable:
+            self.curr_delayed_count += 1
+            if self.curr_delayed_count < self.max_delay_passes:
+                return False
+
+        if _DEBUG_LOG and global_mixed_prefillable:
+            logger.info(
+                f"PrefillDelayer timeout thus not forbid prefill (prefillable: {global_prefillable.sum()})"
+            )
+
+        self.curr_delayed_count = 0
+        return True

python/sglang/srt/managers/scheduler.py

@@ -120,6 +120,7 @@ from sglang.srt.managers.io_struct import (
 )
 from sglang.srt.managers.mm_utils import init_mm_embedding_cache
 from sglang.srt.managers.overlap_utils import FutureMap
+from sglang.srt.managers.prefill_delayer import PrefillDelayer
 from sglang.srt.managers.schedule_batch import (
     FINISH_ABORT,
     ModelWorkerBatch,
@@ -134,7 +135,6 @@ from sglang.srt.managers.schedule_policy import (
     SchedulePolicy,
 )
 from sglang.srt.managers.scheduler_dp_attn_mixin import SchedulerDPAttnMixin
-from sglang.srt.managers.scheduler_enhancer import SchedulerEnhancer
 from sglang.srt.managers.scheduler_input_blocker import SchedulerInputBlocker
 from sglang.srt.managers.scheduler_metrics_mixin import (
     RECORD_STEP_TIME,
@@ -204,7 +204,6 @@ logger = logging.getLogger(__name__)
 TEST_RETRACT = envs.SGLANG_TEST_RETRACT.get()
 TEST_RETRACT_INTERVAL = envs.SGLANG_TEST_RETRACT_INTERVAL.get()
 TEST_RETRACT_NO_PREFILL_BS = envs.SGLANG_TEST_RETRACT_NO_PREFILL_BS.get()
-SCHEDULER_DECREASE_PREFILL_IDLE = envs.SGLANG_SCHEDULER_DECREASE_PREFILL_IDLE.get()
 GRAMMAR_TIMEOUT = float(os.environ.get("SGLANG_GRAMMAR_TIMEOUT", 300))
 
 
@@ -796,14 +795,13 @@ class Scheduler(
             self.enable_priority_scheduling,
             self.schedule_low_priority_values_first,
         )
-        self.schedule_enhancer = None
+        self.prefill_delayer: Optional[PrefillDelayer] = None
-        if SCHEDULER_DECREASE_PREFILL_IDLE:
+        if envs.SGLANG_SCHEDULER_DECREASE_PREFILL_IDLE.get():
-            self.schedule_enhancer = SchedulerEnhancer(
+            self.prefill_delayer = PrefillDelayer(
-                self.dp_size,
+                dp_size=self.dp_size,
-                self.attn_tp_size,
+                attn_tp_size=self.attn_tp_size,
-                self.tp_worker,
+                tp_worker=self.tp_worker,
-                self.max_running_requests,
+                server_args=self.server_args,
-                self.server_args,
             )
         # Enable preemption for priority scheduling.
         self.try_preemption = self.enable_priority_scheduling
@@ -1913,12 +1911,6 @@ class Scheduler(
         return res
 
     def get_new_batch_prefill(self) -> Optional[ScheduleBatch]:
-        if self.schedule_enhancer and not self.schedule_enhancer.get_schedule_decision(
-            self.running_batch
-        ):
-            # Decrease prefill idle as much as possible during high dp load.
-            return None
-
         # Check if the grammar is ready in the grammar queue
         if self.grammar_queue:
             self.move_ready_grammar_requests()
@@ -1927,10 +1919,20 @@ class Scheduler(
             # Reset batch_is_full to try preemption with a prefill adder.
             self.running_batch.batch_is_full = False
 
-        # Handle the cases where prefill is not allowed
+        # The `should_allow_prefill` needs to be called on all ranks since contains communication
-        if (
+        delayer_allow_prefill = (
-            self.running_batch.batch_is_full or len(self.waiting_queue) == 0
+            self.prefill_delayer.should_allow_prefill(
-        ) and self.chunked_req is None:
+                # TODO: consider offline generation cases when there are a lot of waiting requests
+                local_prefillable=len(self.waiting_queue)
+                > 0,
+            )
+            if self.prefill_delayer
+            else True
+        )
+        if (not delayer_allow_prefill) or (
+            (self.running_batch.batch_is_full or len(self.waiting_queue) == 0)
+            and self.chunked_req is None
+        ):
             return None
 
         running_bs = len(self.running_batch.reqs)

python/sglang/srt/managers/scheduler_enhancer.py

@@ -1,59 +0,0 @@
-import torch
-
-
-class SchedulerEnhancer:
-    def __init__(
-        self, dp_size, attn_tp_size, tp_worker, max_running_requests, server_args
-    ):
-        self.dp_size = dp_size
-        self.attn_tp_size = attn_tp_size
-        self.global_batch_size = torch.empty(
-            (self.dp_size, self.attn_tp_size, 1),
-            dtype=torch.int64,
-            device="cpu",
-        )
-        self.cpu_group = tp_worker.get_tp_group().cpu_group
-        self.max_running_requests = max_running_requests
-        self.stable_count = 0
-        # If scheduling is performed 30 times and some dp units are still at full load, the prefill-prioritized scheduling strategy will still be used.
-        self.max_stable_count = 30
-        assert (
-            server_args.schedule_policy == "fcfs"
-        ), f"To use SCHEDULER_DECREASE_PREFILL_IDLE, schedule_policy must be 'fcfs'. '{self.schedule_policy}' is not supported."
-        assert (
-            server_args.enable_dp_attention == True
-        ), f"To use SCHEDULER_DECREASE_PREFILL_IDLE, enable_dp_attention must be enable."
-        assert (
-            server_args.disaggregation_mode == "null"
-        ), f"To use SCHEDULER_DECREASE_PREFILL_IDLE, disaggregation_mode must be null."
-        assert (
-            server_args.disable_overlap_schedule == False
-        ), f"To use SCHEDULER_DECREASE_PREFILL_IDLE, disable_overlap_schedule must be False."
-
-    def get_schedule_info(self, running_batch):
-        local_batch_size = torch.tensor(
-            [
-                running_batch.batch_size(),
-            ],
-            device="cpu",
-            dtype=torch.int64,
-        )
-        torch.distributed.all_gather_into_tensor(
-            self.global_batch_size.flatten(),
-            local_batch_size,
-            group=self.cpu_group,
-        )
-        tp0_info = self.global_batch_size[:, 0, :]
-        return tp0_info
-
-    def get_schedule_decision(self, running_batch):
-        tp0_info = self.get_schedule_info(running_batch)
-        if (
-            int(tp0_info[:, 0].min().item()) < self.max_running_requests
-            and int(tp0_info[:, 0].max().item()) == self.max_running_requests
-        ):
-            self.stable_count += 1
-            if self.stable_count < self.max_stable_count:
-                return False
-        self.stable_count = 0
-        return True

test/srt/run_suite.py

@@ -188,6 +188,7 @@ suites = {
         TestFile("test_profile_v2.py"),
         TestFile("models/test_ministral3_models.py"),
         TestFile("test_mistral_large3_basic.py"),
+        TestFile("test_prefill_delayer.py"),
     ],
 }
 

test/srt/test_prefill_delayer.py

@@ -0,0 +1,76 @@
+import os
+import unittest
+
+from sglang.bench_serving import run_benchmark
+from sglang.srt.environ import envs
+from sglang.srt.utils import kill_process_tree
+from sglang.test.test_utils import (
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    get_benchmark_args,
+    popen_launch_server,
+)
+
+
+class TestPrefillDelayerThroughput(CustomTestCase):
+    def _run_throughput_test(self, with_prefill_delayer: bool):
+        os.environ["SGLANG_PREFILL_DELAYER_DEBUG_LOG"] = "1"
+
+        model = "Qwen/Qwen3-0.6B"
+        base_url = DEFAULT_URL_FOR_TEST
+        other_args = [
+            "--trust-remote-code",
+            "--tp",
+            "8",
+            "--enable-dp-attention",
+            "--dp",
+            "8",
+            "--chunked-prefill-size",
+            "262144",
+            "--mem-fraction-static",
+            "0.6",
+        ]
+
+        # TODO further fix mem leak
+        with envs.SGLANG_SCHEDULER_DECREASE_PREFILL_IDLE.override(
+            with_prefill_delayer
+        ), envs.SGLANG_PREFILL_DELAYER_MAX_DELAY_PASSES.override(
+            100
+        ), envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_IDLE.override(
+            False
+        ):
+            process = popen_launch_server(
+                model,
+                base_url,
+                timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+                other_args=other_args,
+            )
+
+        try:
+            args = get_benchmark_args(
+                base_url=base_url,
+                dataset_name="random",
+                num_prompts=500,
+                random_input_len=30000,
+                random_output_len=256,
+                request_rate=32,
+                tokenizer=model,
+            )
+            res = run_benchmark(args)
+        finally:
+            kill_process_tree(process.pid)
+
+        print(f"=== {with_prefill_delayer=} ===")
+        print(f"Input throughput: {res['input_throughput']:.2f} token/s")
+        print(f"Output throughput: {res['output_throughput']:.2f} token/s")
+
+    def test_1_dp_attention_throughput_with_prefill_delayer(self):
+        self._run_throughput_test(with_prefill_delayer=True)
+
+    def test_2_dp_attention_throughput_without_prefill_delayer(self):
+        self._run_throughput_test(with_prefill_delayer=False)
+
+
+if __name__ == "__main__":
+    unittest.main()