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

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

Browse Source
This commit is contained in:
Liangsheng Yin
2026-06-23 17:04:55 -07:00
committed by GitHub
parent f444b5897b
commit 9ef1830701
6 changed files with 159 additions and 55 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/environ.py
+5 -1
View File
@@ -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
+41
View File
@@ -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
+25 -23
View File
@@ -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.
python/sglang/srt/server_args.py
+16
View File
@@ -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
-31
View File
@@ -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,