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

Added the prefill delayer policy: The prefill deplay range is expanded. (#17456)

Browse Source 
Co-authored-by: sglang-npu-bot <sglangnpu@163.com>
This commit is contained in:
chenxu214
2026-02-28 08:56:49 +08:00
committed by GitHub
parent c6cb0c9649
commit 5f07ff9271
3 changed files with 68 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files

63
python/sglang/srt/managers/prefill_delayer.py
View File

@@ -44,6 +44,7 @@ class PrefillDelayer:
max_delay_passes: int,
token_usage_low_watermark: Optional[float],
metrics_collector: Optional["SchedulerMetricsCollector"] = None,
device: Optional["torch.device"] = "cpu",
):
self._max_delay_passes = max_delay_passes
self._token_usage_low_watermark = token_usage_low_watermark
@@ -52,21 +53,20 @@ class PrefillDelayer:
f"max_delay_passes={self._max_delay_passes} "
f"token_usage_low_watermark={self._token_usage_low_watermark}"
)
# The global_info contains four pieces of information:
# prefillable, token_watermark_force_allow, running_batch, and max_prefill_bs.
self._global_info_buffer = torch.empty(
(dp_size, attn_tp_size, 2),
(dp_size, attn_tp_size, 4),
dtype=torch.int64,
device="cpu",
device=device,
)
self.enable_dp_attention = server_args.enable_dp_attention
self._cpu_group = cpu_group
self._metrics_collector = metrics_collector
self._curr_state: Optional[_State] = None
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."
@@ -75,12 +75,16 @@ class PrefillDelayer:
), "To use PrefillDelayer, disable_overlap_schedule must be False."
def _negotiate_should_allow_prefill(
self, local_prefillable: bool, token_usage: float
self,
local_prefillable: bool,
token_usage: float,
**kwargs,
) -> _NegotiateOutput:
out = self._negotiate_should_allow_prefill_pure(
prev_state=self._curr_state,
local_prefillable=local_prefillable,
token_usage=token_usage,
**kwargs,
)
self._curr_state = out.next_state
return out
@@ -91,6 +95,7 @@ class PrefillDelayer:
prev_state: Optional[_State],
local_prefillable: bool,
token_usage: float,
**kwargs,
) -> _NegotiateOutput:
# Compute local states
local_token_watermark_force_allow = (
@@ -100,10 +105,15 @@ class PrefillDelayer:
)
# Gather global states
global_prefillable, global_token_watermark_force_allow = self._gather_info(
tp0_info = self._gather_info(
local_prefillable=local_prefillable,
local_token_watermark_force_allow=local_token_watermark_force_allow,
**kwargs,
)
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]
# Compute derived global states
if global_prefillable.min().item() > 0:
@@ -123,6 +133,29 @@ class PrefillDelayer:
# Compute outputs
if prefillable_status == "all":
if kwargs is None:
exist_previous_wait = prev_state is not None
return _NegotiateOutput(
next_state=None,
output_allow=True,
output_reason="wait_success" if exist_previous_wait else "no_wait",
**debug_info,
)
max_running_requests = kwargs.get("max_running_requests", 0)
if (
max_running_requests - global_running_batch.max().item()
< global_max_prefill_bs.max().item()
and not self.enable_dp_attention
):
next_state = prev_state or _State()
next_state = next_state.bump_delayed_count()
return _NegotiateOutput(
next_state=next_state,
output_allow=False,
output_reason="delay",
**debug_info,
)
exist_previous_wait = prev_state is not None
return _NegotiateOutput(
next_state=None,
@@ -168,10 +201,15 @@ class PrefillDelayer:
raise NotImplementedError
def _gather_info(
self, local_prefillable: bool, local_token_watermark_force_allow: bool
self, local_prefillable: bool, local_token_watermark_force_allow: bool, **kwargs
):
local_info = torch.tensor(
[int(local_prefillable), int(local_token_watermark_force_allow)],
[
int(local_prefillable),
int(local_token_watermark_force_allow),
kwargs.get("running_batch", 0),
kwargs.get("max_prefill_bs", 0),
],
device="cpu",
dtype=torch.int64,
)
@@ -181,7 +219,7 @@ class PrefillDelayer:
group=self._cpu_group,
)
tp0_info = self._global_info_buffer[:, 0, :]
return tp0_info[:, 0], tp0_info[:, 1]
return tp0_info
class PrefillDelayerSinglePassExecutor:
@@ -204,11 +242,12 @@ class PrefillDelayerSinglePassExecutor:
metrics_collector=self._prefill_delayer._metrics_collector,
)
def negotiate_should_allow_prefill(self, local_prefillable: bool) -> bool:
def negotiate_should_allow_prefill(self, local_prefillable: bool, **kwargs) -> bool:
if not self._called:
self._result = self._prefill_delayer._negotiate_should_allow_prefill(
local_prefillable=local_prefillable,
token_usage=self._token_usage,
**kwargs,
)
return self._result.output_allow

