fix(scheduler): track max prefill batch size over recent real admissions (#34284)

python/sglang/srt/environ.py

@@ -534,6 +534,7 @@ class Envs:
     # "legacy" (rectangular grid, max_extend_len-shaped).
     # Internal/testing only - users should not need to change this.
     SGLANG_PREFILL_TILE_BUDGET_MODE = EnvStr("compact")
+    SGLANG_PREFILL_DELAYER_MAX_PREFILL_BS_WINDOW_SIZE = EnvInt(16)
 
     # ===================================================================
     # Scheduler polling, timeouts, and output

python/sglang/srt/managers/prefill_delayer.py

@@ -1,6 +1,7 @@
 import dataclasses
 import logging
 import time
+from collections import deque
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, NamedTuple, Optional
 
@@ -17,6 +18,33 @@ _DEBUG_LOG = get_bool_env_var("SGLANG_PREFILL_DELAYER_DEBUG_LOG")
 logger = logging.getLogger(__name__)
 
 
+class RecentPrefillBatchSizeTracker:
+    """Track the largest of the latest non-empty prefill attempts.
+
+    The default window keeps 16 attempts. Successful admissions use their
+    actual batch size; rejected attempts use a conservative local estimate.
+    Decode-only and idle scheduler passes do not age the high-watermark.
+    """
+
+    def __init__(self, window_size: int = 16):
+        if window_size <= 0:
+            raise ValueError(f"window_size must be positive, got {window_size}")
+        self._recent_attempt_sizes = deque(maxlen=window_size)
+
+    @property
+    def max_prefill_bs(self) -> int:
+        return max(self._recent_attempt_sizes, default=0)
+
+    def observe_attempt(self, attempted_prefill_bs: int) -> int:
+        if attempted_prefill_bs <= 0:
+            raise ValueError(
+                "attempted_prefill_bs must be positive for a non-empty attempt, "
+                f"got {attempted_prefill_bs}"
+            )
+        self._recent_attempt_sizes.append(attempted_prefill_bs)
+        return self.max_prefill_bs
+
+
 @dataclass(frozen=True)
 class _State:
     delayed_count: int = 0
@@ -347,21 +375,45 @@ class PrefillDelayerSinglePassExecutor:
         self._prefill_delayer = prefill_delayer
         self._token_usage = token_usage
         self._result: Optional[_NegotiateOutput] = None
+        self._attempted_prefill_bs = 0
 
     @property
     def _called(self) -> bool:
         return self._result is not None
 
-    def finalize(self, *, actual_prefill: bool):
+    def finalize(self, *, actual_prefill_bs: int) -> int:
         if not self._called:
             self.negotiate_should_allow_prefill(local_prefillable=False)
 
         _record_single_pass_result(
-            actual_execution=actual_prefill,
+            actual_execution=actual_prefill_bs > 0,
             output=self._result,
             metrics_collector=self._prefill_delayer._metrics_collector,
             debug_log_enabled=self._prefill_delayer._debug_log_enabled,
         )
+        return actual_prefill_bs or self._attempted_prefill_bs
+
+    def _estimate_attempted_prefill_bs(
+        self,
+        *,
+        running_batch: int,
+        max_running_requests: int,
+        waiting_queue_len: int,
+    ) -> int:
+        local_max_running_requests = max_running_requests
+        if not self._prefill_delayer.enable_dp_attention:
+            local_max_running_requests = (
+                max_running_requests + self._prefill_delayer.dp_size - 1
+            ) // self._prefill_delayer.dp_size
+
+        # The delayer negotiates before PrefillAdder materializes can_run_list,
+        # so a rejected pass has no exact batch size. This upper bound is exact
+        # when the waiting queue is the limiter (for example, two queued
+        # requests after a cached BS=10 spike), and it never exceeds the local
+        # request slots available to the candidate batch.
+        free_slots = max(local_max_running_requests - running_batch, 1)
+        non_empty_queue_len = max(waiting_queue_len, 1)
+        return min(non_empty_queue_len, free_slots)
 
     def negotiate_should_allow_prefill(
         self,
@@ -371,6 +423,15 @@ class PrefillDelayerSinglePassExecutor:
         max_running_requests: int = 0,
         waiting_queue_len: int = 0,
     ) -> bool:
+        if local_prefillable:
+            self._attempted_prefill_bs = max(
+                self._attempted_prefill_bs,
+                self._estimate_attempted_prefill_bs(
+                    running_batch=running_batch,
+                    max_running_requests=max_running_requests,
+                    waiting_queue_len=waiting_queue_len,
+                ),
+            )
         if not self._called:
             self._result = self._prefill_delayer._negotiate_should_allow_prefill(
                 local_prefillable=local_prefillable,

python/sglang/srt/managers/scheduler.py

@@ -188,6 +188,7 @@ from sglang.srt.managers.overlap_utils import (
 from sglang.srt.managers.prefill_delayer import (
     PrefillDelayer,
     PrefillDelayerSinglePassExecutor,
+    RecentPrefillBatchSizeTracker,
 )
 from sglang.srt.managers.rust_server import RustServer
 from sglang.srt.managers.schedule_batch import (
@@ -1211,7 +1212,10 @@ class Scheduler(
             self.schedule_low_priority_values_first,
         )
         self.prefill_delayer: Optional[PrefillDelayer] = None
-        self.max_prefill_bs: float = 0.0
+        self.prefill_bs_tracker = RecentPrefillBatchSizeTracker(
+            window_size=envs.SGLANG_PREFILL_DELAYER_MAX_PREFILL_BS_WINDOW_SIZE.get()
+        )
+        self.max_prefill_bs: int = 0
         if get_schedule().enable_prefill_delayer:
             if get_disagg().disaggregation_mode == "decode":
                 logger.info(
@@ -3156,11 +3160,6 @@ class Scheduler(
     def get_new_batch_prefill(self, running_batch: ScheduleBatch) -> NextBatchPlan:
         prefill_delayer_single_pass = None
         if self.prefill_delayer:
-            # Decay the max-prefill-bs high-watermark once per pass so one
-            # unusually large admission burst does not permanently raise the
-            # slot_condition bar in the delayer (0.998/pass ~= half-life of
-            # ~350 forward passes).
-            self.max_prefill_bs *= 0.998
             # Get max usage across all pools for prefill delay decision
             max_pool_usage = (
                 self.pool_stats_observer.get_pool_stats().get_max_pool_usage()
@@ -3175,7 +3174,13 @@ class Scheduler(
         )
 
         if self.prefill_delayer:
-            prefill_delayer_single_pass.finalize(actual_prefill=ret is not None)
+            observed_prefill_bs = prefill_delayer_single_pass.finalize(
+                actual_prefill_bs=ret.batch_size() if ret is not None else 0
+            )
+            if observed_prefill_bs > 0:
+                self.max_prefill_bs = self.prefill_bs_tracker.observe_attempt(
+                    observed_prefill_bs
+                )
 
         return NextBatchPlan(batch_to_run=ret, running_batch=running_batch)
 
@@ -3400,7 +3405,6 @@ class Scheduler(
             self.chunked_req is None or len(can_run_list) != 1
         )
 
-        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 = (