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

feat(scheduler): add adaptive queue-based prefill delayer trigger (#23189)

Browse Source
This commit is contained in:
YAMY
2026-05-08 16:54:30 -07:00
committed by GitHub
parent 6971a03fe6
commit 560829a171
6 changed files with 246 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/advanced_features/server_arguments.md
+2
View File
@@ -145,6 +145,8 @@ Please consult the documentation below and [server_args.py](https://github.com/s
| `--enable-prefill-delayer` | Enable prefill delayer for DP attention to reduce idle time. | `False` | bool flag (set to enable) |
| `--prefill-delayer-max-delay-passes` | Maximum forward passes to delay prefill. | `30` | Type: int |
| `--prefill-delayer-token-usage-low-watermark` | Token usage low watermark for prefill delayer. | `None` | Type: float |
| `--prefill-delayer-queue-min-ratio` | Opt-in to the adaptive queue-based delay trigger (independent of the slot-based one). Defers prefill until the waiting queue reaches `min(running_req * ratio, max_prefill_bs)` so small fragments batch into a larger prefill. Unset keeps the original slot-only behavior. Typical: `0.1`–`0.5`. | `None` | Type: float |
| `--prefill-delayer-max-delay-ms` | Wall-clock cap (ms) on a single queue-trigger delay; once exceeded, prefill is force-released to bound worst-case TTFT. Only consulted when `--prefill-delayer-queue-min-ratio` is set. Typical: `1000`–`5000`. | `5000` | Type: float |
| `--prefill-delayer-forward-passes-buckets` | Custom buckets for prefill delayer forward passes histogram. 0 and max_delay_passes-1 will be auto-added. | `None` | List[float] |
| `--prefill-delayer-wait-seconds-buckets` | Custom buckets for prefill delayer wait seconds histogram. 0 will be auto-added. | `None` | List[float] |
python/sglang/srt/managers/prefill_delayer.py
+67 -8
View File
@@ -48,18 +48,31 @@ class PrefillDelayer:
):
self._max_delay_passes = max_delay_passes
self._token_usage_low_watermark = token_usage_low_watermark
# Queue-based trigger is opt-in: activates only when queue_min_ratio
# is explicitly set. Additive with the slot-based trigger.
self._queue_min_ratio = server_args.prefill_delayer_queue_min_ratio
# Fall back to 5000ms if unset; this is a local safety cap, not a
# semantic default, so we don't surface it via ServerArgs.
self._max_delay_ms = server_args.prefill_delayer_max_delay_ms
if self._max_delay_ms is None:
self._max_delay_ms = 5000.0
self._queue_trigger_enabled = self._queue_min_ratio is not None
logger.info(
f"PrefillDelayer initialized with "
f"max_delay_passes={self._max_delay_passes} "
f"token_usage_low_watermark={self._token_usage_low_watermark}"
f"token_usage_low_watermark={self._token_usage_low_watermark} "
f"queue_min_ratio={self._queue_min_ratio} "
f"max_delay_ms={self._max_delay_ms} "
f"queue_trigger_enabled={self._queue_trigger_enabled}"
)
# The global_info contains four pieces of information:
# prefillable, token_watermark_force_allow, running_batch, and max_prefill_bs.
self.dp_size = dp_size
self.enable_dp_attention = server_args.enable_dp_attention
dp_size_dim = dp_size if self.enable_dp_attention else 1
# Fields packed per rank into the all-gather tensor: prefillable,
# token_watermark_force_allow, running_batch, max_prefill_bs,
# waiting_queue_len.
self._global_info_buffer = torch.empty(
(dp_size_dim, attn_tp_size, 4),
(dp_size_dim, attn_tp_size, 5),
dtype=torch.int64,
device=device,
)
@@ -81,6 +94,7 @@ class PrefillDelayer:
running_batch: int = 0,
max_prefill_bs: int = 0,
max_running_requests: int = 0,
waiting_queue_len: int = 0,
) -> _NegotiateOutput:
out = self._negotiate_should_allow_prefill_pure(
prev_state=self._curr_state,
@@ -89,6 +103,7 @@ class PrefillDelayer:
running_batch=running_batch,
max_prefill_bs=max_prefill_bs,
max_running_requests=max_running_requests,
waiting_queue_len=waiting_queue_len,
)
self._curr_state = out.next_state
return out
@@ -102,6 +117,7 @@ class PrefillDelayer:
running_batch: int = 0,
max_prefill_bs: int = 0,
max_running_requests: int = 0,
waiting_queue_len: int = 0,
) -> _NegotiateOutput:
# Compute local states
local_token_watermark_force_allow = (
@@ -116,11 +132,13 @@ class PrefillDelayer:
local_token_watermark_force_allow=local_token_watermark_force_allow,
running_batch=running_batch,
max_prefill_bs=max_prefill_bs,
waiting_queue_len=waiting_queue_len,
)
global_prefillable = tp0_info[:, 0]
global_token_watermark_force_allow = tp0_info[:, 1]
global_running_batch = tp0_info[:, 2]
global_max_prefill_bs = tp0_info[:, 3]
global_waiting_queue_len = tp0_info[:, 4]
# Compute derived global states
if global_prefillable.min().item() > 0:
@@ -140,14 +158,51 @@ class PrefillDelayer:
# Compute outputs
if prefillable_status == "all":
# Safety valve: low KV usage means GPU is underutilized, skip
# delay. Mirrors the check in the "mixed" branch.
if global_exists_token_watermark_force_allow:
return _NegotiateOutput(
next_state=None,
output_allow=True,
output_reason="token_watermark",
**debug_info,
)
if not self.enable_dp_attention:
max_running_requests = (
max_running_requests + self.dp_size - 1
) // self.dp_size
if (
max_running_requests - global_running_batch.max().item()
< global_max_prefill_bs.max().item()
):
global_running_batch_max = int(global_running_batch.max().item())
global_max_prefill_bs_max = int(global_max_prefill_bs.max().item())
global_waiting_queue_max = int(global_waiting_queue_len.max().item())
# Queue-based trigger: delay prefill until the waiting queue
# reaches queue_min = min(running_req * ratio, max_prefill_bs),
# capped by a wall-clock timeout to bound worst-case TTFT.
# Targets workloads where decode requests finish one-at-a-time
# and fragment prefill into many tiny batches.
queue_condition = False
if self._queue_trigger_enabled and global_running_batch_max > 0:
queue_min_effective = min(
int(global_running_batch_max * self._queue_min_ratio),
global_max_prefill_bs_max,
)
queue_condition = (
queue_min_effective > 0
and global_waiting_queue_max < queue_min_effective
)
if queue_condition and prev_state is not None:
elapsed_ms = (time.perf_counter() - prev_state.start_time) * 1000.0
if elapsed_ms >= self._max_delay_ms:
queue_condition = False
slot_condition = (
max_running_requests - global_running_batch_max
< global_max_prefill_bs_max
)
if slot_condition or queue_condition:
# When the "max_decode_bs - running_bs < max_prefill_bs" condition is met,
# the first merge_batch causes the decoding to fail to reach the maximum batch size.
if self.skip_first_delayer:
@@ -212,6 +267,7 @@ class PrefillDelayer:
local_token_watermark_force_allow: bool,
running_batch: int = 0,
max_prefill_bs: int = 0,
waiting_queue_len: int = 0,
):
local_info = torch.tensor(
[
@@ -219,6 +275,7 @@ class PrefillDelayer:
int(local_token_watermark_force_allow),
running_batch,
max_prefill_bs,
waiting_queue_len,
],
device="cpu",
dtype=torch.int64,
@@ -258,6 +315,7 @@ class PrefillDelayerSinglePassExecutor:
running_batch: int = 0,
max_prefill_bs: int = 0,
max_running_requests: int = 0,
waiting_queue_len: int = 0,
) -> bool:
if not self._called:
self._result = self._prefill_delayer._negotiate_should_allow_prefill(
@@ -266,6 +324,7 @@ class PrefillDelayerSinglePassExecutor:
running_batch=running_batch,
max_prefill_bs=max_prefill_bs,
max_running_requests=max_running_requests,
waiting_queue_len=waiting_queue_len,
)
return self._result.output_allow
python/sglang/srt/managers/schedule_policy.py
+5
View File
@@ -412,6 +412,7 @@ class PrefillAdder:
prefill_max_requests: Optional[int] = None,
prefill_delayer_single_pass: Optional[PrefillDelayerSinglePassExecutor] = None,
dllm_config: Optional[DllmConfig] = None,
waiting_queue_len: int = 0,
):
self.page_size = page_size
self.tree_cache = tree_cache
@@ -466,6 +467,9 @@ class PrefillAdder:
self.prefill_max_requests = prefill_max_requests
self.prefill_delayer_single_pass = prefill_delayer_single_pass
self.max_prefill_bs = max_prefill_bs
# Snapshot of scheduler waiting_queue length at the start of this
# prefill pass. Used by PrefillDelayer's queue-based trigger.
self.waiting_queue_len = waiting_queue_len
def _init_dllm_meta(self, dllm_config: DllmConfig):
self.dllm_block_size = dllm_config.block_size
@@ -806,6 +810,7 @@ class PrefillAdder:
running_batch=self.running_batch.batch_size(),
max_prefill_bs=self.max_prefill_bs,
max_running_requests=self.max_running_requests,
waiting_queue_len=self.waiting_queue_len,
)
):
return AddReqResult.OTHER
python/sglang/srt/managers/scheduler.py
+1
View File
@@ -2647,6 +2647,7 @@ class Scheduler(
prefill_max_requests=self.server_args.prefill_max_requests,
prefill_delayer_single_pass=prefill_delayer_single_pass,
dllm_config=self.dllm_config,
waiting_queue_len=len(self.waiting_queue),
)
if self.chunked_req is not None:
python/sglang/srt/server_args.py
+25
View File
@@ -425,6 +425,8 @@ class ServerArgs:
prefill_delayer_token_usage_low_watermark: Optional[float] = None
prefill_delayer_forward_passes_buckets: Optional[List[float]] = None
prefill_delayer_wait_seconds_buckets: Optional[List[float]] = None
prefill_delayer_queue_min_ratio: Optional[float] = None
prefill_delayer_max_delay_ms: Optional[float] = None
# Runtime options
device: Optional[str] = None
@@ -4733,6 +4735,29 @@ class ServerArgs:
default=None,
help="Custom buckets for prefill delayer wait seconds histogram. 0 will be auto-added.",
)
parser.add_argument(
"--prefill-delayer-queue-min-ratio",
type=float,
default=None,
help=(
"Opt-in to the adaptive queue-based delay trigger (independent of the "
"slot-based one). Delays prefill until the waiting queue reaches "
"min(running_req * ratio, max_prefill_bs) so small fragments batch into a "
"larger prefill. Unset (default) keeps the original slot-only behavior. "
"Typical: 0.1 ~ 0.5."
),
)
parser.add_argument(
"--prefill-delayer-max-delay-ms",
type=float,
default=None,
help=(
"Wall-clock cap (ms) on a single queue-trigger delay; once exceeded, prefill "
"is force-released to bound worst-case TTFT. Only consulted when "
"--prefill-delayer-queue-min-ratio is set. Typical: 1000 ~ 5000; defaults to "
"5000 if unset."
),
)
# Runtime options
parser.add_argument(
test/registered/scheduler/test_prefill_delayer.py
+146
View File
@@ -41,6 +41,16 @@ WORLD_SIZE = os.environ.get("SGLANG_TEST_WORLD_SIZE", "8")
class NegotiateCall:
prefillable: List[bool]
token_usage: List[float]
# Optional scheduler state; when None, _run_negotiate_test does not pass
# the kwarg and the delayer falls back to the historical behavior of
# reading kwargs.get(..., 0).
running_batch: Optional[List[int]] = None
max_prefill_bs: Optional[List[int]] = None
waiting_queue_len: Optional[List[int]] = None
max_running_requests: Optional[int] = None
# Inter-call sleep (seconds). Used to exercise the queue-trigger
# wall-clock timeout.
sleep_before_s: float = 0.0
@dataclass
@@ -51,6 +61,10 @@ class NegotiateTestCase:
calls: List[NegotiateCall]
expected_allow: bool
expected_reason: str
# Queue-trigger knobs (new in the queue-based delayer). Leave both None
# to exercise the legacy slot-only code paths.
queue_min_ratio: Optional[float] = None
max_delay_ms: Optional[float] = None
def _run_negotiate_test(rank, test_cases):
@@ -66,15 +80,31 @@ def _run_negotiate_test(rank, test_cases):
enable_dp_attention=True,
disaggregation_mode="null",
disable_overlap_schedule=False,
prefill_delayer_queue_min_ratio=case.queue_min_ratio,
prefill_delayer_max_delay_ms=case.max_delay_ms,
),
max_delay_passes=case.max_delay_passes,
token_usage_low_watermark=case.token_usage_low_watermark,
)
for call in case.calls:
if call.sleep_before_s > 0:
time.sleep(call.sleep_before_s)
extra_kwargs = {}
if call.running_batch is not None:
extra_kwargs["running_batch"] = call.running_batch[rank]
if call.max_prefill_bs is not None:
extra_kwargs["max_prefill_bs"] = call.max_prefill_bs[rank]
if call.waiting_queue_len is not None:
extra_kwargs["waiting_queue_len"] = call.waiting_queue_len[rank]
if call.max_running_requests is not None:
extra_kwargs["max_running_requests"] = call.max_running_requests
result = delayer._negotiate_should_allow_prefill(
local_prefillable=call.prefillable[rank],
token_usage=call.token_usage[rank],
**extra_kwargs,
)
assert (result.output_allow, result.output_reason) == (
@@ -200,6 +230,122 @@ _NEGOTIATE_TEST_CASES = [
expected_allow=True,
expected_reason="wait_timeout",
),
# Queue-based trigger: waiting queue below queue_min = min(running * R,
# max_prefill_bs) should defer prefill. With R=0.5, running=100 and
# max_prefill_bs=80, queue_min = min(50, 80) = 50, and queue_len=10 < 50.
NegotiateTestCase(
name="queue_trigger_delay",
max_delay_passes=100,
token_usage_low_watermark=0.8,
queue_min_ratio=0.5,
max_delay_ms=5000,
calls=[
NegotiateCall(
prefillable=[True, True, True, True],
token_usage=[0.9, 0.9, 0.9, 0.9],
running_batch=[100, 100, 100, 100],
max_prefill_bs=[80, 80, 80, 80],
waiting_queue_len=[10, 10, 10, 10],
max_running_requests=1024,
),
# skip_first_delayer consumes the first would-be delay; a second
# identical call must actually delay.
NegotiateCall(
prefillable=[True, True, True, True],
token_usage=[0.9, 0.9, 0.9, 0.9],
running_batch=[100, 100, 100, 100],
max_prefill_bs=[80, 80, 80, 80],
waiting_queue_len=[10, 10, 10, 10],
max_running_requests=1024,
),
],
expected_allow=False,
expected_reason="delay",
),
# Waiting queue at or above queue_min: queue trigger must not fire.
NegotiateTestCase(
name="queue_trigger_above_threshold",
max_delay_passes=100,
token_usage_low_watermark=0.8,
queue_min_ratio=0.5,
max_delay_ms=5000,
calls=[
NegotiateCall(
prefillable=[True, True, True, True],
token_usage=[0.9, 0.9, 0.9, 0.9],
running_batch=[100, 100, 100, 100],
max_prefill_bs=[80, 80, 80, 80],
waiting_queue_len=[64, 64, 64, 64],
max_running_requests=1024,
)
],
expected_allow=True,
expected_reason="no_wait",
),
# queue_min_ratio unset: queue trigger is opt-in and must stay disabled
# even when running_batch and queue_len would otherwise trigger it.
NegotiateTestCase(
name="queue_trigger_disabled_when_ratio_unset",
max_delay_passes=100,
token_usage_low_watermark=0.8,
queue_min_ratio=None,
max_delay_ms=None,
calls=[
NegotiateCall(
prefillable=[True, True, True, True],
token_usage=[0.9, 0.9, 0.9, 0.9],
running_batch=[100, 100, 100, 100],
max_prefill_bs=[80, 80, 80, 80],
waiting_queue_len=[1, 1, 1, 1],
max_running_requests=1024,
)
],
expected_allow=True,
expected_reason="no_wait",
),
# max_delay_ms wall-clock timeout: once a single queue-trigger delay
# exceeds the cap, prefill must be force-released.
# Call sequence:
# 1) queue_condition holds but skip_first_delayer consumes it
# (no state recorded, falls through to allow)
# 2) queue_condition holds -> delay, records start_time in state
# 3) after sleeping past max_delay_ms, elapsed >= cap -> force release
NegotiateTestCase(
name="queue_trigger_wall_clock_timeout",
max_delay_passes=100,
token_usage_low_watermark=0.8,
queue_min_ratio=0.5,
max_delay_ms=50,
calls=[
NegotiateCall(
prefillable=[True, True, True, True],
token_usage=[0.9, 0.9, 0.9, 0.9],
running_batch=[100, 100, 100, 100],
max_prefill_bs=[80, 80, 80, 80],
waiting_queue_len=[10, 10, 10, 10],
max_running_requests=1024,
),
NegotiateCall(
prefillable=[True, True, True, True],
token_usage=[0.9, 0.9, 0.9, 0.9],
running_batch=[100, 100, 100, 100],
max_prefill_bs=[80, 80, 80, 80],
waiting_queue_len=[10, 10, 10, 10],
max_running_requests=1024,
),
NegotiateCall(
prefillable=[True, True, True, True],
token_usage=[0.9, 0.9, 0.9, 0.9],
running_batch=[100, 100, 100, 100],
max_prefill_bs=[80, 80, 80, 80],
waiting_queue_len=[10, 10, 10, 10],
max_running_requests=1024,
sleep_before_s=0.2, # > max_delay_ms (50ms)
),
],
expected_allow=True,
expected_reason="wait_success",
),
]