10
python/sglang/srt/managers/schedule_policy.py
View File

@@ -381,6 +381,8 @@ class PrefillAdder:
rem_chunk_tokens: Optional[int],
mixed_with_decode_tokens: int = 0,
priority_scheduling_preemption_threshold: int = 0,
max_prefill_bs: int = 0,
max_running_requests: Optional[int] = None,
prefill_max_requests: Optional[int] = None,
prefill_delayer_single_pass: Optional[PrefillDelayerSinglePassExecutor] = None,
dllm_config: Optional[DllmConfig] = None,
@@ -427,8 +429,10 @@ class PrefillAdder:
priority_scheduling_preemption_threshold
)
self.nsa_prefill_cp_in_seq_split = is_nsa_prefill_cp_in_seq_split()
self.max_running_requests = max_running_requests
self.prefill_max_requests = prefill_max_requests
self.prefill_delayer_single_pass = prefill_delayer_single_pass
self.max_prefill_bs = max_prefill_bs
def _init_dllm_meta(self, dllm_config: DllmConfig):
self.dllm_block_size = dllm_config.block_size
@@ -767,10 +771,12 @@ class PrefillAdder:
if input_tokens >= self.rem_input_tokens and len(self.can_run_list) != 0:
return AddReqResult.OTHER
if (self.prefill_delayer_single_pass is not None) and (
not self.prefill_delayer_single_pass.negotiate_should_allow_prefill(
local_prefillable=True
local_prefillable=True,
running_batch=self.running_batch.batch_size(),
max_prefill_bs=self.max_prefill_bs,
max_running_requests=self.max_running_requests,
)
):
return AddReqResult.OTHER

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

@@ -794,6 +794,7 @@ class Scheduler(
self.schedule_low_priority_values_first,
)
self.prefill_delayer: Optional[PrefillDelayer] = None
self.max_prefill_bs: int = 0
if self.server_args.enable_prefill_delayer:
self.prefill_delayer = PrefillDelayer(
dp_size=self.dp_size,
@@ -805,6 +806,11 @@ class Scheduler(
),
max_delay_passes=self.server_args.prefill_delayer_max_delay_passes,
token_usage_low_watermark=self.server_args.prefill_delayer_token_usage_low_watermark,
device=(
self.tp_group.device
if self.server_args.disable_overlap_schedule
else "cpu"
),
)
# Enable preemption for priority scheduling.
self.try_preemption = self.enable_priority_scheduling
@@ -2049,6 +2055,8 @@ class Scheduler(
chunked_prefill_size,
running_bs if self.is_mixed_chunk else 0,
self.priority_scheduling_preemption_threshold,
max_prefill_bs=self.max_prefill_bs,
max_running_requests=self.max_running_requests,
prefill_max_requests=self.server_args.prefill_max_requests,
prefill_delayer_single_pass=prefill_delayer_single_pass,
dllm_config=self.dllm_config,
@@ -2163,6 +2171,7 @@ class Scheduler(
self.spec_algorithm,
chunked_req=self.chunked_req,
)
self.max_prefill_bs = max(self.max_prefill_bs, len(can_run_list))
if self.enable_hierarchical_cache:
# todo (zhiqiang): disable cuda graph execution if hicache loading triggered
new_batch.hicache_consumer_index = (