Refactor prefill delayer for clarity and extensibility (#16811)

2026-01-09 22:46:41 +08:00
parent 068abe7e40
commit dd24ba9084
2 changed files with 148 additions and 62 deletions
--- a/python/sglang/srt/managers/prefill_delayer.py
+++ b/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
-class _DelayInfo:
+@dataclass(frozen=True)
+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,
-        attn_tp_size,
-        tp_worker,
+        dp_size: int,
+        attn_tp_size: int,
+        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:
-        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
+    def _negotiate_should_allow_prefill(
+        self, local_prefillable: bool
+    ) -> _NegotiateOutput:
+        out = self._negotiate_should_allow_prefill_pure(
+            prev_state=self._curr_state,
+            local_prefillable=local_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 self._curr_delay_info is None:
-                self._curr_delay_info = _DelayInfo()
-            self._curr_delay_info.delayed_count += 1
-            if self._curr_delay_info.delayed_count < self.max_delay_passes:
-                return False
-
-        is_timeout = global_mixed_prefillable
-        if _DEBUG_LOG and is_timeout:
-            logger.info(
-                f"PrefillDelayer timeout thus not forbid prefill (prefillable: {global_prefillable.sum()})"
+        if prefillable_status == "all":
+            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,
            )
-
-        self._record_metrics(is_timeout=is_timeout)
-        self._curr_delay_info = None
-        return True
-
-    def _record_metrics(self, is_timeout: bool) -> None:
-        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 == "none":
+            return _NegotiateOutput(
+                next_state=None,
+                output_allow=True,
+                output_reason="",
+                **debug_info,
            )
+        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, :]
-        return tp0_info
+        tp0_info = self._global_info_buffer[:, 0, :]
+        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):
-        return self
-
-    def __exit__(self, exc_type, exc_val, exc_tb):
+    def finalize(self, *, actual_prefill: bool):
        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"),
+        )
--- a/python/sglang/srt/managers/scheduler.py
+++ b/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 (
-            PrefillDelayerSinglePassExecutor(self.prefill_delayer)
-            if self.prefill_delayer
-            else nullcontext()
-        ) as prefill_delayer_single_pass:
-            return self._get_new_batch_prefill_raw(
-                prefill_delayer_single_pass=prefill_delayer_single_pass
+        prefill_delayer_single_pass = None
+        if self.prefill_delayer:
+            prefill_delayer_single_pass = PrefillDelayerSinglePassExecutor(
+                self.prefill_delayer
            )

+        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]: