chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Fix wrong prefill log. (#18570)

Browse Source
This commit is contained in:
Liangsheng Yin
2026-02-10 15:54:03 -08:00
committed by GitHub
parent 2bfab1bb67
commit 93fca0bbc3
4 changed files with 42 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
python/sglang/srt/managers/schedule_batch.py
View File

@@ -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):

10
python/sglang/srt/managers/scheduler.py
View File

@@ -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

46
python/sglang/srt/managers/scheduler_metrics_mixin.py
View File

@@ -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:

10
python/sglang/srt/managers/scheduler_output_processor_mixin.py
View File

@@ -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,
)