diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 5f8bad4ac..d0089ebe1 100644 --- a/python/sglang/srt/environ.py +++ b/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 diff --git a/python/sglang/srt/managers/prefill_delayer.py b/python/sglang/srt/managers/prefill_delayer.py index 650048791..c00726d5b 100644 --- a/python/sglang/srt/managers/prefill_delayer.py +++ b/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, diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 3b66940bd..09b0a9421 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/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 = ( diff --git a/test/registered/npu/performance/kimi_k2_6/test_npu_kimi_k2_6_w4a8_8p_in3k5_out1k5_20ms.py b/test/registered/npu/performance/kimi_k2_6/test_npu_kimi_k2_6_w4a8_8p_in3k5_out1k5_20ms.py index 8ad487567..d8b798cb6 100644 --- a/test/registered/npu/performance/kimi_k2_6/test_npu_kimi_k2_6_w4a8_8p_in3k5_out1k5_20ms.py +++ b/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 = [ diff --git a/test/registered/npu/performance/minimax_m2_5/test_npu_minimax_m2_5_w8a8_4p_in64k_out1k_prefix90_50ms.py b/test/registered/npu/performance/minimax_m2_5/test_npu_minimax_m2_5_w8a8_4p_in64k_out1k_prefix90_50ms.py index d37cf0f55..7ef54643f 100644 --- a/test/registered/npu/performance/minimax_m2_5/test_npu_minimax_m2_5_w8a8_4p_in64k_out1k_prefix90_50ms.py +++ b/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', '')}", } diff --git a/test/registered/npu/performance/qwen3_6_27b/test_npu_qwen3_6_27b_1p_in1080p_30_out256_50ms.py b/test/registered/npu/performance/qwen3_6_27b/test_npu_qwen3_6_27b_1p_in1080p_30_out256_50ms.py index abe2bedb6..01039457a 100644 --- a/test/registered/npu/performance/qwen3_6_27b/test_npu_qwen3_6_27b_1p_in1080p_30_out256_50ms.py +++ b/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", } diff --git a/test/registered/unit/managers/test_prefill_delayer_high_watermark.py b/test/registered/unit/managers/test_prefill_delayer_high_watermark.py new file mode 100644 index 000000000..e97a8a742 --- /dev/null +++ b/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()