diff --git a/docs/advanced_features/server_arguments.md b/docs/advanced_features/server_arguments.md index 8ad1c0881..6e8074ab9 100644 --- a/docs/advanced_features/server_arguments.md +++ b/docs/advanced_features/server_arguments.md @@ -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] | diff --git a/python/sglang/srt/managers/prefill_delayer.py b/python/sglang/srt/managers/prefill_delayer.py index bc83f366a..a20ffd263 100644 --- a/python/sglang/srt/managers/prefill_delayer.py +++ b/python/sglang/srt/managers/prefill_delayer.py @@ -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 diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index 764d086d1..84ec14c36 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -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 diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 5c04dc4ee..37d7d05f3 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -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: diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 5da9e4337..4a1fe0540 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -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( diff --git a/test/registered/scheduler/test_prefill_delayer.py b/test/registered/scheduler/test_prefill_delayer.py index f1a9c8420..0bd8b7997 100644 --- a/test/registered/scheduler/test_prefill_delayer.py +++ b/test/registered/scheduler/test_prefill_delayer.py @@ -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", + ), ]