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

test/registered/npu/performance/kimi_k2_6/test_npu_kimi_k2_6_w4a8_8p_in3k5_out1k5_20ms.py

@@ -27,6 +27,7 @@ KIMI_K2_6_ENVS = {
     "HCCL_OP_EXPANSION_MODE": "AIV",
     "SGLANG_NPU_USE_MLAPO": "1",
     "SGLANG_NPU_USE_MULTI_STREAM": "1",
+    "SGLANG_PREFILL_DELAYER_MAX_PREFILL_BS_WINDOW_SIZE": "64",
 }
 
 KIMI_K2_6_OTHER_ARGS = [

test/registered/npu/performance/minimax_m2_5/test_npu_minimax_m2_5_w8a8_4p_in64k_out1k_prefix90_50ms.py

@@ -25,6 +25,7 @@ MINIMAX_M2_5_W8A8_4P_IN64K_OUT1K_PREFIX90_ENVS = {
     "DEEP_NORMAL_MODE_USE_INT8_QUANT": "1",
     "DEEPEP_HCCL_BUFFSIZE": "1024",
     "SGLANG_EXTERNAL_MODEL_PACKAGE": "custom_eagle3",
+    "SGLANG_PREFILL_DELAYER_MAX_PREFILL_BS_WINDOW_SIZE": "64",
     "PYTHONPATH": f"{MINIMAX_M2_5_EAGLE3_MODEL_PATH}:{os.environ.get('PYTHONPATH', '')}",
 }
 

test/registered/npu/performance/qwen3_6_27b/test_npu_qwen3_6_27b_1p_in1080p_30_out256_50ms.py

@@ -25,6 +25,7 @@ QWEN3_6_27B_1080P_ENVS = {
     "SGLANG_NPU_PROFILING": "0",
     "SGLANG_NPU_PROFILING_STAGE": "prefill",
     "SGLANG_ENABLE_OVERLAP_PLAN_STREAM": "1",
+    "SGLANG_PREFILL_DELAYER_MAX_PREFILL_BS_WINDOW_SIZE": "64",
     "ASCEND_USE_FIA": "1",
 }
 

test/registered/unit/managers/test_prefill_delayer_high_watermark.py

@@ -0,0 +1,143 @@
+import unittest
+from unittest.mock import MagicMock
+
+from sglang.srt.managers.prefill_delayer import (
+    PrefillDelayerSinglePassExecutor,
+    RecentPrefillBatchSizeTracker,
+    _NegotiateOutput,
+)
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=2, suite="base-a-test-cpu")
+
+
+class TestPrefillDelayerHighWatermark(CustomTestCase):
+    def test_peak_expires_after_recent_attempt_window(self):
+        tracker = RecentPrefillBatchSizeTracker(window_size=4)
+
+        self.assertEqual(tracker.observe_attempt(100), 100)
+        for attempted_prefill_bs in [2, 1, 2]:
+            self.assertEqual(tracker.observe_attempt(attempted_prefill_bs), 100)
+
+        self.assertEqual(tracker.observe_attempt(2), 2)
+
+    def test_recurring_small_peak_remains_effective(self):
+        tracker = RecentPrefillBatchSizeTracker(window_size=4)
+
+        for attempted_prefill_bs in [2, 1, 1, 2, 1, 1, 2]:
+            self.assertEqual(tracker.observe_attempt(attempted_prefill_bs), 2)
+
+    def test_rejected_non_empty_attempt_advances_window(self):
+        tracker = RecentPrefillBatchSizeTracker(window_size=4)
+        self.assertEqual(tracker.observe_attempt(10), 10)
+
+        delayer = MagicMock()
+        delayer.enable_dp_attention = True
+        delayer.dp_size = 1
+        delayer._metrics_collector = None
+        delayer._debug_log_enabled = False
+        delayer._negotiate_should_allow_prefill.return_value = _NegotiateOutput(
+            next_state=None,
+            input_estimation="all",
+            output_allow=False,
+            output_reason="delay",
+            num_prefillable=1,
+            num_token_watermark_force_allow=0,
+        )
+
+        for _ in range(4):
+            executor = PrefillDelayerSinglePassExecutor(delayer, token_usage=0.9)
+            self.assertFalse(
+                executor.negotiate_should_allow_prefill(
+                    local_prefillable=True,
+                    running_batch=15,
+                    max_prefill_bs=tracker.max_prefill_bs,
+                    max_running_requests=20,
+                    waiting_queue_len=2,
+                )
+            )
+            attempted_prefill_bs = executor.finalize(actual_prefill_bs=0)
+            self.assertEqual(attempted_prefill_bs, 2)
+            tracker.observe_attempt(attempted_prefill_bs)
+
+        self.assertEqual(tracker.max_prefill_bs, 2)
+
+    def test_print_rejected_attempt_estimates_after_unusual_peak(self):
+        window_size = 16
+
+        for steady_prefill_bs in (2, 3, 4):
+            with self.subTest(steady_prefill_bs=steady_prefill_bs):
+                tracker = RecentPrefillBatchSizeTracker(window_size=window_size)
+                self.assertEqual(tracker.observe_attempt(100), 100)
+
+                delayer = MagicMock()
+                delayer.enable_dp_attention = True
+                delayer.dp_size = 1
+                delayer._metrics_collector = None
+                delayer._debug_log_enabled = False
+                delayer._negotiate_should_allow_prefill.return_value = _NegotiateOutput(
+                    next_state=None,
+                    input_estimation="all",
+                    output_allow=False,
+                    output_reason="delay",
+                    num_prefillable=1,
+                    num_token_watermark_force_allow=0,
+                )
+
+                print(
+                    f"\nunusual_prefill_bs=100, "
+                    f"subsequent_prefill_bs={steady_prefill_bs}, "
+                    f"window_size={window_size}",
+                    flush=True,
+                )
+                print(
+                    "round | waiting_bs | estimated_rejected_bs | "
+                    "high_watermark_before | high_watermark_after",
+                    flush=True,
+                )
+
+                for round_index in range(1, window_size + 1):
+                    high_watermark_before = tracker.max_prefill_bs
+                    executor = PrefillDelayerSinglePassExecutor(
+                        delayer, token_usage=0.9
+                    )
+                    self.assertFalse(
+                        executor.negotiate_should_allow_prefill(
+                            local_prefillable=True,
+                            running_batch=20,
+                            max_prefill_bs=high_watermark_before,
+                            max_running_requests=128,
+                            waiting_queue_len=steady_prefill_bs,
+                        )
+                    )
+                    estimated_prefill_bs = executor.finalize(actual_prefill_bs=0)
+                    high_watermark_after = tracker.observe_attempt(estimated_prefill_bs)
+
+                    print(
+                        f"{round_index:>5} | {steady_prefill_bs:>10} | "
+                        f"{estimated_prefill_bs:>21} | "
+                        f"{high_watermark_before:>21} | "
+                        f"{high_watermark_after:>20}",
+                        flush=True,
+                    )
+                    self.assertEqual(estimated_prefill_bs, steady_prefill_bs)
+                    expected_high_watermark = (
+                        100 if round_index < window_size else steady_prefill_bs
+                    )
+                    self.assertEqual(high_watermark_after, expected_high_watermark)
+
+    def test_rejects_empty_attempts(self):
+        with self.assertRaisesRegex(ValueError, "window_size must be positive"):
+            RecentPrefillBatchSizeTracker(window_size=0)
+
+        tracker = RecentPrefillBatchSizeTracker(window_size=4)
+
+        with self.assertRaisesRegex(ValueError, "non-empty attempt"):
+            tracker.observe_attempt(0)
+
+        self.assertEqual(tracker.max_prefill_bs, 0)
+
+
+if __name__ == "__main__":
+    unittest.main()