chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Refactor and fix prefill delayer (scheduler enhancer) (#16269)

Browse Source
This commit is contained in:
fzyzcjy
2026-01-04 07:15:06 +08:00
committed by GitHub
parent 7bc13c906c
commit 2337b1bbb0
6 changed files with 172 additions and 79 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
python/sglang/srt/environ.py
View File

@@ -223,6 +223,7 @@ class Envs:
SGLANG_DYNAMIC_CHUNKING_SMOOTH_FACTOR = EnvFloat(0.75)
SGLANG_SCHEDULER_SKIP_ALL_GATHER = EnvBool(False)
SGLANG_SCHEDULER_DECREASE_PREFILL_IDLE = EnvBool(False)
SGLANG_PREFILL_DELAYER_MAX_DELAY_PASSES = EnvInt(30)
SGLANG_DATA_PARALLEL_BUDGET_INTERVAL = EnvInt(1)
# Test: pd-disaggregation

72
python/sglang/srt/managers/prefill_delayer.py Normal file
View File

@@ -0,0 +1,72 @@
import logging
import torch
from sglang.srt.environ import envs
from sglang.srt.utils import get_bool_env_var
_DEBUG_LOG = get_bool_env_var("SGLANG_PREFILL_DELAYER_DEBUG_LOG")
logger = logging.getLogger(__name__)
class PrefillDelayer:
def __init__(self, dp_size, attn_tp_size, tp_worker, server_args):
self.global_info = torch.empty(
(dp_size, attn_tp_size, 1),
dtype=torch.int64,
device="cpu",
)
self.cpu_group = tp_worker.get_tp_group().cpu_group
self.curr_delayed_count = 0
self.max_delay_passes = envs.SGLANG_PREFILL_DELAYER_MAX_DELAY_PASSES.get()
assert (
server_args.schedule_policy == "fcfs"
), f"To use PrefillDelayer, schedule_policy must be 'fcfs'. '{server_args.schedule_policy}' is not supported."
assert (
server_args.enable_dp_attention
), "To use PrefillDelayer, enable_dp_attention must be enabled."
assert (
server_args.disaggregation_mode == "null"
), "To use PrefillDelayer, disaggregation_mode must be null."
assert (
not server_args.disable_overlap_schedule
), "To use PrefillDelayer, disable_overlap_schedule must be False."
def _gather_info(self, local_prefillable: int):
local_info = torch.tensor(
[local_prefillable],
device="cpu",
dtype=torch.int64,
)
torch.distributed.all_gather_into_tensor(
self.global_info.flatten(),
local_info,
group=self.cpu_group,
)
tp0_info = self.global_info[:, 0, :]
return tp0_info
def should_allow_prefill(self, local_prefillable: int) -> bool:
tp0_info = self._gather_info(local_prefillable=local_prefillable)
global_prefillable = tp0_info[:, 0]
global_exists_not_prefillable = global_prefillable.min().item() == 0
global_exists_prefillable = global_prefillable.max().item() > 0
global_mixed_prefillable = (
global_exists_not_prefillable and global_exists_prefillable
)
if global_mixed_prefillable:
self.curr_delayed_count += 1
if self.curr_delayed_count < self.max_delay_passes:
return False
if _DEBUG_LOG and global_mixed_prefillable:
logger.info(
f"PrefillDelayer timeout thus not forbid prefill (prefillable: {global_prefillable.sum()})"
)
self.curr_delayed_count = 0
return True

42
python/sglang/srt/managers/scheduler.py
View File

@@ -120,6 +120,7 @@ from sglang.srt.managers.io_struct import (
)
from sglang.srt.managers.mm_utils import init_mm_embedding_cache
from sglang.srt.managers.overlap_utils import FutureMap
from sglang.srt.managers.prefill_delayer import PrefillDelayer
from sglang.srt.managers.schedule_batch import (
FINISH_ABORT,
ModelWorkerBatch,
@@ -134,7 +135,6 @@ from sglang.srt.managers.schedule_policy import (
SchedulePolicy,
)
from sglang.srt.managers.scheduler_dp_attn_mixin import SchedulerDPAttnMixin
from sglang.srt.managers.scheduler_enhancer import SchedulerEnhancer
from sglang.srt.managers.scheduler_input_blocker import SchedulerInputBlocker
from sglang.srt.managers.scheduler_metrics_mixin import (
RECORD_STEP_TIME,
@@ -204,7 +204,6 @@ logger = logging.getLogger(__name__)
TEST_RETRACT = envs.SGLANG_TEST_RETRACT.get()
TEST_RETRACT_INTERVAL = envs.SGLANG_TEST_RETRACT_INTERVAL.get()
TEST_RETRACT_NO_PREFILL_BS = envs.SGLANG_TEST_RETRACT_NO_PREFILL_BS.get()
SCHEDULER_DECREASE_PREFILL_IDLE = envs.SGLANG_SCHEDULER_DECREASE_PREFILL_IDLE.get()
GRAMMAR_TIMEOUT = float(os.environ.get("SGLANG_GRAMMAR_TIMEOUT", 300))
@@ -796,14 +795,13 @@ class Scheduler(
self.enable_priority_scheduling,
self.schedule_low_priority_values_first,
)
self.schedule_enhancer = None
if SCHEDULER_DECREASE_PREFILL_IDLE:
self.schedule_enhancer = SchedulerEnhancer(
self.dp_size,
self.attn_tp_size,
self.tp_worker,
self.max_running_requests,
self.server_args,
self.prefill_delayer: Optional[PrefillDelayer] = None
if envs.SGLANG_SCHEDULER_DECREASE_PREFILL_IDLE.get():
self.prefill_delayer = PrefillDelayer(
dp_size=self.dp_size,
attn_tp_size=self.attn_tp_size,
tp_worker=self.tp_worker,
server_args=self.server_args,
)
# Enable preemption for priority scheduling.
self.try_preemption = self.enable_priority_scheduling
@@ -1913,12 +1911,6 @@ class Scheduler(
return res
def get_new_batch_prefill(self) -> Optional[ScheduleBatch]:
if self.schedule_enhancer and not self.schedule_enhancer.get_schedule_decision(
self.running_batch
):
# Decrease prefill idle as much as possible during high dp load.
return None
# Check if the grammar is ready in the grammar queue
if self.grammar_queue:
self.move_ready_grammar_requests()
@@ -1927,10 +1919,20 @@ class Scheduler(
# Reset batch_is_full to try preemption with a prefill adder.
self.running_batch.batch_is_full = False
# Handle the cases where prefill is not allowed
if (
self.running_batch.batch_is_full or len(self.waiting_queue) == 0
) and self.chunked_req is None:
# The `should_allow_prefill` needs to be called on all ranks since contains communication
delayer_allow_prefill = (
self.prefill_delayer.should_allow_prefill(
# TODO: consider offline generation cases when there are a lot of waiting requests
local_prefillable=len(self.waiting_queue)
> 0,
)
if self.prefill_delayer
else True
)
if (not delayer_allow_prefill) or (
(self.running_batch.batch_is_full or len(self.waiting_queue) == 0)
and self.chunked_req is None
):
return None
running_bs = len(self.running_batch.reqs)

59
python/sglang/srt/managers/scheduler_enhancer.py
View File

@@ -1,59 +0,0 @@
import torch
class SchedulerEnhancer:
def __init__(
self, dp_size, attn_tp_size, tp_worker, max_running_requests, server_args
):
self.dp_size = dp_size
self.attn_tp_size = attn_tp_size
self.global_batch_size = torch.empty(
(self.dp_size, self.attn_tp_size, 1),
dtype=torch.int64,
device="cpu",
)
self.cpu_group = tp_worker.get_tp_group().cpu_group
self.max_running_requests = max_running_requests
self.stable_count = 0
# If scheduling is performed 30 times and some dp units are still at full load, the prefill-prioritized scheduling strategy will still be used.
self.max_stable_count = 30
assert (
server_args.schedule_policy == "fcfs"
), f"To use SCHEDULER_DECREASE_PREFILL_IDLE, schedule_policy must be 'fcfs'. '{self.schedule_policy}' is not supported."
assert (
server_args.enable_dp_attention == True
), f"To use SCHEDULER_DECREASE_PREFILL_IDLE, enable_dp_attention must be enable."
assert (
server_args.disaggregation_mode == "null"
), f"To use SCHEDULER_DECREASE_PREFILL_IDLE, disaggregation_mode must be null."
assert (
server_args.disable_overlap_schedule == False
), f"To use SCHEDULER_DECREASE_PREFILL_IDLE, disable_overlap_schedule must be False."
def get_schedule_info(self, running_batch):
local_batch_size = torch.tensor(
[
running_batch.batch_size(),
],
device="cpu",
dtype=torch.int64,
)
torch.distributed.all_gather_into_tensor(
self.global_batch_size.flatten(),
local_batch_size,
group=self.cpu_group,
)
tp0_info = self.global_batch_size[:, 0, :]
return tp0_info
def get_schedule_decision(self, running_batch):
tp0_info = self.get_schedule_info(running_batch)
if (
int(tp0_info[:, 0].min().item()) < self.max_running_requests
and int(tp0_info[:, 0].max().item()) == self.max_running_requests
):
self.stable_count += 1
if self.stable_count < self.max_stable_count:
return False
self.stable_count = 0
return True