From 9a7d6be56775307bada7cadc5513cea733feccf7 Mon Sep 17 00:00:00 2001 From: Ratish P <114130421+Ratish1@users.noreply.github.com> Date: Wed, 18 Feb 2026 04:05:35 +0530 Subject: [PATCH] cleanup prefill metrics logging to fix dp-attn metrics (#18778) Co-authored-by: Liangsheng Yin --- python/sglang/srt/disaggregation/decode.py | 2 +- python/sglang/srt/disaggregation/prefill.py | 2 +- python/sglang/srt/managers/scheduler.py | 8 ++--- .../srt/managers/scheduler_dp_attn_mixin.py | 10 ++---- .../srt/managers/scheduler_metrics_mixin.py | 32 ++++++------------- .../scheduler_output_processor_mixin.py | 2 ++ .../sglang/srt/managers/scheduler_pp_mixin.py | 2 +- python/sglang/srt/metrics/collector.py | 6 +++- 8 files changed, 24 insertions(+), 40 deletions(-) diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index aa69eaabe..fa653448b 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/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) diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index eadd04e21..fbc801635 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/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) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 2edec253f..3435fcaef 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/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, log_stats=False - ) + new_batch = self.maybe_prepare_mlp_sync_batch(new_batch) 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, need_sync=need_mlp_sync - ) + ret = self.maybe_prepare_mlp_sync_batch(ret, need_sync=need_mlp_sync) if ret: trace_event_batch("schedule", ret.reqs) diff --git a/python/sglang/srt/managers/scheduler_dp_attn_mixin.py b/python/sglang/srt/managers/scheduler_dp_attn_mixin.py index 656b6e7c2..12ef400db 100644 --- a/python/sglang/srt/managers/scheduler_dp_attn_mixin.py +++ b/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. - Should be called after get_new_batch_prefill() to ensure proper pairing. + Helper to prepare MLP sync batch for DP attention. + 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: diff --git a/python/sglang/srt/managers/scheduler_metrics_mixin.py b/python/sglang/srt/managers/scheduler_metrics_mixin.py index ad1330efc..30b2732b9 100644 --- a/python/sglang/srt/managers/scheduler_metrics_mixin.py +++ b/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 ): diff --git a/python/sglang/srt/managers/scheduler_output_processor_mixin.py b/python/sglang/srt/managers/scheduler_output_processor_mixin.py index 8c22a9a71..482bc6ca6 100644 --- a/python/sglang/srt/managers/scheduler_output_processor_mixin.py +++ b/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( diff --git a/python/sglang/srt/managers/scheduler_pp_mixin.py b/python/sglang/srt/managers/scheduler_pp_mixin.py index 3d199b230..1a65a3c3d 100644 --- a/python/sglang/srt/managers/scheduler_pp_mixin.py +++ b/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 diff --git a/python/sglang/srt/metrics/collector.py b/python/sglang/srt/metrics/collector.py index ac6570803..255d41ccc 100644 --- a/python/sglang/srt/metrics/collector.py +++ b/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(