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

Refactor / Unify event loop across PD-Disagg, Overlap, DP-Attn cases (#12839)

Browse Source 
Co-authored-by: cctry <17473714+cctry@users.noreply.github.com>
This commit is contained in:
Liangsheng Yin
2025-11-10 00:42:50 +08:00
committed by GitHub
parent f5b3ccd9a5
commit 4f65a64666
11 changed files with 142 additions and 144 deletions
Download Patch File Download Diff File Expand all files Collapse all files

142
python/sglang/srt/disaggregation/decode.py
View File

@@ -48,6 +48,7 @@ from sglang.srt.disaggregation.utils import (
)
from sglang.srt.layers.dp_attention import get_attention_tp_size
from sglang.srt.managers.schedule_batch import FINISH_ABORT, RequestStage, ScheduleBatch
from sglang.srt.managers.utils import GenerationBatchResult
from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache
from sglang.srt.mem_cache.memory_pool import (
@@ -58,11 +59,7 @@ from sglang.srt.mem_cache.memory_pool import (
ReqToTokenPool,
SWAKVPool,
)
from sglang.srt.tracing.trace import (
trace_event_batch,
trace_slice_batch,
trace_slice_end,
)
from sglang.srt.tracing.trace import trace_event_batch, trace_slice_end
from sglang.srt.utils import get_int_env_var
from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter
@@ -587,11 +584,11 @@ class DecodePreallocQueue:
need_space_for_single_req,
)
# Note: if the last fake extend just finishes, and we enter `pop_preallocated` immediately in the next iteration
# Note: if the last prebuilt extend just finishes, and we enter `pop_preallocated` immediately in the next iteration
# the extend batch is not in any queue, so we need to explicitly add the tokens slots here
if (
self.scheduler.last_batch
and self.scheduler.last_batch.forward_mode.is_extend()
and self.scheduler.last_batch.forward_mode.is_prebuilt_extend()
):
allocatable_tokens -= self.num_reserved_decode_tokens * len(
self.scheduler.last_batch.reqs
@@ -773,29 +770,12 @@ class DecodeTransferQueue:
indices_to_remove.add(i)
decode_req.req.time_stats.wait_queue_entry_time = time.perf_counter()
decode_req.req.check_finished()
if decode_req.req.finished():
# finish immediately
decode_req.req.time_stats.forward_entry_time = (
decode_req.req.time_stats.completion_time
) = time.perf_counter()
self.scheduler.stream_output(
[decode_req.req], decode_req.req.return_logprob
)
self.tree_cache.cache_finished_req(decode_req.req)
trace_slice_end(
RequestStage.DECODE_QUICK_FINISH,
decode_req.req.rid,
thread_finish_flag=True,
)
else:
transferred_reqs.append(decode_req.req)
trace_slice_end(
RequestStage.DECODE_TRANSFERRED,
decode_req.req.rid,
auto_next_anon=True,
)
transferred_reqs.append(decode_req.req)
trace_slice_end(
RequestStage.DECODE_TRANSFERRED,
decode_req.req.rid,
auto_next_anon=True,
)
elif poll in [
KVPoll.Bootstrapping,
@@ -835,31 +815,9 @@ class SchedulerDisaggregationDecodeMixin:
if batch:
# Generate fake extend output.
if batch.forward_mode.is_extend():
# Note: Logprobs should be handled on the prefill engine.
self.stream_output(
batch.reqs, any(req.return_logprob for req in batch.reqs)
)
trace_slice_batch(RequestStage.DECODE_FAKE_OUTPUT, batch.reqs)
if self.require_mlp_sync:
self._prepare_idle_batch_and_run()
else:
if self.require_mlp_sync:
self.prepare_mlp_sync_batch(batch)
result = self.run_batch(batch)
self.process_batch_result(batch, result)
elif self.require_mlp_sync:
batch, _ = self._prepare_idle_batch_and_run()
queue_size = (
len(self.waiting_queue)
+ len(self.disagg_decode_transfer_queue.queue)
+ len(self.disagg_decode_prealloc_queue.queue)
)
if self.server_args.disaggregation_decode_enable_offload_kvcache:
queue_size += len(self.decode_offload_manager.ongoing_offload)
if batch is None and queue_size == 0:
result = self.run_batch(batch)
self.process_batch_result(batch, result)
else:
self.self_check_during_idle()
self.last_batch = batch
@@ -868,7 +826,6 @@ class SchedulerDisaggregationDecodeMixin:
def event_loop_overlap_disagg_decode(self: Scheduler):
self.result_queue = deque()
self.last_batch: Optional[ScheduleBatch] = None
self.last_batch_in_queue = False # last batch is modified in-place, so we need another variable to track if it's extend
while True:
@@ -878,68 +835,28 @@ class SchedulerDisaggregationDecodeMixin:
self.process_decode_queue()
batch = self.get_next_disagg_decode_batch_to_run()
self.cur_batch = batch
last_batch_in_queue = False
batch_result = None
if batch:
# Generate fake extend output.
if batch.forward_mode.is_extend():
# Note: Logprobs should be handled on the prefill engine.
self.stream_output(
batch.reqs, any(req.return_logprob for req in batch.reqs)
)
trace_slice_batch(RequestStage.DECODE_FAKE_OUTPUT, batch.reqs)
if self.require_mlp_sync:
batch_, batch_result = self._prepare_idle_batch_and_run(
delay_process=True
)
if batch_:
self.result_queue.append((batch_.copy(), batch_result))
last_batch_in_queue = True
else:
if self.require_mlp_sync:
self.prepare_mlp_sync_batch(batch)
batch_result = self.run_batch(batch)
self.result_queue.append((batch.copy(), batch_result))
last_batch_in_queue = True
batch_result = self.run_batch(batch)
self.result_queue.append((batch.copy(), batch_result))
elif self.require_mlp_sync:
batch, batch_result = self._prepare_idle_batch_and_run(
delay_process=True
)
if batch:
self.result_queue.append((batch.copy(), batch_result))
last_batch_in_queue = True
# Process the results of the previous batch but skip if the last batch is extend
if self.last_batch and self.last_batch_in_queue:
if self.last_batch:
tmp_batch, tmp_result = self.result_queue.popleft()
self.process_batch_result(tmp_batch, tmp_result)
self.launch_batch_sample_if_needed(batch_result)
queue_size = (
len(self.waiting_queue)
+ len(self.disagg_decode_transfer_queue.queue)
+ len(self.disagg_decode_prealloc_queue.queue)
)
if self.server_args.disaggregation_decode_enable_offload_kvcache:
queue_size += len(self.decode_offload_manager.ongoing_offload)
if batch is None and queue_size == 0:
elif batch is None:
self.self_check_during_idle()
self.launch_batch_sample_if_needed(batch_result)
self.last_batch = batch
self.last_batch_in_queue = last_batch_in_queue
def _prepare_idle_batch_and_run(self: Scheduler, delay_process=False):
batch = self.prepare_mlp_sync_batch(None)
result = None
if batch:
result = self.run_batch(batch)
if not delay_process:
self.process_batch_result(batch, result)
return batch, result
def _run_batch_prebuilt_extend(
self: Scheduler, batch: ScheduleBatch
) -> GenerationBatchResult:
if batch.inner_idle_batch is not None:
return self.run_batch(batch.inner_idle_batch)
return GenerationBatchResult()
def get_next_disagg_decode_batch_to_run(
self: Scheduler,
@@ -947,7 +864,7 @@ class SchedulerDisaggregationDecodeMixin:
"""Create fake completed prefill if possible and merge with running batch"""
# Merge the prefill batch into the running batch
last_batch = self.last_batch
if last_batch and last_batch.forward_mode.is_extend():
if last_batch and last_batch.forward_mode.is_prebuilt_extend():
# chunked prefill doesn't happen in decode instance.
assert self.chunked_req is None
# Filter finished batches.
@@ -971,6 +888,13 @@ class SchedulerDisaggregationDecodeMixin:
self.running_batch = self.update_running_batch(self.running_batch)
ret = self.running_batch if not self.running_batch.is_empty() else None
# 1. decode + None -> decode + idle
# 2. decode + prebuilt -> decode + idle (idle forward, prebuilt returns)
# 3. prebuilt + None -> None (None forward, prebuilt returns) + None
# 4. prebuilt + decode + None -> idle (idle forward, prebuilt returns) + decode + idle
if self.require_mlp_sync:
ret = self.prepare_mlp_sync_batch(ret)
if ret:
attrs = {"bid": hex(id(ret)), "batch_size": ret.batch_size()}
trace_event_batch("schedule", ret.reqs, attrs=attrs)

2
python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py
View File

@@ -26,7 +26,7 @@ class ScheduleBatchDisaggregationDecodeMixin:
Adapted from .prepare_for_extend().
"""
self.forward_mode = ForwardMode.EXTEND
self.forward_mode = ForwardMode.PREBUILT_EXTEND
reqs = self.reqs
input_ids = [r.fill_ids[len(r.prefix_indices) :] for r in reqs]
extend_num_tokens = sum(len(ids) for ids in input_ids)

1
python/sglang/srt/managers/schedule_batch.py
View File

@@ -1010,6 +1010,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
orig_seq_lens: torch.Tensor = None # shape: [b], int32
# For DP attention
inner_idle_batch: Optional[ScheduleBatch] = None
global_num_tokens: Optional[List[int]] = None
global_num_tokens_for_logprob: Optional[List[int]] = None
is_extend_in_batch: bool = False

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

@@ -1963,6 +1963,10 @@ class Scheduler(
for req in batch.reqs:
req.time_stats.prefill_start_time = current_time
# Place holder handling for pd-disagg decode event loop
if batch.forward_mode.is_prebuilt_extend():
return self._run_batch_prebuilt_extend(batch)
# Run forward
if self.is_generation:
batch_or_worker_batch = batch
@@ -2093,10 +2097,10 @@ class Scheduler(
if batch.forward_mode.is_decode():
self.process_batch_result_decode(batch, result)
trace_slice_batch(RequestStage.DECODE_LOOP, batch.reqs)
elif batch.forward_mode.is_extend():
self.process_batch_result_prefill(batch, result)
elif batch.forward_mode.is_prebuilt_extend():
self.process_batch_result_prebuilt_extend(batch)
elif batch.forward_mode.is_idle():
if self.enable_overlap:
if result.copy_done is not None:

65
python/sglang/srt/managers/scheduler_dp_attn_mixin.py
View File

@@ -30,6 +30,8 @@ class MLPSyncBatchInfo:
tp0_info: torch.Tensor = None
global_num_tokens: list[int] = None
global_num_tokens_for_logprob: list[int] = None
tbo_split_seq_index: torch.Tensor = None
global_forward_mode: int = None
def _get_local_tensor(self, device, dtype=torch.int64) -> torch.Tensor:
return torch.tensor(
@@ -67,6 +69,32 @@ class MLPSyncBatchInfo:
self.is_extend_in_batch = bool(tp0_info[:, 3].max().item())
def _update_gather_batch(
batch: ScheduleBatch,
mlp_sync_info: MLPSyncBatchInfo,
require_mlp_tp_gather: bool,
):
local_batch = batch
# TODO: handle the case when moe_dense_tp_size != 1
if not require_mlp_tp_gather:
local_batch.global_num_tokens = [mlp_sync_info.num_tokens]
local_batch.global_num_tokens_for_logprob = [
mlp_sync_info.num_tokens_for_logprob
]
else:
local_batch.global_num_tokens = mlp_sync_info.global_num_tokens
local_batch.global_num_tokens_for_logprob = (
mlp_sync_info.global_num_tokens_for_logprob
)
local_batch.is_extend_in_batch = mlp_sync_info.is_extend_in_batch
local_batch.tbo_split_seq_index = mlp_sync_info.tbo_split_seq_index
local_batch.global_forward_mode = mlp_sync_info.global_forward_mode
# Check forward mode for cuda graph
local_batch.can_run_dp_cuda_graph = mlp_sync_info.can_cuda_graph
def prepare_mlp_sync_batch_raw(
local_batch: ScheduleBatch,
dp_size: int,
@@ -79,7 +107,7 @@ def prepare_mlp_sync_batch_raw(
offload_tags: set[str],
):
# Check if other DP workers have running batches
if local_batch is None:
if local_batch is None or local_batch.forward_mode.is_prebuilt_extend():
num_tokens = 0
num_tokens_for_logprob = 0
elif local_batch.forward_mode.is_decode():
@@ -101,7 +129,9 @@ def prepare_mlp_sync_batch_raw(
num_tokens_for_logprob = local_batch.batch_size()
can_cuda_graph = (
local_batch is None or local_batch.forward_mode.is_decode_or_idle()
local_batch is None
or local_batch.forward_mode.is_decode_or_idle()
or local_batch.forward_mode.is_prebuilt_extend()
) and not disable_cuda_graph
is_extend_in_batch = local_batch.forward_mode.is_extend() if local_batch else False
@@ -128,29 +158,24 @@ def prepare_mlp_sync_batch_raw(
)
mlp_sync_info.all_gather(device=device, group=group)
tbo_split_seq_index, global_forward_mode = tbo_preparer.compute_output(
mlp_sync_info.tp0_info[:, 4:6],
mlp_sync_info.tbo_split_seq_index, mlp_sync_info.global_forward_mode = (
tbo_preparer.compute_output(
mlp_sync_info.tp0_info[:, 4:6],
)
)
if local_batch is None and max(mlp_sync_info.global_num_tokens) > 0:
need_idle_batch = max(mlp_sync_info.global_num_tokens) > 0
if need_idle_batch and local_batch is None:
local_batch = get_idle_batch()
if local_batch is not None:
# TODO: handle the case when moe_dense_tp_size != 1
if not require_mlp_tp_gather:
local_batch.global_num_tokens = [num_tokens]
local_batch.global_num_tokens_for_logprob = [num_tokens_for_logprob]
else:
local_batch.global_num_tokens = mlp_sync_info.global_num_tokens
local_batch.global_num_tokens_for_logprob = (
mlp_sync_info.global_num_tokens_for_logprob
)
local_batch.is_extend_in_batch = mlp_sync_info.is_extend_in_batch
local_batch.tbo_split_seq_index = tbo_split_seq_index
local_batch.global_forward_mode = global_forward_mode
batch_to_gather = local_batch
if need_idle_batch and local_batch.forward_mode.is_prebuilt_extend():
# NOTE: for prebuilt batch, we add an inner idle batch to run MLP sync
local_batch.inner_idle_batch = get_idle_batch()
batch_to_gather = local_batch.inner_idle_batch
# Check forward mode for cuda graph
local_batch.can_run_dp_cuda_graph = mlp_sync_info.can_cuda_graph
if local_batch is not None:
_update_gather_batch(batch_to_gather, mlp_sync_info, require_mlp_tp_gather)
return local_batch

22
python/sglang/srt/managers/scheduler_output_processor_mixin.py
View File

@@ -20,7 +20,7 @@ from sglang.srt.managers.schedule_batch import (
RequestStage,
ScheduleBatch,
)
from sglang.srt.tracing.trace import trace_slice
from sglang.srt.tracing.trace import trace_slice, trace_slice_batch, trace_slice_end
from sglang.srt.utils.common import ceil_div
if TYPE_CHECKING:
@@ -42,6 +42,26 @@ class SchedulerOutputProcessorMixin:
We put them into a separate file to make the `scheduler.py` shorter.
"""
def process_batch_result_prebuilt_extend(self: Scheduler, batch: ScheduleBatch):
assert self.disaggregation_mode == DisaggregationMode.DECODE
for req in batch.reqs:
for req in batch.reqs:
req.check_finished()
if req.finished():
req.time_stats.forward_entry_time = (
req.time_stats.completion_time
) = time.perf_counter()
trace_slice_end(
RequestStage.DECODE_QUICK_FINISH,
req.rid,
thread_finish_flag=True,
)
self.tree_cache.cache_finished_req(req)
# Note: Logprobs should be handled on the prefill engine.
trace_slice_batch(RequestStage.DECODE_FAKE_OUTPUT, batch.reqs)
self.stream_output(batch.reqs, batch.return_logprob)
def process_batch_result_prefill(
self: Scheduler,
batch: ScheduleBatch,

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

@@ -217,6 +217,17 @@ class SchedulerRuntimeCheckerMixin:
self.tree_cache.sanity_check()
def self_check_during_idle(self: Scheduler):
if self.disaggregation_mode == DisaggregationMode.DECODE:
queue_size = (
len(self.waiting_queue)
+ len(self.disagg_decode_transfer_queue.queue)
+ len(self.disagg_decode_prealloc_queue.queue)
)
if self.server_args.disaggregation_decode_enable_offload_kvcache:
queue_size += len(self.decode_offload_manager.ongoing_offload)
if queue_size:
return
self.check_memory()
self.check_tree_cache()
self.new_token_ratio = self.init_new_token_ratio

7
python/sglang/srt/model_executor/forward_batch_info.py
View File

@@ -79,6 +79,10 @@ class ForwardMode(IntEnum):
DRAFT_EXTEND_V2 = auto()
# Used in disaggregated decode worker
# Represent a batch of requests having their KV cache ready to start decoding
PREBUILT_EXTEND = auto()
# Split Prefill for PD multiplexing
SPLIT_PREFILL = auto()
@@ -148,6 +152,9 @@ class ForwardMode(IntEnum):
and not self.is_draft_extend()
)
def is_prebuilt_extend(self):
return self == ForwardMode.PREBUILT_EXTEND
@total_ordering
class CaptureHiddenMode(IntEnum):

6
python/sglang/srt/two_batch_overlap.py
View File

@@ -68,7 +68,7 @@ def get_token_num_per_seq(
# TODO: may smartly disable TBO when batch size is too small b/c it will slow down
def compute_split_seq_index(
forward_mode: "ForwardMode",
forward_mode: ForwardMode,
num_tokens: int,
extend_lens: Optional[Sequence[int]],
token_num_per_seq: Optional[int],
@@ -79,7 +79,7 @@ def compute_split_seq_index(
elif forward_mode.is_target_verify() or forward_mode.is_decode():
assert token_num_per_seq is not None
return (num_tokens // token_num_per_seq) // 2
elif forward_mode.is_idle():
elif forward_mode.is_idle() or forward_mode.is_prebuilt_extend():
assert num_tokens == 0
return 0
else:
@@ -381,6 +381,8 @@ class TboDPAttentionPreparer:
or local_batch.forward_mode.is_decode()
):
num_tokens = local_batch.batch_size() * token_num_per_seq
elif local_batch.forward_mode.is_prebuilt_extend():
num_tokens = 0
else:
num_tokens = local_batch.extend_num_tokens
self.local_tbo_split_seq_index = compute_split_seq_index(