minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
hanwlax
2026-08-15 10:48:03 +08:00
committed by GitHub
parent 6eb941a34c
commit 8720a72814
7 changed files with 222 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/npu/performance/kimi_k2_6/test_npu_kimi_k2_6_w4a8_8p_in3k5_out1k5_20ms.py
+1
View File
@@ -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
+1
View File
@@ -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
+1
View File
@@ -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
+143
View File
@@ -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()