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

[Refactor / Style] Unify all event loops (except for PP) (#12959)

Browse Source
This commit is contained in:
Liangsheng Yin
2025-11-10 14:49:12 +08:00
committed by GitHub
parent e123648b36
commit dc8a5a1ce7
5 changed files with 102 additions and 100 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
python/sglang/srt/disaggregation/base/conn.py
View File

@@ -141,6 +141,18 @@ class BaseKVReceiver(ABC):
"""
...
def clear(self):
"""
Clear any internal states.
"""
pass
def abort(self):
"""
Abort the current transfer.
"""
pass
class BaseKVBootstrapServer(ABC):
@abstractmethod

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

@@ -693,6 +693,50 @@ class DecodeTransferQueue:
def extend(self, decode_reqs: List[DecodeRequest]) -> None:
self.queue.extend(decode_reqs)
def _commit_transfer_to_req(self, decode_req: DecodeRequest) -> None:
idx = decode_req.metadata_buffer_index
(
output_id,
cached_tokens,
output_token_logprobs_val,
output_token_logprobs_idx,
output_top_logprobs_val,
output_top_logprobs_idx,
output_topk_p,
output_topk_index,
output_hidden_states,
) = self.metadata_buffers.get_buf(idx)
decode_req.req.output_ids.append(output_id[0].item())
decode_req.req.cached_tokens = cached_tokens[0].item()
if not self.spec_algorithm.is_none():
decode_req.req.output_topk_p = output_topk_p
decode_req.req.output_topk_index = output_topk_index
decode_req.req.hidden_states_tensor = output_hidden_states
if decode_req.req.return_logprob:
decode_req.req.output_token_logprobs_val.append(
output_token_logprobs_val[0].item()
)
decode_req.req.output_token_logprobs_idx.append(
output_token_logprobs_idx[0].item()
)
decode_req.req.output_top_logprobs_val.append(
output_top_logprobs_val[: decode_req.req.top_logprobs_num].tolist()
)
decode_req.req.output_top_logprobs_idx.append(
output_top_logprobs_idx[: decode_req.req.top_logprobs_num].tolist()
)
decode_req.kv_receiver.clear()
decode_req.kv_receiver = None
trace_slice_end(
RequestStage.DECODE_TRANSFERRED,
decode_req.req.rid,
auto_next_anon=True,
)
decode_req.req.time_stats.wait_queue_entry_time = time.perf_counter()
def pop_transferred(self) -> List[Req]:
if not self.queue:
return []
@@ -725,58 +769,9 @@ class DecodeTransferQueue:
self.scheduler.metrics_collector.increment_transfer_failed_reqs()
continue
elif poll == KVPoll.Success:
idx = decode_req.metadata_buffer_index
(
output_id,
cached_tokens,
output_token_logprobs_val,
output_token_logprobs_idx,
output_top_logprobs_val,
output_top_logprobs_idx,
output_topk_p,
output_topk_index,
output_hidden_states,
) = self.metadata_buffers.get_buf(idx)
decode_req.req.output_ids.append(output_id[0].item())
decode_req.req.cached_tokens = cached_tokens[0].item()
if not self.spec_algorithm.is_none():
decode_req.req.output_topk_p = output_topk_p
decode_req.req.output_topk_index = output_topk_index
decode_req.req.hidden_states_tensor = output_hidden_states
if decode_req.req.return_logprob:
decode_req.req.output_token_logprobs_val.append(
output_token_logprobs_val[0].item()
)
decode_req.req.output_token_logprobs_idx.append(
output_token_logprobs_idx[0].item()
)
decode_req.req.output_top_logprobs_val.append(
output_top_logprobs_val[
: decode_req.req.top_logprobs_num
].tolist()
)
decode_req.req.output_top_logprobs_idx.append(
output_top_logprobs_idx[
: decode_req.req.top_logprobs_num
].tolist()
)
if hasattr(decode_req.kv_receiver, "clear"):
decode_req.kv_receiver.clear()
decode_req.kv_receiver = None
self._commit_transfer_to_req(decode_req)
indices_to_remove.add(i)
decode_req.req.time_stats.wait_queue_entry_time = time.perf_counter()
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,
KVPoll.WaitingForInput,

84
python/sglang/srt/disaggregation/prefill.py
View File

@@ -310,6 +310,37 @@ class SchedulerDisaggregationPrefillMixin:
Mixin for Scheduler to handle disaggregation prefill
"""
def get_next_disagg_prefill_batch_to_run(
self: Scheduler,
) -> Optional[ScheduleBatch]:
if self.last_batch and self.last_batch.forward_mode.is_extend():
if self.chunked_req:
# Move the chunked request out of the batch so that we can merge
# only finished requests to running_batch.
self.last_batch.filter_batch(chunked_req_to_exclude=self.chunked_req)
self.tree_cache.cache_unfinished_req(self.chunked_req, chunked=True)
if self.enable_overlap:
# Delay KV transfer to process_batch_result_disagg_prefill when overlap is enabled to ensure results are resolved
self.chunked_req.tmp_end_idx = min(
len(self.chunked_req.fill_ids),
len(self.chunked_req.origin_input_ids),
)
else:
self.send_kv_chunk(self.chunked_req)
# chunked request keeps its rid but will get a new req_pool_idx
self.req_to_token_pool.free(self.chunked_req.req_pool_idx)
self.running_batch.batch_is_full = False
batch = self.get_new_batch_prefill()
if self.require_mlp_sync:
batch = self.prepare_mlp_sync_batch(batch)
if batch:
attrs = {"bid": hex(id(batch)), "batch_size": batch.batch_size()}
trace_event_batch("schedule", batch.reqs, attrs=attrs)
return batch
@torch.no_grad()
def event_loop_normal_disagg_prefill(self: Scheduler) -> None:
"""A normal scheduler loop for prefill worker in disaggregation mode."""
@@ -320,26 +351,17 @@ class SchedulerDisaggregationPrefillMixin:
self.waiting_queue.extend(
self.disagg_prefill_bootstrap_queue.pop_bootstrapped()
)
self.process_prefill_chunk()
batch = self.get_new_batch_prefill()
if batch:
attrs = {"bid": hex(id(batch)), "batch_size": batch.batch_size()}
trace_event_batch("schedule", batch.reqs, attrs=attrs)
if self.require_mlp_sync:
batch = self.prepare_mlp_sync_batch(batch)
batch = self.get_next_disagg_prefill_batch_to_run()
self.cur_batch = batch
if batch:
result = self.run_batch(batch)
self.process_batch_result_disagg_prefill(batch, result)
if len(self.disagg_prefill_inflight_queue) > 0:
self.process_disagg_prefill_inflight_queue()
if batch is None and len(self.disagg_prefill_inflight_queue) == 0:
else:
self.self_check_during_idle()
self.process_disagg_prefill_inflight_queue()
self.last_batch = batch
# HACK (byronhsu): reset the batch_is_full flag because we never enter update_running_batch which resets it
# Otherwise, it hangs under high concurrency
@@ -355,14 +377,7 @@ class SchedulerDisaggregationPrefillMixin:
self.waiting_queue.extend(
self.disagg_prefill_bootstrap_queue.pop_bootstrapped()
)
self.process_prefill_chunk()
batch = self.get_new_batch_prefill()
if batch:
attrs = {"bid": hex(id(batch)), "batch_size": batch.batch_size()}
trace_event_batch("schedule", batch.reqs, attrs=attrs)
if self.require_mlp_sync:
batch = self.prepare_mlp_sync_batch(batch)
batch = self.get_next_disagg_prefill_batch_to_run()
self.cur_batch = batch
batch_result = None
@@ -373,15 +388,13 @@ class SchedulerDisaggregationPrefillMixin:
if self.last_batch:
tmp_batch, tmp_result = self.result_queue.popleft()
self.process_batch_result_disagg_prefill(tmp_batch, tmp_result)
elif batch is None:
self.self_check_during_idle()
if len(self.disagg_prefill_inflight_queue) > 0:
self.process_disagg_prefill_inflight_queue()
self.process_disagg_prefill_inflight_queue()
self.launch_batch_sample_if_needed(batch_result)
if batch is None and len(self.disagg_prefill_inflight_queue) == 0:
self.self_check_during_idle()
self.last_batch = batch
# HACK (byronhsu): reset the batch_is_full flag because we never enter update_running_batch which resets it
# Otherwise, it hangs under high concurrency
@@ -599,25 +612,6 @@ class SchedulerDisaggregationPrefillMixin:
return transferred_rids
def process_prefill_chunk(self: Scheduler) -> None:
if self.last_batch and self.last_batch.forward_mode.is_extend():
if self.chunked_req:
# Move the chunked request out of the batch so that we can merge
# only finished requests to running_batch.
self.last_batch.filter_batch(chunked_req_to_exclude=self.chunked_req)
self.tree_cache.cache_unfinished_req(self.chunked_req, chunked=True)
if self.enable_overlap:
# Delay KV transfer to process_batch_result_disagg_prefill when overlap is enabled to ensure results are resolved
self.chunked_req.tmp_end_idx = min(
len(self.chunked_req.fill_ids),
len(self.chunked_req.origin_input_ids),
)
else:
self.send_kv_chunk(self.chunked_req)
# chunked request keeps its rid but will get a new req_pool_idx
self.req_to_token_pool.free(self.chunked_req.req_pool_idx)
self.running_batch.batch_is_full = False
def send_kv_chunk(
self: Scheduler,
req: Req,

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

@@ -2443,15 +2443,13 @@ class Scheduler(
for decode_req in self.disagg_decode_prealloc_queue.queue:
if recv_req.abort_all or decode_req.req.rid.startswith(recv_req.rid):
logger.debug(f"Abort prealloc queue request. {decode_req.req.rid=}")
if hasattr(decode_req.kv_receiver, "abort"):
decode_req.kv_receiver.abort()
decode_req.kv_receiver.abort()
# Abort requests waiting for kvcache to release tree cache
for decode_req in self.disagg_decode_transfer_queue.queue:
if recv_req.abort_all or decode_req.req.rid.startswith(recv_req.rid):
logger.debug(f"Abort transfer queue request. {decode_req.req.rid=}")
if hasattr(decode_req.kv_receiver, "abort"):
decode_req.kv_receiver.abort()
decode_req.kv_receiver.abort()
# Delete requests in the running batch
if self.cur_batch is self.running_batch or self.cur_batch is None:

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

@@ -217,7 +217,10 @@ class SchedulerRuntimeCheckerMixin:
self.tree_cache.sanity_check()
def self_check_during_idle(self: Scheduler):
if self.disaggregation_mode == DisaggregationMode.DECODE:
if self.disaggregation_mode == DisaggregationMode.PREFILL:
if len(self.disagg_prefill_inflight_queue) > 0:
return
elif self.disaggregation_mode == DisaggregationMode.DECODE:
queue_size = (
len(self.waiting_queue)
+ len(self.disagg_decode_transfer_queue.queue)