From 7e63feee6f37480b62938e5c002676d777eb4115 Mon Sep 17 00:00:00 2001 From: Lianmin Zheng Date: Wed, 24 Jun 2026 15:50:05 -0700 Subject: [PATCH] Add scheduler metrics extension hooks (#29207) Co-authored-by: Yinghai Lu --- python/sglang/srt/managers/schedule_batch.py | 6 +++ .../scheduler_components/metrics_reporter.py | 42 +++++++++++++++---- .../sglang/srt/model_executor/model_runner.py | 7 +++- .../srt/model_executor/runner/eager_runner.py | 7 +++- 4 files changed, 52 insertions(+), 10 deletions(-) diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index f1dc81d17..2fad05156 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -2822,6 +2822,12 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): # merge_batch) on the original don't corrupt this snapshot. return ScheduleBatch( reqs=self.reqs[:], + # Per-request extend/prefix lens, snapshotted (sliced like reqs) so the + # deferred prefill-stats report reads them after the original batch has + # moved on. prepare_for_extend sets these; mix_with_running mutates them + # in place. None for decode batches (no extend), which the reader skips. + extend_lens=self.extend_lens[:] if self.extend_lens is not None else None, + prefix_lens=self.prefix_lens[:] if self.prefix_lens is not None else None, req_to_token_pool=self.req_to_token_pool, req_pool_indices=self.req_pool_indices, model_config=self.model_config, diff --git a/python/sglang/srt/managers/scheduler_components/metrics_reporter.py b/python/sglang/srt/managers/scheduler_components/metrics_reporter.py index d7bfcbc02..dc2a332dd 100644 --- a/python/sglang/srt/managers/scheduler_components/metrics_reporter.py +++ b/python/sglang/srt/managers/scheduler_components/metrics_reporter.py @@ -435,8 +435,10 @@ class SchedulerMetricsReporter: num_attn_heads * head_dim * act_bytes * num_layers ) - def _estimate_prefill_perf(self, num_tokens: int) -> Tuple[float, float, float]: - tokens = max(0, int(num_tokens)) + def _estimate_prefill_perf(self, batch) -> Tuple[float, float, float]: + if batch is None or batch.extend_lens is None: + return 0.0, 0.0, 0.0 + tokens = max(0, int(sum(batch.extend_lens))) if tokens == 0: return 0.0, 0.0, 0.0 @@ -484,6 +486,19 @@ class SchedulerMetricsReporter: ) return flops, read_bytes, write_bytes + def _prefill_sol_suffix(self, batch, elapsed_s: float) -> str: + """Hook: model-specific speed-of-light % suffix for the prefill log line. + ``batch`` carries the per-request extend/prefix lengths a subclass needs + for an exact attention pair-count. No model arch here, so returns ""; + a subclass may override it.""" + return "" + + def _decode_sol_suffix(self, batch, elapsed_s: float) -> str: + """Hook: model-specific speed-of-light % suffix for the decode log line. + ``elapsed_s`` is per-iteration. No model arch here, so returns ""; + a subclass may override it.""" + return "" + def reset_metrics(self): self.forward_ct_decode = 0 self.num_generated_tokens = 0 @@ -553,9 +568,18 @@ class SchedulerMetricsReporter: msg += f"input throughput (token/s): {self.last_input_throughput:.2f}" if self.enable_mfu_metrics and gap_latency > 0: - flops, _, _ = self._estimate_prefill_perf(prefill_stats.log_input_tokens) - tflops_per_s = flops / gap_latency / 1e12 - msg += f", est. prefill TFLOPS/s (per GPU): {tflops_per_s:.2f}" + # Prefer the SoL suffix when it carries content: it scores FLOPs against + # each forward's actual GPU span (device timer). The wall-clock est. + # TFLOPS below divides FLOPs by gap_latency -- the inter-log interval on + # the async scheduler loop, which is decoupled from this forward's + # execution -- so it disagrees with the SoL. Omit it when SoL is present. + sol_suffix = self._prefill_sol_suffix(batch, gap_latency) + if sol_suffix: + msg += sol_suffix + else: + flops, _, _ = self._estimate_prefill_perf(batch) + tflops_per_s = flops / gap_latency / 1e12 + msg += f", est. prefill TFLOPS/s (per GPU): {tflops_per_s:.2f}" if ENABLE_METRICS_DEVICE_TIMER: msg += f", fwd occupancy: {self.fwd_occupancy:.2f}%" @@ -572,9 +596,7 @@ class SchedulerMetricsReporter: dp_cooperation_info=dp_cooperation_info, ) if self.enable_mfu_metrics: - flops, read_bytes, write_bytes = self._estimate_prefill_perf( - prefill_stats.log_input_tokens - ) + flops, read_bytes, write_bytes = self._estimate_prefill_perf(batch) self.metrics_collector.increment_estimated_perf( num_flops_per_gpu=flops, num_read_bytes_per_gpu=read_bytes, @@ -765,6 +787,10 @@ class SchedulerMetricsReporter: f"est. read BW (GB/s per GPU): {read_gb_per_s:.2f}, " f"est. write BW (GB/s per GPU): {write_gb_per_s:.2f}" ) + msg += self._decode_sol_suffix( + batch, + gap_latency / max(1, self.scheduler.server_args.decode_log_interval), + ) self._mfu_log_flops = 0.0 self._mfu_log_read_bytes = 0.0 self._mfu_log_write_bytes = 0.0 diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index ff964c1f6..3c77948da 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -3052,13 +3052,18 @@ class ModelRunner(ModelRunnerKVCacheMixin): and self.prefill_cuda_graph_runner.can_run_graph(forward_batch) and get_cp_strategy() is None ): + category = ( + "target_verify" + if forward_batch.forward_mode.is_target_verify() + else "extend" + ) # Prefill cuda graph (piecewise). kwargs = self._extend_forward_kwargs(forward_batch, pp_proxy_tensors) # TODO: device_timer.wrap is too broad here — it also includes # load_batch time. Move timing into the prefill cuda graph runner # to capture only the model.forward part. ctx = ( - self.device_timer.wrap(metadata={"category": "extend"}) + self.device_timer.wrap(metadata={"category": category}) if self.device_timer else contextlib.nullcontext() ) diff --git a/python/sglang/srt/model_executor/runner/eager_runner.py b/python/sglang/srt/model_executor/runner/eager_runner.py index 489e5d7d5..8a61fe40f 100644 --- a/python/sglang/srt/model_executor/runner/eager_runner.py +++ b/python/sglang/srt/model_executor/runner/eager_runner.py @@ -282,8 +282,13 @@ class EagerRunner(BaseRunner): kwargs["input_embeds"] = sharded_hidden_states forward_positions = sharded_positions + category = ( + "target_verify" + if forward_batch.forward_mode.is_target_verify() + else "extend" + ) ctx = ( - model_runner.device_timer.wrap(metadata={"category": "extend"}) + model_runner.device_timer.wrap(metadata={"category": category}) if model_runner.device_timer else contextlib.nullcontext() )