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

[Disagg] Route disagg prefill results through process_batch_result (#19364)

Browse Source
This commit is contained in:
sglang-bot
2026-02-25 18:38:39 -08:00
committed by GitHub
parent 914ed34757
commit 6e82183f5a
3 changed files with 7 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/disaggregation/prefill.py
+3 -5
View File
@@ -370,7 +370,7 @@ class SchedulerDisaggregationPrefillMixin:
# Launch the current batch # Launch the current batch
if batch: if batch:
result = self.run_batch(batch) result = self.run_batch(batch)
self.process_batch_result_disagg_prefill(batch, result) self.process_batch_result(batch, result)
else: else:
self.self_check_during_idle() self.self_check_during_idle()
@@ -405,7 +405,7 @@ class SchedulerDisaggregationPrefillMixin:
# Process the last batch # Process the last batch
if self.last_batch: if self.last_batch:
tmp_batch, tmp_result = self.result_queue.popleft() tmp_batch, tmp_result = self.result_queue.popleft()
self.process_batch_result_disagg_prefill(tmp_batch, tmp_result) self.process_batch_result(tmp_batch, tmp_result)
elif batch is None: elif batch is None:
# When the server is idle, do self-check and re-init some states # When the server is idle, do self-check and re-init some states
self.self_check_during_idle() self.self_check_during_idle()
@@ -533,9 +533,7 @@ class SchedulerDisaggregationPrefillMixin:
self.send_kv_chunk(req, last_chunk=False, end_idx=req.tmp_end_idx) self.send_kv_chunk(req, last_chunk=False, end_idx=req.tmp_end_idx)
req.time_stats.set_last_chunked_prefill_finish_time() req.time_stats.set_last_chunked_prefill_finish_time()
self.maybe_send_health_check_signal() if self.current_scheduler_metrics_enabled:
if self.current_scheduler_metrics_enabled and batch.prefill_stats is not None:
can_run_cuda_graph = getattr(result, "can_run_cuda_graph", False) can_run_cuda_graph = getattr(result, "can_run_cuda_graph", False)
self.log_prefill_stats( self.log_prefill_stats(
prefill_stats=batch.prefill_stats, prefill_stats=batch.prefill_stats,
python/sglang/srt/managers/scheduler.py
+2
View File
@@ -2441,6 +2441,8 @@ class Scheduler(
elif batch.forward_mode.is_extend(): elif batch.forward_mode.is_extend():
if batch.is_dllm(): if batch.is_dllm():
self.process_batch_result_dllm(batch, result) self.process_batch_result_dllm(batch, result)
elif self.disaggregation_mode == DisaggregationMode.PREFILL:
self.process_batch_result_disagg_prefill(batch, result)
else: else:
self.process_batch_result_prefill(batch, result) self.process_batch_result_prefill(batch, result)
elif batch.forward_mode.is_prebuilt(): elif batch.forward_mode.is_prebuilt():
python/sglang/srt/managers/scheduler_pp_mixin.py
+2 -5
View File
@@ -13,7 +13,7 @@ import torch.distributed
from tqdm import tqdm from tqdm import tqdm
from sglang.srt.disaggregation.base.conn import KVPoll from sglang.srt.disaggregation.base.conn import KVPoll
from sglang.srt.disaggregation.utils import DisaggregationMode, poll_and_all_reduce from sglang.srt.disaggregation.utils import poll_and_all_reduce
from sglang.srt.distributed.parallel_state import P2PWork from sglang.srt.distributed.parallel_state import P2PWork
from sglang.srt.environ import envs from sglang.srt.environ import envs
from sglang.srt.layers.dp_attention import ( from sglang.srt.layers.dp_attention import (
@@ -980,10 +980,7 @@ class SchedulerPPMixin:
def _pp_process_batch_result( def _pp_process_batch_result(
self: Scheduler, batch: ScheduleBatch, output_result: GenerationBatchResult self: Scheduler, batch: ScheduleBatch, output_result: GenerationBatchResult
): ):
if self.disaggregation_mode == DisaggregationMode.PREFILL: self.process_batch_result(batch, output_result)
self.process_batch_result_disagg_prefill(batch, output_result)
else:
self.process_batch_result(batch, output_result)
def _pp_send_output_to_next_stage( def _pp_send_output_to_next_stage(
self: Scheduler, self: Scheduler,