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

Minor code cleanup / improvement for PREBUILT_EXTEND mode (#12948)

Browse Source
This commit is contained in:
Liangsheng Yin
2025-11-10 11:59:54 +08:00
committed by GitHub
parent 83f89cc615
commit 61c7fe7aed
7 changed files with 50 additions and 55 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -588,7 +588,7 @@ class DecodePreallocQueue:
# 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_prebuilt_extend()
and self.scheduler.last_batch.forward_mode.is_prebuilt()
):
allocatable_tokens -= self.num_reserved_decode_tokens * len(
self.scheduler.last_batch.reqs
@@ -850,11 +850,14 @@ class SchedulerDisaggregationDecodeMixin:
self.launch_batch_sample_if_needed(batch_result)
self.last_batch = batch
def _run_batch_prebuilt_extend(
def _run_batch_prebuilt(
self: Scheduler, batch: ScheduleBatch
) -> GenerationBatchResult:
if batch.inner_idle_batch is not None:
return self.run_batch(batch.inner_idle_batch)
idle_batch = batch.inner_idle_batch
# Reset the inner idle batch to avoid reusing it.
batch.inner_idle_batch = None
return self.run_batch(idle_batch)
return GenerationBatchResult()
@@ -864,7 +867,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_prebuilt_extend():
if last_batch and last_batch.forward_mode.is_prebuilt():
# chunked prefill doesn't happen in decode instance.
assert self.chunked_req is None
# Filter finished batches.
@@ -947,8 +950,8 @@ class SchedulerDisaggregationDecodeMixin:
)
# construct fake completed prefill
new_batch.prepare_for_prebuilt_extend()
new_batch.process_prebuilt_extend(self.server_args, self.model_config)
new_batch.prepare_for_prebuilt()
new_batch.process_prebuilt(self.server_args, self.model_config)
return new_batch

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

@@ -20,13 +20,13 @@ if TYPE_CHECKING:
class ScheduleBatchDisaggregationDecodeMixin:
def prepare_for_prebuilt_extend(self: ScheduleBatch):
def prepare_for_prebuilt(self: ScheduleBatch):
"""
Prepare a prebuilt extend by populate metadata
Adapted from .prepare_for_extend().
"""
self.forward_mode = ForwardMode.PREBUILT_EXTEND
self.forward_mode = ForwardMode.PREBUILT
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)
@@ -100,7 +100,7 @@ class ScheduleBatchDisaggregationDecodeMixin:
self.model_config.vocab_size,
)
def process_prebuilt_extend(
def process_prebuilt(
self: ScheduleBatch, server_args: ServerArgs, model_config: ModelConfig
):
"""Assign the buffered last input id to schedule batch"""

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

@@ -1964,8 +1964,8 @@ class Scheduler(
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)
if batch.forward_mode.is_prebuilt():
return self._run_batch_prebuilt(batch)
# Run forward
if self.is_generation:
@@ -2099,8 +2099,8 @@ class Scheduler(
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_prebuilt():
self.process_batch_result_prebuilt(batch)
elif batch.forward_mode.is_idle():
if self.enable_overlap:
if result.copy_done is not None:

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

@@ -74,25 +74,21 @@ def _update_gather_batch(
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
]
batch.global_num_tokens = [mlp_sync_info.num_tokens]
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 = (
batch.global_num_tokens = mlp_sync_info.global_num_tokens
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
batch.is_extend_in_batch = mlp_sync_info.is_extend_in_batch
batch.tbo_split_seq_index = mlp_sync_info.tbo_split_seq_index
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
batch.can_run_dp_cuda_graph = mlp_sync_info.can_cuda_graph
def prepare_mlp_sync_batch_raw(
@@ -107,7 +103,7 @@ def prepare_mlp_sync_batch_raw(
offload_tags: set[str],
):
# Check if other DP workers have running batches
if local_batch is None or local_batch.forward_mode.is_prebuilt_extend():
if local_batch is None or local_batch.forward_mode.is_prebuilt():
num_tokens = 0
num_tokens_for_logprob = 0
elif local_batch.forward_mode.is_decode():
@@ -131,7 +127,7 @@ def prepare_mlp_sync_batch_raw(
can_cuda_graph = (
local_batch is None
or local_batch.forward_mode.is_decode_or_idle()
or local_batch.forward_mode.is_prebuilt_extend()
or local_batch.forward_mode.is_prebuilt()
) and not disable_cuda_graph
is_extend_in_batch = local_batch.forward_mode.is_extend() if local_batch else False
@@ -165,16 +161,13 @@ def prepare_mlp_sync_batch_raw(
)
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()
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
if local_batch is not None:
if need_idle_batch:
batch_to_gather = local_batch
if local_batch is None:
batch_to_gather = local_batch = get_idle_batch()
elif local_batch.forward_mode.is_prebuilt():
# NOTE: for prebuilt batch, we add an inner idle batch to run MLP sync
batch_to_gather = local_batch.inner_idle_batch = get_idle_batch()
_update_gather_batch(batch_to_gather, mlp_sync_info, require_mlp_tp_gather)
return local_batch

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

@@ -42,21 +42,20 @@ 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):
def process_batch_result_prebuilt(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)
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)

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

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

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

@@ -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() or forward_mode.is_prebuilt_extend():
elif forward_mode.is_idle() or forward_mode.is_prebuilt():
assert num_tokens == 0
return 0
else:
@@ -381,7 +381,7 @@ 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():
elif local_batch.forward_mode.is_prebuilt():
num_tokens = 0
else:
num_tokens = local_batch.extend_num_tokens