diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index cc5851a2a..9cc2f7710 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -636,7 +636,6 @@ class Envs: # Overlap Spec V2 SGLANG_ENABLE_OVERLAP_PLAN_STREAM = EnvBool(False) - SGLANG_DFLASH_PREFILL_REFILL_TARGET = EnvInt(None) # Spec Config SGLANG_SPEC_ENABLE_STRICT_FILTER_CHECK = EnvBool(True) @@ -994,6 +993,11 @@ _warn_deprecated_env_to_cli_flag( "SGLANG_PREFILL_DELAYER_TOKEN_USAGE_LOW_WATERMARK", "Please use '--prefill-delayer-token-usage-low-watermark' instead.", ) +_warn_deprecated_env_to_cli_flag( + "SGLANG_DFLASH_PREFILL_REFILL_TARGET", + "DFlash now auto-enables the min-free-slots delay; unset this env. To " + "override the threshold, use '--min-free-slots-delay'.", +) # Import cuda_coredump to trigger auto-injection of CUDA env vars # when SGLANG_CUDA_COREDUMP=1. Best-effort; for strict guarantees, diff --git a/python/sglang/srt/managers/min_free_slots_delayer.py b/python/sglang/srt/managers/min_free_slots_delayer.py new file mode 100644 index 000000000..9f65d9ec6 --- /dev/null +++ b/python/sglang/srt/managers/min_free_slots_delayer.py @@ -0,0 +1,41 @@ +from typing import Optional + + +def resolve_min_free_slots( + user_value: Optional[int], + max_running_requests: int, + is_dflash: bool = False, +) -> Optional[int]: + """Resolve the min-free-slots threshold (None = disabled). + + A user value (>1) is capped to the DFlash formula so the trigger never + delays more aggressively than the legacy heuristic. When unset, DFlash + workloads fall back to the formula (preserving the always-on behavior); + other workloads stay disabled. Also disabled when max_running_requests < 8. + """ + max_running_requests = max(0, int(max_running_requests)) + formula = min(4, max(2, (max_running_requests + 5) // 6)) + if user_value is None: + user_value = formula if is_dflash else None + + if user_value is None or user_value <= 1: + return None + if max_running_requests < 8: + return None + return min(user_value, formula) + + +class MinFreeSlotsDelayer: + """Delay fresh prefill admissions until at least ``min_free_slots`` running- + request slots free up, batching them into one admission instead of one at a + time. Useful when each admission is expensive (e.g. DFlash's draft prefill). + + Per-rank local: running-batch slots are private to each DP rank, so a rank + with free slots does not wait for a congested peer. + """ + + def __init__(self, min_free_slots: int): + self._min_free_slots = min_free_slots + + def should_delay(self, *, running_bs: int, num_allocatable_reqs: int) -> bool: + return running_bs > 0 and num_allocatable_reqs < self._min_free_slots diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index a2765c954..d36cea7c0 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -150,6 +150,10 @@ from sglang.srt.managers.io_struct import ( sock_send, ) from sglang.srt.managers.load_snapshot import LoadSnapshot, create_load_snapshot_writer +from sglang.srt.managers.min_free_slots_delayer import ( + MinFreeSlotsDelayer, + resolve_min_free_slots, +) from sglang.srt.managers.multimodal_processor import get_mm_processor, import_processors from sglang.srt.managers.overlap_utils import ( decide_needs_cpu_seq_lens, @@ -237,11 +241,7 @@ from sglang.srt.plugins import load_plugins from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo from sglang.srt.server_args import PortArgs, ServerArgs, get_global_server_args from sglang.srt.session.session_controller import SessionController -from sglang.srt.speculative.dflash_utils import ( - resolve_dflash_prefill_refill_target, - should_delay_dflash_prefill_for_batching, - validate_dflash_request, -) +from sglang.srt.speculative.dflash_utils import validate_dflash_request from sglang.srt.speculative.spec_info import SpeculativeAlgorithm from sglang.srt.utils import ( DynamicGradMode, @@ -894,11 +894,18 @@ class Scheduler( _, _, ) = self.tp_worker.get_worker_info() - self.dflash_prefill_refill_target = ( - resolve_dflash_prefill_refill_target(self.max_running_requests) - if self.spec_algorithm.is_dflash() - else 1 + # DFlash auto-enables the legacy formula; other workloads opt in via + # --min-free-slots-delay. Built independently of the prefill delayer. + self.min_free_slots_delayer: Optional[MinFreeSlotsDelayer] = None + min_free_slots = resolve_min_free_slots( + self.server_args.min_free_slots_delay, + self.max_running_requests, + is_dflash=self.spec_algorithm.is_dflash(), ) + if min_free_slots is not None: + self.min_free_slots_delayer = MinFreeSlotsDelayer( + min_free_slots=min_free_slots + ) if not get_global_server_args().pp_max_micro_batch_size: get_global_server_args().pp_max_micro_batch_size = max( self.max_running_requests // self.ps.pp_size, 1 @@ -2720,19 +2727,6 @@ class Scheduler( res = min(res, self.req_to_token_pool.available_size()) return res - def _should_delay_dflash_prefill_for_batching(self, running_bs: int) -> bool: - if not self.spec_algorithm.is_dflash(): - return False - if running_bs <= 0 or self.chunked_req is not None: - return False - - return should_delay_dflash_prefill_for_batching( - running_bs=running_bs, - num_allocatable_reqs=self.get_num_allocatable_reqs(running_bs), - max_running_requests=self.max_running_requests, - prefill_refill_target=self.dflash_prefill_refill_target, - ) - def get_new_batch_prefill(self) -> Optional[ScheduleBatch]: prefill_delayer_single_pass = None if self.prefill_delayer: @@ -2775,7 +2769,15 @@ class Scheduler( return None running_bs = len(self.running_batch.reqs) - if self._should_delay_dflash_prefill_for_batching(running_bs): + # Skipped during a chunked prefill: that pass must proceed regardless. + if ( + self.min_free_slots_delayer is not None + and self.chunked_req is None + and self.min_free_slots_delayer.should_delay( + running_bs=running_bs, + num_allocatable_reqs=self.get_num_allocatable_reqs(running_bs), + ) + ): return None # Ignore the check if self.chunked_req is not None. diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 82f8a3d74..eaa41f3a9 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -1261,6 +1261,22 @@ class ServerArgs: ), ] = None + # ------------------------------------------------------------------------- + # Min free slots delay (prefill refill batching) + # ------------------------------------------------------------------------- + min_free_slots_delay: A[ + Optional[int], + ( + "Hold new prefills until at least N running-request slots have freed " + "up, so they are admitted in one batch instead of one at a time. " + "Useful when each admission is disproportionately expensive, e.g. " + "speculative decoding with a separate draft prefill pass. Capped to " + "the DFlash formula (disabled when max-running-requests < 8; " + "min(4, max(2, (max-run + 5) // 6))). DFlash workloads auto-enable " + "this with the formula when unset; other workloads stay disabled." + ), + ] = None + # ------------------------------------------------------------------------- # LoRA # ------------------------------------------------------------------------- diff --git a/python/sglang/srt/speculative/dflash_utils.py b/python/sglang/srt/speculative/dflash_utils.py index e63efe4ea..1c3c798ea 100644 --- a/python/sglang/srt/speculative/dflash_utils.py +++ b/python/sglang/srt/speculative/dflash_utils.py @@ -9,7 +9,6 @@ from typing import Any, List, Optional, Tuple import torch import torch.nn.functional as F -from sglang.srt.environ import envs from sglang.srt.layers.quantization.unquant import UnquantizedLinearMethod from sglang.srt.layers.sampler import apply_custom_logit_processor from sglang.srt.managers.schedule_batch import Req @@ -56,36 +55,6 @@ def is_dflash_sampling_verify_available() -> bool: return _DFLASH_SAMPLING_VERIFY_AVAILABLE -def resolve_dflash_prefill_refill_target(max_running_requests: int) -> int: - """Choose how many free running-request slots DFlash waits for before refill.""" - override = envs.SGLANG_DFLASH_PREFILL_REFILL_TARGET.get() - if override is not None: - return override - - max_running_requests = max(0, int(max_running_requests)) - if max_running_requests < 8: - return 1 - return min(4, max(2, (max_running_requests + 5) // 6)) - - -def should_delay_dflash_prefill_for_batching( - *, - running_bs: int, - num_allocatable_reqs: int, - max_running_requests: int, - prefill_refill_target: int, -) -> bool: - if running_bs <= 0: - return False - - target_prefill_bs = int(prefill_refill_target) - if target_prefill_bs <= 1: - return False - - target_prefill_bs = min(target_prefill_bs, int(max_running_requests)) - return int(num_allocatable_reqs) < target_prefill_bs - - def scale_kv_cell_size_per_token_for_dflash( *, target_cell_size_per_token: int, diff --git a/test/registered/scheduler/test_min_free_slots_delayer.py b/test/registered/scheduler/test_min_free_slots_delayer.py new file mode 100644 index 000000000..4dee7bf3b --- /dev/null +++ b/test/registered/scheduler/test_min_free_slots_delayer.py @@ -0,0 +1,72 @@ +import unittest + +from sglang.srt.managers.min_free_slots_delayer import ( + MinFreeSlotsDelayer, + resolve_min_free_slots, +) +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + + +class TestResolveMinFreeSlots(unittest.TestCase): + """Unit tests for resolve_min_free_slots threshold resolution.""" + + def test_unset_non_dflash_disables(self): + # Unset + not DFlash -> trigger stays disabled. + self.assertIsNone(resolve_min_free_slots(None, 512, is_dflash=False)) + + def test_unset_dflash_auto_enables(self): + # Unset + DFlash -> falls back to the legacy formula (full mapping). + self.assertEqual(resolve_min_free_slots(None, 512, is_dflash=True), 4) + self.assertEqual(resolve_min_free_slots(None, 8, is_dflash=True), 2) + + def test_unset_dflash_small_cluster_disables(self): + # DFlash auto-default still respects the < 8 guard. + self.assertIsNone(resolve_min_free_slots(None, 7, is_dflash=True)) + self.assertIsNone(resolve_min_free_slots(None, 0, is_dflash=True)) + + def test_le_one_disables(self): + # <= 1 can never batch, so it is a no-op. + self.assertIsNone(resolve_min_free_slots(1, 512)) + self.assertIsNone(resolve_min_free_slots(0, 512)) + + def test_small_cluster_disables(self): + # max_running_requests < 8 disables, matching DFlash. + self.assertIsNone(resolve_min_free_slots(4, 7)) + + def test_caps_to_formula(self): + # Capped down so it never delays more aggressively than DFlash. + self.assertEqual(resolve_min_free_slots(10, 512), 4) + self.assertEqual(resolve_min_free_slots(10, 8), 2) # (8 + 5) // 6 = 2 + + def test_respects_smaller_user_value(self): + # Below the formula cap is taken as-is. + self.assertEqual(resolve_min_free_slots(3, 512), 3) + self.assertEqual(resolve_min_free_slots(2, 8), 2) + + def test_user_value_overrides_dflash_default(self): + # An explicit user value wins over the DFlash auto-default. + self.assertEqual(resolve_min_free_slots(3, 512, is_dflash=True), 3) + + +class TestMinFreeSlotsDelayer(unittest.TestCase): + """Unit tests for the per-rank local should_delay decision.""" + + def test_delays_below_threshold(self): + delayer = MinFreeSlotsDelayer(min_free_slots=4) + self.assertTrue(delayer.should_delay(running_bs=100, num_allocatable_reqs=2)) + + def test_no_delay_at_or_above_threshold(self): + delayer = MinFreeSlotsDelayer(min_free_slots=4) + self.assertFalse(delayer.should_delay(running_bs=100, num_allocatable_reqs=4)) + self.assertFalse(delayer.should_delay(running_bs=100, num_allocatable_reqs=8)) + + def test_no_delay_when_idle(self): + # Nothing running: no decode batch to protect, prefill at once. + delayer = MinFreeSlotsDelayer(min_free_slots=4) + self.assertFalse(delayer.should_delay(running_bs=0, num_allocatable_reqs=0)) + + +if __name__ == "__main__": + unittest.main()