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

Support offline generation scenario for prefill delayer (#16363)

Browse Source
This commit is contained in:
fzyzcjy
2026-01-04 11:05:10 +08:00
committed by GitHub
parent 5d4b7c78bf
commit e797f0c570
4 changed files with 138 additions and 55 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -1,4 +1,5 @@
import logging
from typing import Optional
import torch
@@ -35,21 +36,7 @@ class PrefillDelayer:
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:
def _negotiate_should_allow_prefill(self, local_prefillable: bool) -> bool:
tp0_info = self._gather_info(local_prefillable=local_prefillable)
global_prefillable = tp0_info[:, 0]
global_exists_not_prefillable = global_prefillable.min().item() == 0
@@ -70,3 +57,42 @@ class PrefillDelayer:
self.curr_delayed_count = 0
return True
def _gather_info(self, local_prefillable: bool):
local_info = torch.tensor(
[int(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
class PrefillDelayerSinglePassExecutor:
def __init__(self, prefill_delayer: PrefillDelayer):
self._prefill_delayer = prefill_delayer
self._result: Optional[bool] = None
@property
def _called(self) -> bool:
return self._result is not None
def __enter__(self):
return self
def __exit__(self, exc_type, exc_val, exc_tb):
if not self._called:
self.negotiate_should_allow_prefill(local_prefillable=False)
return False
def negotiate_should_allow_prefill(self, local_prefillable: bool) -> bool:
if not self._called:
self._result = self._prefill_delayer._negotiate_should_allow_prefill(
local_prefillable=local_prefillable
)
return self._result

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

@@ -1,5 +1,7 @@
from __future__ import annotations
from sglang.srt.managers.prefill_delayer import PrefillDelayerSinglePassExecutor
# Copyright 2023-2024 SGLang Team
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
@@ -327,6 +329,7 @@ class PrefillAdder:
mixed_with_decode_tokens: int = 0,
priority_scheduling_preemption_threshold: int = 0,
prefill_max_requests: Optional[int] = None,
prefill_delayer_single_pass: Optional[PrefillDelayerSinglePassExecutor] = None,
):
self.page_size = page_size
self.tree_cache = tree_cache
@@ -366,6 +369,7 @@ class PrefillAdder:
)
self.nsa_prefill_cp_in_seq_split = is_nsa_prefill_cp_in_seq_split()
self.prefill_max_requests = prefill_max_requests
self.prefill_delayer_single_pass = prefill_delayer_single_pass
def _get_running_request_total_token_offset(self, req: Req) -> int:
return (
@@ -614,6 +618,13 @@ 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
)
):
return AddReqResult.OTHER
if self.rem_chunk_tokens is None or input_tokens <= self.rem_chunk_tokens:
# Non-chunked prefill
self.can_run_list.append(req)

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

@@ -22,6 +22,7 @@ import sys
import time
from collections import deque
from concurrent import futures
from contextlib import nullcontext
from dataclasses import dataclass
from http import HTTPStatus
from typing import Any, Deque, Dict, List, Optional, Tuple, Union
@@ -120,7 +121,10 @@ 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.prefill_delayer import (
PrefillDelayer,
PrefillDelayerSinglePassExecutor,
)
from sglang.srt.managers.schedule_batch import (
FINISH_ABORT,
ModelWorkerBatch,
@@ -1911,6 +1915,18 @@ class Scheduler(
return res
def get_new_batch_prefill(self) -> Optional[ScheduleBatch]:
with (
PrefillDelayerSinglePassExecutor(self.prefill_delayer)
if self.prefill_delayer
else nullcontext()
) as prefill_delayer_single_pass:
return self._get_new_batch_prefill_raw(
prefill_delayer_single_pass=prefill_delayer_single_pass
)
def _get_new_batch_prefill_raw(
self, prefill_delayer_single_pass: Optional[PrefillDelayerSinglePassExecutor]
) -> Optional[ScheduleBatch]:
# Check if the grammar is ready in the grammar queue
if self.grammar_queue:
self.move_ready_grammar_requests()
@@ -1919,19 +1935,8 @@ class Scheduler(
# Reset batch_is_full to try preemption with a prefill adder.
self.running_batch.batch_is_full = False
# 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))
self.running_batch.batch_is_full or len(self.waiting_queue) == 0
) and self.chunked_req is None:
return None
@@ -1981,6 +1986,7 @@ class Scheduler(
running_bs if self.is_mixed_chunk else 0,
self.priority_scheduling_preemption_threshold,
prefill_max_requests=self.server_args.prefill_max_requests,
prefill_delayer_single_pass=prefill_delayer_single_pass,
)
if self.chunked_req is not None: