[Scheduler] Extract DFlash prefill refill into a standalone MinFreeSlotsDelayer (#29089)

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,

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

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 (
+from sglang.srt.speculative.dflash_utils import validate_dflash_request
-    resolve_dflash_prefill_refill_target,
-    should_delay_dflash_prefill_for_batching,
-    validate_dflash_request,
-)
 from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
 from sglang.srt.utils import (
     DynamicGradMode,
@@ -894,10 +894,17 @@ class Scheduler(
             _,
             _,
         ) = self.tp_worker.get_worker_info()
-        self.dflash_prefill_refill_target = (
+        # DFlash auto-enables the legacy formula; other workloads opt in via
-            resolve_dflash_prefill_refill_target(self.max_running_requests)
+        # --min-free-slots-delay. Built independently of the prefill delayer.
-            if self.spec_algorithm.is_dflash()
+        self.min_free_slots_delayer: Optional[MinFreeSlotsDelayer] = None
-            else 1
+        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(
@@ -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.

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
     # -------------------------------------------------------------------------

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,

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()