Added the prefill delayer policy: The prefill deplay range is expanded. (#17456)

Co-authored-by: sglang-npu-bot <sglangnpu@163.com>

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
 

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

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 = (