diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index ff380e002..eadfdf566 100644 --- a/python/sglang/srt/environ.py +++ b/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") diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index fdb817517..db884c4a2 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/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() diff --git a/python/sglang/srt/managers/scheduler_enhancer.py b/python/sglang/srt/managers/scheduler_enhancer.py new file mode 100644 index 000000000..d7799aa8c --- /dev/null +++ b/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