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

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,

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():

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,