From 8a9ca41fda772548775781b86d8b05d02632fd1f Mon Sep 17 00:00:00 2001
From: Hudson Xing <1277646412@qq.com>
Date: Wed, 31 Dec 2025 12:00:24 +0800
Subject: [PATCH] Move log_prefill_stats_late to correct location in PP mode
 (#15946)

---
 python/sglang/srt/disaggregation/decode.py    |  3 +--
 python/sglang/srt/disaggregation/prefill.py   |  6 +-----
 python/sglang/srt/managers/scheduler.py       | 13 ++++++------
 .../srt/managers/scheduler_dp_attn_mixin.py   | 21 +++++++++++++++++++
 .../sglang/srt/managers/scheduler_pp_mixin.py |  3 +--
 5 files changed, 31 insertions(+), 15 deletions(-)

diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py
index 7726f6b64..199885244 100644
--- a/python/sglang/srt/disaggregation/decode.py
+++ b/python/sglang/srt/disaggregation/decode.py
@@ -937,8 +937,7 @@ class SchedulerDisaggregationDecodeMixin:
         # 2. decode + prebuilt -> decode + idle (idle forward, prebuilt returns)
         # 3. prebuilt + None -> None (None forward, prebuilt returns) + None
         # 4. prebuilt + decode + None -> idle (idle forward, prebuilt returns) + decode + idle
-        if self.require_mlp_sync:
-            ret = self.prepare_mlp_sync_batch(ret)
+        ret = self.maybe_prepare_mlp_sync_batch_and_log_stats(ret)
 
         if ret:
             trace_event_batch("schedule", ret.reqs)
diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py
index 2af278205..ac11013f8 100644
--- a/python/sglang/srt/disaggregation/prefill.py
+++ b/python/sglang/srt/disaggregation/prefill.py
@@ -325,14 +325,11 @@ class SchedulerDisaggregationPrefillMixin:
         self.process_prefill_chunk()
 
         batch = self.get_new_batch_prefill()
-        if self.require_mlp_sync:
-            batch = self.prepare_mlp_sync_batch(batch)
+        batch = self.maybe_prepare_mlp_sync_batch_and_log_stats(batch)
 
         if batch:
             trace_event_batch("schedule", batch.reqs)
 
-        self.log_prefill_stats_late(batch)
-
         return batch
 
     @torch.no_grad()
@@ -523,7 +520,6 @@ class SchedulerDisaggregationPrefillMixin:
                 )
 
         self.maybe_send_health_check_signal()
-        self.log_prefill_stats_late(batch)
 
     def process_disagg_prefill_inflight_queue(
         self: Scheduler, rids_to_check: Optional[List[str]] = None
diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py
index 28a423612..d99038135 100644
--- a/python/sglang/srt/managers/scheduler.py
+++ b/python/sglang/srt/managers/scheduler.py
@@ -1835,7 +1835,9 @@ class Scheduler(
             # Before merging the new batch into running batch:
             # 1. All new batches are none -> need_mlp_sync remains true (sync is needed for decode batch).
             # 2. All new batches are some (prefill / idle) -> we do not need prepare mlp sync one more time.
-            new_batch = self.prepare_mlp_sync_batch(new_batch)
+            new_batch = self.maybe_prepare_mlp_sync_batch_and_log_stats(
+                new_batch, log_stats=False
+            )
             need_mlp_sync = new_batch is None
 
         if new_batch is not None:
@@ -1849,15 +1851,14 @@ class Scheduler(
             else:
                 ret = None
 
-        # Handle DP attention
-        if need_mlp_sync:
-            ret = self.prepare_mlp_sync_batch(ret)
+        # Handle DP attention and log stats
+        ret = self.maybe_prepare_mlp_sync_batch_and_log_stats(
+            ret, need_sync=need_mlp_sync
+        )
 
         if ret:
             trace_event_batch("schedule", ret.reqs)
 
-        self.log_prefill_stats_late(ret)
-
         return ret
 
     def get_num_allocatable_reqs(self, running_bs):
diff --git a/python/sglang/srt/managers/scheduler_dp_attn_mixin.py b/python/sglang/srt/managers/scheduler_dp_attn_mixin.py
index 9e8600c88..9c92cd9c3 100644
--- a/python/sglang/srt/managers/scheduler_dp_attn_mixin.py
+++ b/python/sglang/srt/managers/scheduler_dp_attn_mixin.py
@@ -207,6 +207,27 @@ class SchedulerDPAttnMixin:
             offload_tags=self.offload_tags,
         )
 
+    def maybe_prepare_mlp_sync_batch_and_log_stats(
+        self: Scheduler,
+        batch: Optional[ScheduleBatch],
+        need_sync: Optional[bool] = None,
+        log_stats: bool = True,
+    ) -> Optional[ScheduleBatch]:
+        """
+        Helper to pair log_prefill_stats with log_prefill_stats_late.
+        Should be called after get_new_batch_prefill() to ensure proper pairing.
+
+        Args:
+            batch: The batch to process
+            need_sync: If specified, overrides self.require_mlp_sync for prepare_mlp_sync_batch decision
+            log_stats: Whether to call log_prefill_stats_late. Set to False for intermediate calls.
+        """
+        if need_sync if need_sync is not None else self.require_mlp_sync:
+            batch = self.prepare_mlp_sync_batch(batch)
+        if log_stats:
+            self.log_prefill_stats_late(batch)
+        return batch
+
     def get_idle_batch(self: Scheduler) -> ScheduleBatch:
         idle_batch = ScheduleBatch.init_new(
             [],
diff --git a/python/sglang/srt/managers/scheduler_pp_mixin.py b/python/sglang/srt/managers/scheduler_pp_mixin.py
index 945c4e497..a1a58fe38 100644
--- a/python/sglang/srt/managers/scheduler_pp_mixin.py
+++ b/python/sglang/srt/managers/scheduler_pp_mixin.py
@@ -219,8 +219,7 @@ class SchedulerPPMixin:
 
                 self.process_prefill_chunk()
                 batch = self.get_new_batch_prefill()
-                if self.require_mlp_sync:
-                    batch = self.prepare_mlp_sync_batch(batch)
+                batch = self.maybe_prepare_mlp_sync_batch_and_log_stats(batch)
                 self.mbs[mb_id] = batch
                 self.running_mbs[mb_id] = self.running_batch