diff --git a/python/sglang/srt/managers/prefill_delayer.py b/python/sglang/srt/managers/prefill_delayer.py index 37a1f63f2..d3399c419 100644 --- a/python/sglang/srt/managers/prefill_delayer.py +++ b/python/sglang/srt/managers/prefill_delayer.py @@ -44,6 +44,7 @@ class PrefillDelayer: max_delay_passes: int, token_usage_low_watermark: Optional[float], metrics_collector: Optional["SchedulerMetricsCollector"] = None, + device: Optional["torch.device"] = "cpu", ): self._max_delay_passes = max_delay_passes self._token_usage_low_watermark = token_usage_low_watermark @@ -52,21 +53,20 @@ class PrefillDelayer: f"max_delay_passes={self._max_delay_passes} " f"token_usage_low_watermark={self._token_usage_low_watermark}" ) - + # The global_info contains four pieces of information: + # prefillable, token_watermark_force_allow, running_batch, and max_prefill_bs. self._global_info_buffer = torch.empty( - (dp_size, attn_tp_size, 2), + (dp_size, attn_tp_size, 4), dtype=torch.int64, - device="cpu", + device=device, ) + self.enable_dp_attention = server_args.enable_dp_attention self._cpu_group = cpu_group self._metrics_collector = metrics_collector self._curr_state: Optional[_State] = None - 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." @@ -75,12 +75,16 @@ class PrefillDelayer: ), "To use PrefillDelayer, disable_overlap_schedule must be False." def _negotiate_should_allow_prefill( - self, local_prefillable: bool, token_usage: float + self, + local_prefillable: bool, + token_usage: float, + **kwargs, ) -> _NegotiateOutput: out = self._negotiate_should_allow_prefill_pure( prev_state=self._curr_state, local_prefillable=local_prefillable, token_usage=token_usage, + **kwargs, ) self._curr_state = out.next_state return out @@ -91,6 +95,7 @@ class PrefillDelayer: prev_state: Optional[_State], local_prefillable: bool, token_usage: float, + **kwargs, ) -> _NegotiateOutput: # Compute local states local_token_watermark_force_allow = ( @@ -100,10 +105,15 @@ class PrefillDelayer: ) # Gather global states - global_prefillable, global_token_watermark_force_allow = self._gather_info( + tp0_info = self._gather_info( local_prefillable=local_prefillable, local_token_watermark_force_allow=local_token_watermark_force_allow, + **kwargs, ) + global_prefillable = tp0_info[:, 0] + global_token_watermark_force_allow = tp0_info[:, 1] + global_running_batch = tp0_info[:, 2] + global_max_prefill_bs = tp0_info[:, 3] # Compute derived global states if global_prefillable.min().item() > 0: @@ -123,6 +133,29 @@ class PrefillDelayer: # Compute outputs if prefillable_status == "all": + if kwargs is None: + exist_previous_wait = prev_state is not None + return _NegotiateOutput( + next_state=None, + output_allow=True, + output_reason="wait_success" if exist_previous_wait else "no_wait", + **debug_info, + ) + + max_running_requests = kwargs.get("max_running_requests", 0) + if ( + max_running_requests - global_running_batch.max().item() + < global_max_prefill_bs.max().item() + and not self.enable_dp_attention + ): + next_state = prev_state or _State() + next_state = next_state.bump_delayed_count() + return _NegotiateOutput( + next_state=next_state, + output_allow=False, + output_reason="delay", + **debug_info, + ) exist_previous_wait = prev_state is not None return _NegotiateOutput( next_state=None, @@ -168,10 +201,15 @@ class PrefillDelayer: raise NotImplementedError def _gather_info( - self, local_prefillable: bool, local_token_watermark_force_allow: bool + self, local_prefillable: bool, local_token_watermark_force_allow: bool, **kwargs ): local_info = torch.tensor( - [int(local_prefillable), int(local_token_watermark_force_allow)], + [ + int(local_prefillable), + int(local_token_watermark_force_allow), + kwargs.get("running_batch", 0), + kwargs.get("max_prefill_bs", 0), + ], device="cpu", dtype=torch.int64, ) @@ -181,7 +219,7 @@ class PrefillDelayer: group=self._cpu_group, ) tp0_info = self._global_info_buffer[:, 0, :] - return tp0_info[:, 0], tp0_info[:, 1] + return tp0_info class PrefillDelayerSinglePassExecutor: @@ -204,11 +242,12 @@ class PrefillDelayerSinglePassExecutor: metrics_collector=self._prefill_delayer._metrics_collector, ) - def negotiate_should_allow_prefill(self, local_prefillable: bool) -> bool: + def negotiate_should_allow_prefill(self, local_prefillable: bool, **kwargs) -> bool: if not self._called: self._result = self._prefill_delayer._negotiate_should_allow_prefill( local_prefillable=local_prefillable, token_usage=self._token_usage, + **kwargs, ) return self._result.output_allow diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index 1b137e52b..cbb32868b 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -381,6 +381,8 @@ class PrefillAdder: rem_chunk_tokens: Optional[int], mixed_with_decode_tokens: int = 0, priority_scheduling_preemption_threshold: int = 0, + max_prefill_bs: int = 0, + max_running_requests: Optional[int] = None, prefill_max_requests: Optional[int] = None, prefill_delayer_single_pass: Optional[PrefillDelayerSinglePassExecutor] = None, dllm_config: Optional[DllmConfig] = None, @@ -427,8 +429,10 @@ class PrefillAdder: priority_scheduling_preemption_threshold ) self.nsa_prefill_cp_in_seq_split = is_nsa_prefill_cp_in_seq_split() + self.max_running_requests = max_running_requests self.prefill_max_requests = prefill_max_requests self.prefill_delayer_single_pass = prefill_delayer_single_pass + self.max_prefill_bs = max_prefill_bs def _init_dllm_meta(self, dllm_config: DllmConfig): self.dllm_block_size = dllm_config.block_size @@ -767,10 +771,12 @@ class PrefillAdder: if input_tokens >= self.rem_input_tokens and len(self.can_run_list) != 0: return AddReqResult.OTHER - if (self.prefill_delayer_single_pass is not None) and ( not self.prefill_delayer_single_pass.negotiate_should_allow_prefill( - local_prefillable=True + local_prefillable=True, + running_batch=self.running_batch.batch_size(), + max_prefill_bs=self.max_prefill_bs, + max_running_requests=self.max_running_requests, ) ): return AddReqResult.OTHER diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index ab51a1a52..714d38b48 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -794,6 +794,7 @@ class Scheduler( self.schedule_low_priority_values_first, ) self.prefill_delayer: Optional[PrefillDelayer] = None + self.max_prefill_bs: int = 0 if self.server_args.enable_prefill_delayer: self.prefill_delayer = PrefillDelayer( dp_size=self.dp_size, @@ -805,6 +806,11 @@ class Scheduler( ), max_delay_passes=self.server_args.prefill_delayer_max_delay_passes, token_usage_low_watermark=self.server_args.prefill_delayer_token_usage_low_watermark, + device=( + self.tp_group.device + if self.server_args.disable_overlap_schedule + else "cpu" + ), ) # Enable preemption for priority scheduling. self.try_preemption = self.enable_priority_scheduling @@ -2049,6 +2055,8 @@ class Scheduler( chunked_prefill_size, running_bs if self.is_mixed_chunk else 0, self.priority_scheduling_preemption_threshold, + max_prefill_bs=self.max_prefill_bs, + max_running_requests=self.max_running_requests, prefill_max_requests=self.server_args.prefill_max_requests, prefill_delayer_single_pass=prefill_delayer_single_pass, dllm_config=self.dllm_config, @@ -2163,6 +2171,7 @@ class Scheduler( self.spec_algorithm, chunked_req=self.chunked_req, ) + self.max_prefill_bs = max(self.max_prefill_bs, len(can_run_list)) if self.enable_hierarchical_cache: # todo (zhiqiang): disable cuda graph execution if hicache loading triggered new_batch.hicache_consumer_index = (