From 6e82183f5a28ca6a7f4dc0db4075c63d2dd6e23e Mon Sep 17 00:00:00 2001 From: sglang-bot Date: Wed, 25 Feb 2026 18:38:39 -0800 Subject: [PATCH] [Disagg] Route disagg prefill results through `process_batch_result` (#19364) --- python/sglang/srt/disaggregation/prefill.py | 8 +++----- python/sglang/srt/managers/scheduler.py | 2 ++ python/sglang/srt/managers/scheduler_pp_mixin.py | 7 ++----- 3 files changed, 7 insertions(+), 10 deletions(-) diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index a42cbe771..0ad5d6f19 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -370,7 +370,7 @@ class SchedulerDisaggregationPrefillMixin: # Launch the current batch if batch: result = self.run_batch(batch) - self.process_batch_result_disagg_prefill(batch, result) + self.process_batch_result(batch, result) else: self.self_check_during_idle() @@ -405,7 +405,7 @@ class SchedulerDisaggregationPrefillMixin: # Process the last batch if self.last_batch: 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: # When the server is idle, do self-check and re-init some states 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) req.time_stats.set_last_chunked_prefill_finish_time() - self.maybe_send_health_check_signal() - - if self.current_scheduler_metrics_enabled and batch.prefill_stats is not None: + if self.current_scheduler_metrics_enabled: can_run_cuda_graph = getattr(result, "can_run_cuda_graph", False) self.log_prefill_stats( prefill_stats=batch.prefill_stats, diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index bc5a35951..b885729c4 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -2441,6 +2441,8 @@ class Scheduler( elif batch.forward_mode.is_extend(): if batch.is_dllm(): self.process_batch_result_dllm(batch, result) + elif self.disaggregation_mode == DisaggregationMode.PREFILL: + self.process_batch_result_disagg_prefill(batch, result) else: self.process_batch_result_prefill(batch, result) elif batch.forward_mode.is_prebuilt(): diff --git a/python/sglang/srt/managers/scheduler_pp_mixin.py b/python/sglang/srt/managers/scheduler_pp_mixin.py index 1a65a3c3d..9a4396f2f 100644 --- a/python/sglang/srt/managers/scheduler_pp_mixin.py +++ b/python/sglang/srt/managers/scheduler_pp_mixin.py @@ -13,7 +13,7 @@ import torch.distributed from tqdm import tqdm 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.environ import envs from sglang.srt.layers.dp_attention import ( @@ -980,10 +980,7 @@ class SchedulerPPMixin: def _pp_process_batch_result( self: Scheduler, batch: ScheduleBatch, output_result: GenerationBatchResult ): - if self.disaggregation_mode == DisaggregationMode.PREFILL: - self.process_batch_result_disagg_prefill(batch, output_result) - else: - self.process_batch_result(batch, output_result) + self.process_batch_result(batch, output_result) def _pp_send_output_to_next_stage( self: Scheduler,