From 93fca0bbc3135f473441611fc900c90a2f09fe20 Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Tue, 10 Feb 2026 15:54:03 -0800 Subject: [PATCH] Fix wrong prefill log. (#18570) --- python/sglang/srt/managers/schedule_batch.py | 3 ++ python/sglang/srt/managers/scheduler.py | 10 ++++ .../srt/managers/scheduler_metrics_mixin.py | 46 +++++++++++-------- .../scheduler_output_processor_mixin.py | 10 +--- 4 files changed, 42 insertions(+), 27 deletions(-) diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 062e9a576..519235db6 100644 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -92,6 +92,7 @@ if TYPE_CHECKING: from typing import Any, Dict from sglang.srt.configs.model_config import ModelConfig + from sglang.srt.managers.scheduler_metrics_mixin import PrefillStats from sglang.srt.speculative.eagle_info import EagleDraftInput from sglang.srt.speculative.spec_info import SpecInput, SpeculativeAlgorithm @@ -1304,6 +1305,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): # Metrics dp_cooperation_info: Optional[DPCooperationInfo] = None + prefill_stats: Optional[PrefillStats] = None @classmethod def init_new( @@ -2243,6 +2245,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): mamba_track_mask=self.mamba_track_mask, mamba_track_seqlens=self.mamba_track_seqlens, dp_cooperation_info=self.dp_cooperation_info, + prefill_stats=self.prefill_stats, ) def maybe_evict_swa(self): diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 3747df78a..d0cce5875 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -151,6 +151,7 @@ from sglang.srt.managers.scheduler_dp_attn_mixin import SchedulerDPAttnMixin from sglang.srt.managers.scheduler_input_blocker import SchedulerInputBlocker from sglang.srt.managers.scheduler_metrics_mixin import ( RECORD_STEP_TIME, + PrefillStats, SchedulerMetricsMixin, ) from sglang.srt.managers.scheduler_output_processor_mixin import ( @@ -2121,6 +2122,15 @@ class Scheduler( new_batch.prepare_for_extend() + # Record prefill stats for logging after forward + new_batch.prefill_stats = PrefillStats( + log_input_tokens=adder.log_input_tokens, + log_hit_tokens=adder.log_hit_tokens, + new_token_ratio=adder.new_token_ratio, + running_bs=len(self.running_batch.reqs), + num_new_seqs=len(can_run_list), + ) + # Mixed-style chunked prefill if ( self.is_mixed_chunk diff --git a/python/sglang/srt/managers/scheduler_metrics_mixin.py b/python/sglang/srt/managers/scheduler_metrics_mixin.py index b73f81d75..78f1ba31c 100644 --- a/python/sglang/srt/managers/scheduler_metrics_mixin.py +++ b/python/sglang/srt/managers/scheduler_metrics_mixin.py @@ -1,10 +1,11 @@ from __future__ import annotations +import dataclasses import logging import time from collections import defaultdict from contextlib import contextmanager -from typing import TYPE_CHECKING, Dict, List, Optional, Union +from typing import TYPE_CHECKING, Dict, Optional, Union from sglang.srt.disaggregation.kv_events import EventPublisherFactory, KVEventBatch from sglang.srt.disaggregation.utils import DisaggregationMode @@ -20,8 +21,7 @@ from sglang.srt.managers.io_struct import ( QueueMetrics, SpeculativeMetrics, ) -from sglang.srt.managers.schedule_policy import PrefillAdder -from sglang.srt.managers.scheduler import Req, ScheduleBatch +from sglang.srt.managers.scheduler import ScheduleBatch from sglang.srt.managers.utils import GenerationBatchResult from sglang.srt.metrics.collector import ( SchedulerMetricsCollector, @@ -42,6 +42,17 @@ LOG_FORWARD_ITERS = envs.SGLANG_LOG_FORWARD_ITERS.get() ENABLE_METRICS_DEVICE_TIMER = envs.SGLANG_ENABLE_METRICS_DEVICE_TIMER.get() +@dataclasses.dataclass +class PrefillStats: + """Stats for logging prefill batch metrics.""" + + log_input_tokens: int + log_hit_tokens: int + new_token_ratio: float + running_bs: int + num_new_seqs: int # len(can_run_list) + + class KvMetrics: def __init__(self): self.request_active_slots = None @@ -148,21 +159,18 @@ class SchedulerMetricsMixin: def log_prefill_stats( self: Scheduler, - adder: PrefillAdder, - can_run_list: List[Req], - running_bs: int, - running_bs_offline_batch: int, + prefill_stats: PrefillStats, can_run_cuda_graph: bool, ): 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 = adder.log_input_tokens + 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=adder.log_input_tokens, - adder_log_hit_tokens=adder.log_hit_tokens, + 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 @@ -204,16 +212,16 @@ class SchedulerMetricsMixin: num_used, token_usage, _, _ = self._get_token_info() token_usage_msg = f"token usage: {token_usage:.2f}, " - self.stats.new_token_ratio = adder.new_token_ratio + self.stats.new_token_ratio = prefill_stats.new_token_ratio iter_msg = f" [{self.forward_ct + 1}]" if LOG_FORWARD_ITERS else "" msg = ( f"Prefill batch{iter_msg}, " - f"#new-seq: {len(can_run_list)}, " - f"#new-token: {adder.log_input_tokens}, " - f"#cached-token: {adder.log_hit_tokens}, " + f"#new-seq: {prefill_stats.num_new_seqs}, " + f"#new-token: {prefill_stats.log_input_tokens}, " + f"#cached-token: {prefill_stats.log_hit_tokens}, " f"{token_usage_msg}" - f"#running-req: {running_bs}, " + f"#running-req: {prefill_stats.running_bs}, " f"#queue-req: {len(self.waiting_queue)}, " ) @@ -240,13 +248,13 @@ class SchedulerMetricsMixin: if self.enable_metrics: # Basics - total_tokens = adder.log_input_tokens + adder.log_hit_tokens + total_tokens = prefill_stats.log_input_tokens + prefill_stats.log_hit_tokens cache_hit_rate = ( - adder.log_hit_tokens / total_tokens if total_tokens > 0 else 0.0 + prefill_stats.log_hit_tokens / total_tokens if total_tokens > 0 else 0.0 ) - self.stats.num_running_reqs = running_bs - self.stats.num_running_reqs_offline_batch = running_bs_offline_batch + self.stats.num_running_reqs = prefill_stats.running_bs + self.stats.num_running_reqs_offline_batch = 0 self.stats.num_used_tokens = num_used self.stats.token_usage = token_usage if self.is_hybrid_swa: diff --git a/python/sglang/srt/managers/scheduler_output_processor_mixin.py b/python/sglang/srt/managers/scheduler_output_processor_mixin.py index ae2459708..5baaa5aed 100644 --- a/python/sglang/srt/managers/scheduler_output_processor_mixin.py +++ b/python/sglang/srt/managers/scheduler_output_processor_mixin.py @@ -343,10 +343,7 @@ class SchedulerOutputProcessorMixin: if self.current_scheduler_metrics_enabled: can_run_cuda_graph = getattr(result, "can_run_cuda_graph", False) self.log_prefill_stats( - adder=self.adder, - can_run_list=self.can_run_list, - running_bs=self.running_bs, - running_bs_offline_batch=0, + prefill_stats=batch.prefill_stats, can_run_cuda_graph=can_run_cuda_graph, ) @@ -422,10 +419,7 @@ class SchedulerOutputProcessorMixin: if self.current_scheduler_metrics_enabled: can_run_cuda_graph = getattr(result, "can_run_cuda_graph", False) self.log_prefill_stats( - adder=self.adder, - can_run_list=self.can_run_list, - running_bs=self.running_bs, - running_bs_offline_batch=0, + prefill_stats=batch.prefill_stats, can_run_cuda_graph=can_run_cuda_graph, )