cleanup prefill metrics logging to fix dp-attn metrics (#18778)

Co-authored-by: Liangsheng Yin <lsyincs@gmail.com>

python/sglang/srt/disaggregation/decode.py

@@ -1000,7 +1000,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
-        ret = self.maybe_prepare_mlp_sync_batch_and_log_stats(ret)
+        ret = self.maybe_prepare_mlp_sync_batch(ret)
 
         if ret:
             trace_event_batch("schedule", ret.reqs)

python/sglang/srt/disaggregation/prefill.py

@@ -351,7 +351,7 @@ class SchedulerDisaggregationPrefillMixin:
         self.process_prefill_chunk()
 
         batch = self.get_new_batch_prefill()
-        batch = self.maybe_prepare_mlp_sync_batch_and_log_stats(batch)
+        batch = self.maybe_prepare_mlp_sync_batch(batch)
 
         if batch:
             trace_event_batch("schedule", batch.reqs)

python/sglang/srt/managers/scheduler.py

@@ -1917,9 +1917,7 @@ 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.maybe_prepare_mlp_sync_batch_and_log_stats(
+            new_batch = self.maybe_prepare_mlp_sync_batch(new_batch)
-                new_batch, log_stats=False
-            )
             need_mlp_sync = new_batch is None
 
         if new_batch is not None:
@@ -1934,9 +1932,7 @@ class Scheduler(
                 ret = None
 
         # Handle DP attention and log stats
-        ret = self.maybe_prepare_mlp_sync_batch_and_log_stats(
+        ret = self.maybe_prepare_mlp_sync_batch(ret, need_sync=need_mlp_sync)
-            ret, need_sync=need_mlp_sync
-        )
 
         if ret:
             trace_event_batch("schedule", ret.reqs)

python/sglang/srt/managers/scheduler_dp_attn_mixin.py

@@ -238,25 +238,21 @@ class SchedulerDPAttnMixin:
             offload_tags=self.offload_tags,
         )
 
-    def maybe_prepare_mlp_sync_batch_and_log_stats(
+    def maybe_prepare_mlp_sync_batch(
         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.
+        Helper to prepare MLP sync batch for DP attention.
-        Should be called after get_new_batch_prefill() to ensure proper pairing.
+        Should be called after get_new_batch_prefill().
 
         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:

python/sglang/srt/managers/scheduler_metrics_mixin.py

@@ -5,7 +5,7 @@ import logging
 import time
 from collections import defaultdict
 from contextlib import contextmanager
-from typing import TYPE_CHECKING, Dict, Optional, Union
+from typing import TYPE_CHECKING, Optional, Union
 
 from sglang.srt.disaggregation.kv_events import EventPublisherFactory, KVEventBatch
 from sglang.srt.disaggregation.utils import DisaggregationMode
@@ -24,6 +24,7 @@ from sglang.srt.managers.io_struct import (
 from sglang.srt.managers.scheduler import ScheduleBatch
 from sglang.srt.managers.utils import GenerationBatchResult
 from sglang.srt.metrics.collector import (
+    DPCooperationInfo,
     SchedulerMetricsCollector,
     SchedulerStats,
     compute_routing_key_stats,
@@ -93,9 +94,6 @@ class SchedulerMetricsMixin:
         self.kv_transfer_alloc_ms: float = 0.0
         self.kv_transfer_total_mb: float = 0.0
 
-        # Only for `log_prefill_stats` to pass information to `log_prefill_stats_late`
-        self.temp_prefill_info: Optional[Dict] = None
-
         self.stats = SchedulerStats()
 
         # Metrics
@@ -163,18 +161,13 @@ class SchedulerMetricsMixin:
         self: Scheduler,
         prefill_stats: PrefillStats,
         can_run_cuda_graph: bool,
+        dp_cooperation_info: Optional[DPCooperationInfo] = None,
     ):
         gap_latency = time.perf_counter() - self.last_prefill_stats_tic
         self.last_prefill_stats_tic = time.perf_counter()
         self.last_input_throughput = self.last_prefill_tokens / gap_latency
         self.last_prefill_tokens = prefill_stats.log_input_tokens
 
-        assert self.temp_prefill_info is None
-        self.temp_prefill_info = dict(
-            adder_log_input_tokens=prefill_stats.log_input_tokens,
-            adder_log_hit_tokens=prefill_stats.log_hit_tokens,
-        )
-
         # TODO: generalize this for various memory pools
         if self.is_hybrid_swa:
             (
@@ -249,6 +242,12 @@ class SchedulerMetricsMixin:
         logger.info(msg)
 
         if self.enable_metrics:
+            self.metrics_collector.increment_realtime_tokens(
+                prefill_compute_tokens=prefill_stats.log_input_tokens,
+                prefill_cache_tokens=prefill_stats.log_hit_tokens,
+                dp_cooperation_info=dp_cooperation_info,
+            )
+
             # Basics
             total_tokens = prefill_stats.log_input_tokens + prefill_stats.log_hit_tokens
             cache_hit_rate = (
@@ -302,19 +301,6 @@ class SchedulerMetricsMixin:
             self._emit_kv_metrics()
         self._publish_kv_events()
 
-    def log_prefill_stats_late(self: Scheduler, batch: Optional[ScheduleBatch]):
-        """This should be called after `batch` has gathered enough metadata."""
-
-        info = self.temp_prefill_info
-        self.temp_prefill_info = None
-
-        if self.enable_metrics and batch is not None and info is not None:
-            self.metrics_collector.increment_realtime_tokens(
-                prefill_compute_tokens=info["adder_log_input_tokens"],
-                prefill_cache_tokens=info["adder_log_hit_tokens"],
-                dp_cooperation_info=batch.dp_cooperation_info,
-            )
-
     def log_decode_stats(
         self: Scheduler, can_run_cuda_graph: bool, running_batch: ScheduleBatch = None
     ):

python/sglang/srt/managers/scheduler_output_processor_mixin.py

@@ -331,6 +331,7 @@ class SchedulerOutputProcessorMixin:
             self.log_prefill_stats(
                 prefill_stats=batch.prefill_stats,
                 can_run_cuda_graph=can_run_cuda_graph,
+                dp_cooperation_info=batch.dp_cooperation_info,
             )
 
     def _resolve_spec_overlap_token_ids(
@@ -410,6 +411,7 @@ class SchedulerOutputProcessorMixin:
             self.log_prefill_stats(
                 prefill_stats=batch.prefill_stats,
                 can_run_cuda_graph=can_run_cuda_graph,
+                dp_cooperation_info=batch.dp_cooperation_info,
             )
 
     def process_batch_result_decode(

python/sglang/srt/managers/scheduler_pp_mixin.py

@@ -224,7 +224,7 @@ class SchedulerPPMixin:
 
                 self.process_prefill_chunk()
                 batch = self.get_new_batch_prefill()
-                batch = self.maybe_prepare_mlp_sync_batch_and_log_stats(batch)
+                batch = self.maybe_prepare_mlp_sync_batch(batch)
                 self.mbs[mb_id] = batch
                 self.running_mbs[mb_id] = self.running_batch
 

python/sglang/srt/metrics/collector.py

@@ -277,8 +277,10 @@ class DPCooperationInfo:
     @staticmethod
     def create(forward_modes: List[int]):
         return DPCooperationInfo(
+            # Count ranks that are doing any extend-like work.
+            # With overlap scheduling, prefill can appear as MIXED rather than EXTEND.
             num_prefill_ranks=sum(
-                1 for mode in forward_modes if mode == ForwardMode.EXTEND.value
+                1 for mode in forward_modes if ForwardMode(mode).is_extend()
             ),
         )
 
@@ -945,6 +947,8 @@ class SchedulerMetricsCollector:
             ("prefill_cache", prefill_cache_tokens),
             ("decode", decode_tokens),
         ]:
+            if delta == 0:
+                continue
             self.realtime_tokens_total.labels(**self.labels, mode=mode).inc(delta)
             if dp_cooperation_info is not None:
                 self.dp_cooperation_realtime_tokens_total.labels(