Fix wrong prefill log. (sgl-project#18570)

Author: hnyls2002

python/sglang/srt/managers/schedule_batch.py

@@ -92,6 +92,7 @@
     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 @@ def copy(self):
             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):

python/sglang/srt/managers/scheduler.py

@@ -151,6 +151,7 @@
 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 @@ def _get_new_batch_prefill_raw(
 
         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

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 @@
     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 @@
 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 @@ def reset_metrics(self: Scheduler):
 
     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 @@ def log_prefill_stats(
             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 @@ def log_prefill_stats(
 
         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:

python/sglang/srt/managers/scheduler_output_processor_mixin.py

@@ -343,10 +343,7 @@ def process_batch_result_prefill(
         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 @@ def process_batch_result_dllm(
         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,
             )