[scheduler] enhance scheduler in dp_attention mixed case with spec (#14201)

python/sglang/srt/environ.py

@@ -195,6 +195,7 @@ class Envs:
     SGLANG_SCHEDULER_MAX_RECV_PER_POLL = EnvInt(-1)
     SGLANG_EXPERIMENTAL_CPP_RADIX_TREE = EnvBool(False)
     SGLANG_DYNAMIC_CHUNKING_SMOOTH_FACTOR = EnvFloat(0.75)
+    SGLANG_SCHEDULER_DECREASE_PREFILL_IDLE = EnvBool(False)
 
     # Test: pd-disaggregation
     SGLANG_TEST_PD_DISAGG_BACKEND = EnvStr("mooncake")

python/sglang/srt/managers/scheduler.py

@@ -132,6 +132,7 @@ 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,
@@ -201,6 +202,7 @@ 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))
 
 
@@ -507,6 +509,15 @@ class Scheduler(
             self.enable_priority_scheduling,
             self.schedule_low_priority_values_first,
         )
+        self.schedule_enhancer = None
+        if SCHEDULER_DECREASE_PREFILL_IDLE:
+            self.schedule_enhancer = SchedulerEnhancer(
+                self.dp_size,
+                self.attn_tp_size,
+                self.tp_worker,
+                self.max_running_requests,
+                server_args,
+            )
         # Enable preemption for priority scheduling.
         self.try_preemption = self.enable_priority_scheduling
         self.init_new_token_ratio = min(
@@ -1745,6 +1756,11 @@ 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()

python/sglang/srt/managers/scheduler_enhancer.py

@@ -0,0 +1,59 @@
+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