Refactor prefill delayer for clarity and extensibility (#16811)

python/sglang/srt/managers/prefill_delayer.py

@@ -1,11 +1,11 @@
+import dataclasses
 import logging
 import time
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, Optional
+from typing import TYPE_CHECKING, NamedTuple, Optional
 
 import torch
 
-from sglang.srt.environ import envs
 from sglang.srt.utils import get_bool_env_var
 
 if TYPE_CHECKING:
@@ -16,32 +16,48 @@ _DEBUG_LOG = get_bool_env_var("SGLANG_PREFILL_DELAYER_DEBUG_LOG")
 logger = logging.getLogger(__name__)
 
 
-@dataclass
+@dataclass(frozen=True)
-class _DelayInfo:
+class _State:
     delayed_count: int = 0
     start_time: float = field(default_factory=time.perf_counter)
 
+    def bump_delayed_count(self) -> "_State":
+        return dataclasses.replace(self, delayed_count=self.delayed_count + 1)
+
+
+class _NegotiateOutput(NamedTuple):
+    next_state: Optional[_State]
+    input_estimation: str
+    output_allow: bool
+    output_reason: str
+    num_prefillable: int
+
 
 class PrefillDelayer:
     def __init__(
         self,
-        dp_size,
+        dp_size: int,
-        attn_tp_size,
+        attn_tp_size: int,
-        tp_worker,
+        cpu_group,
         server_args,
+        max_delay_passes: int,
         metrics_collector: Optional["SchedulerMetricsCollector"] = None,
     ):
-        self.global_info = torch.empty(
+        self._max_delay_passes = max_delay_passes
+        logger.info(
+            f"PrefillDelayer initialized with max_delay_passes={self._max_delay_passes}"
+        )
+
+        self._global_info_buffer = torch.empty(
             (dp_size, attn_tp_size, 1),
             dtype=torch.int64,
             device="cpu",
         )
-        self.cpu_group = tp_worker.get_tp_group().cpu_group
+        self._cpu_group = cpu_group
 
-        self.max_delay_passes = envs.SGLANG_PREFILL_DELAYER_MAX_DELAY_PASSES.get()
         self._metrics_collector = metrics_collector
 
-        self._curr_delay_info: Optional[_DelayInfo] = None
+        self._curr_state: Optional[_State] = None
 
         assert (
             server_args.enable_dp_attention
@@ -53,40 +69,69 @@ class PrefillDelayer:
             not server_args.disable_overlap_schedule
         ), "To use PrefillDelayer, disable_overlap_schedule must be False."
 
-    def _negotiate_should_allow_prefill(self, local_prefillable: bool) -> bool:
+    def _negotiate_should_allow_prefill(
-        tp0_info = self._gather_info(local_prefillable=local_prefillable)
+        self, local_prefillable: bool
-        global_prefillable = tp0_info[:, 0]
+    ) -> _NegotiateOutput:
-        global_exists_not_prefillable = global_prefillable.min().item() == 0
+        out = self._negotiate_should_allow_prefill_pure(
-        global_exists_prefillable = global_prefillable.max().item() > 0
+            prev_state=self._curr_state,
-        global_mixed_prefillable = (
+            local_prefillable=local_prefillable,
-            global_exists_not_prefillable and global_exists_prefillable
+        )
+        self._curr_state = out.next_state
+        return out
+
+    def _negotiate_should_allow_prefill_pure(
+        self,
+        prev_state: Optional[_State],
+        local_prefillable: bool,
+    ) -> _NegotiateOutput:
+        global_prefillable = self._gather_info(local_prefillable=local_prefillable)
+
+        if global_prefillable.min().item() > 0:
+            prefillable_status = "all"
+        elif global_prefillable.max().item() == 0:
+            prefillable_status = "none"
+        else:
+            prefillable_status = "mixed"
+        debug_info = dict(
+            input_estimation=prefillable_status,
+            num_prefillable=global_prefillable.sum().item(),
         )
 
-        if global_mixed_prefillable:
+        if prefillable_status == "all":
-            if self._curr_delay_info is None:
+            exist_previous_wait = prev_state is not None
-                self._curr_delay_info = _DelayInfo()
+            return _NegotiateOutput(
-            self._curr_delay_info.delayed_count += 1
+                next_state=None,
-            if self._curr_delay_info.delayed_count < self.max_delay_passes:
+                output_allow=True,
-                return False
+                output_reason="wait_success" if exist_previous_wait else "no_wait",
-
+                **debug_info,
-        is_timeout = global_mixed_prefillable
-        if _DEBUG_LOG and is_timeout:
-            logger.info(
-                f"PrefillDelayer timeout thus not forbid prefill (prefillable: {global_prefillable.sum()})"
             )
-
+        elif prefillable_status == "none":
-        self._record_metrics(is_timeout=is_timeout)
+            return _NegotiateOutput(
-        self._curr_delay_info = None
+                next_state=None,
-        return True
+                output_allow=True,
-
+                output_reason="",
-    def _record_metrics(self, is_timeout: bool) -> None:
+                **debug_info,
-        if self._curr_delay_info is not None and self._metrics_collector is not None:
-            wait_seconds = time.perf_counter() - self._curr_delay_info.start_time
-            self._metrics_collector.observe_prefill_delayer_wait(
-                forward_passes=self._curr_delay_info.delayed_count,
-                wait_seconds=wait_seconds,
-                is_timeout=is_timeout,
             )
+        elif prefillable_status == "mixed":
+            prev_delayed_count = prev_state.delayed_count if prev_state else 0
+            if prev_delayed_count < self._max_delay_passes - 1:
+                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,
+                )
+            else:
+                return _NegotiateOutput(
+                    next_state=None,
+                    output_allow=True,
+                    output_reason="wait_timeout",
+                    **debug_info,
+                )
+        else:
+            raise NotImplementedError
 
     def _gather_info(self, local_prefillable: bool):
         local_info = torch.tensor(
@@ -95,34 +140,69 @@ class PrefillDelayer:
             dtype=torch.int64,
         )
         torch.distributed.all_gather_into_tensor(
-            self.global_info.flatten(),
+            self._global_info_buffer.flatten(),
             local_info,
-            group=self.cpu_group,
+            group=self._cpu_group,
         )
-        tp0_info = self.global_info[:, 0, :]
+        tp0_info = self._global_info_buffer[:, 0, :]
-        return tp0_info
+        return tp0_info[:, 0]
 
 
 class PrefillDelayerSinglePassExecutor:
     def __init__(self, prefill_delayer: PrefillDelayer):
         self._prefill_delayer = prefill_delayer
-        self._result: Optional[bool] = None
+        self._result: Optional[_NegotiateOutput] = None
 
     @property
     def _called(self) -> bool:
         return self._result is not None
 
-    def __enter__(self):
+    def finalize(self, *, actual_prefill: bool):
-        return self
-
-    def __exit__(self, exc_type, exc_val, exc_tb):
         if not self._called:
             self.negotiate_should_allow_prefill(local_prefillable=False)
-        return False
+
+        _record_single_pass_result(
+            actual_execution=actual_prefill,
+            output=self._result,
+            metrics_collector=self._prefill_delayer._metrics_collector,
+        )
 
     def negotiate_should_allow_prefill(self, local_prefillable: bool) -> bool:
         if not self._called:
             self._result = self._prefill_delayer._negotiate_should_allow_prefill(
-                local_prefillable=local_prefillable
+                local_prefillable=local_prefillable,
             )
-        return self._result
+        return self._result.output_allow
+
+
+def _record_single_pass_result(
+    actual_execution: bool,
+    output: _NegotiateOutput,
+    metrics_collector: Optional["SchedulerMetricsCollector"],
+) -> None:
+    if _DEBUG_LOG:
+        if output.output_allow and (output.output_reason == "wait_timeout"):
+            logger.info(
+                f"PrefillDelayer timeout thus not forbid prefill "
+                f"(num_prefillable={output.num_prefillable}, "
+                f"actual_execution={actual_execution})"
+            )
+        else:
+            assert output.output_reason in {
+                "",
+                "wait_success",
+                "no_wait",
+                "delay",
+            }
+
+    if metrics_collector is not None:
+        if (s := output.next_state) is not None:
+            wait_seconds = time.perf_counter() - s.start_time
+            forward_passes = s.delayed_count
+        else:
+            wait_seconds = forward_passes = 0
+        metrics_collector.observe_prefill_delayer_wait(
+            forward_passes=forward_passes,
+            wait_seconds=wait_seconds,
+            is_timeout=(output.output_reason == "wait_timeout"),
+        )

python/sglang/srt/managers/scheduler.py

@@ -21,7 +21,6 @@ import sys
 import time
 from collections import deque
 from concurrent import futures
-from contextlib import nullcontext
 from dataclasses import dataclass
 from http import HTTPStatus
 from typing import Any, Deque, Dict, List, Optional, Tuple, Union
@@ -770,11 +769,12 @@ class Scheduler(
             self.prefill_delayer = PrefillDelayer(
                 dp_size=self.dp_size,
                 attn_tp_size=self.attn_tp_size,
-                tp_worker=self.tp_worker,
+                cpu_group=self.tp_worker.get_tp_group().cpu_group,
                 server_args=self.server_args,
                 metrics_collector=(
                     self.metrics_collector if self.enable_metrics else None
                 ),
+                max_delay_passes=envs.SGLANG_PREFILL_DELAYER_MAX_DELAY_PASSES.get(),
             )
         # Enable preemption for priority scheduling.
         self.try_preemption = self.enable_priority_scheduling
@@ -1844,15 +1844,21 @@ class Scheduler(
         return res
 
     def get_new_batch_prefill(self) -> Optional[ScheduleBatch]:
-        with (
+        prefill_delayer_single_pass = None
-            PrefillDelayerSinglePassExecutor(self.prefill_delayer)
+        if self.prefill_delayer:
-            if self.prefill_delayer
+            prefill_delayer_single_pass = PrefillDelayerSinglePassExecutor(
-            else nullcontext()
+                self.prefill_delayer
-        ) as prefill_delayer_single_pass:
-            return self._get_new_batch_prefill_raw(
-                prefill_delayer_single_pass=prefill_delayer_single_pass
             )
 
+        ret = self._get_new_batch_prefill_raw(
+            prefill_delayer_single_pass=prefill_delayer_single_pass
+        )
+
+        if self.prefill_delayer:
+            prefill_delayer_single_pass.finalize(actual_prefill=ret is not None)
+
+        return ret
+
     def _get_new_batch_prefill_raw(
         self, prefill_delayer_single_pass: Optional[PrefillDelayerSinglePassExecutor]
     ) -> Optional[ScheduleBatch]: