diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 4fb5d8f88..a5f1d21e4 100644 --- a/python/sglang/srt/environ.py +++ b/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 diff --git a/python/sglang/srt/managers/prefill_delayer.py b/python/sglang/srt/managers/prefill_delayer.py new file mode 100644 index 000000000..e33f68518 --- /dev/null +++ b/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 diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 2b6f078c6..ddd8fdcb3 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/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 - if SCHEDULER_DECREASE_PREFILL_IDLE: - self.schedule_enhancer = SchedulerEnhancer( - self.dp_size, - self.attn_tp_size, - self.tp_worker, - self.max_running_requests, - self.server_args, + self.prefill_delayer: Optional[PrefillDelayer] = None + if envs.SGLANG_SCHEDULER_DECREASE_PREFILL_IDLE.get(): + self.prefill_delayer = PrefillDelayer( + dp_size=self.dp_size, + attn_tp_size=self.attn_tp_size, + tp_worker=self.tp_worker, + 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 - if ( - self.running_batch.batch_is_full or len(self.waiting_queue) == 0 - ) and self.chunked_req is None: + # The `should_allow_prefill` needs to be called on all ranks since contains communication + delayer_allow_prefill = ( + self.prefill_delayer.should_allow_prefill( + # 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) diff --git a/python/sglang/srt/managers/scheduler_enhancer.py b/python/sglang/srt/managers/scheduler_enhancer.py deleted file mode 100644 index d7799aa8c..000000000 --- a/python/sglang/srt/managers/scheduler_enhancer.py +++ /dev/null @@ -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 diff --git a/test/srt/run_suite.py b/test/srt/run_suite.py index a826122a5..890b691c5 100644 --- a/test/srt/run_suite.py +++ b/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"), ], } diff --git a/test/srt/test_prefill_delayer.py b/test/srt/test_prefill_delayer.py new file mode 100644 index 000000000..d0e75d9c6 --- /dev/null +++ b/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()