From ecab3f322e5e2b52c929a5aaf962b3bb4b7a8cb4 Mon Sep 17 00:00:00 2001 From: Lianmin Zheng Date: Tue, 23 Jun 2026 15:46:22 -0700 Subject: [PATCH] Revert "Improve MFU metrics for prefill and verify timing" (#29079) --- python/sglang/srt/managers/schedule_batch.py | 5 --- .../scheduler_components/metrics_reporter.py | 37 ++++--------------- .../sglang/srt/model_executor/model_runner.py | 7 +--- .../srt/model_executor/runner/eager_runner.py | 7 +--- 4 files changed, 10 insertions(+), 46 deletions(-) diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 7ee1a1728..f1dc81d17 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -2822,11 +2822,6 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): # merge_batch) on the original don't corrupt this snapshot. return ScheduleBatch( reqs=self.reqs[:], - # Per-request extend/prefix lens, snapshotted like reqs so the - # deferred prefill-stats report reads stable values after the - # original batch has moved on. - 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 295ca3949..d7bfcbc02 100644 --- a/python/sglang/srt/managers/scheduler_components/metrics_reporter.py +++ b/python/sglang/srt/managers/scheduler_components/metrics_reporter.py @@ -435,13 +435,8 @@ class SchedulerMetricsReporter: num_attn_heads * head_dim * act_bytes * num_layers ) - def _estimate_prefill_perf( - self, batch: Optional[ScheduleBatch] - ) -> 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))) + def _estimate_prefill_perf(self, num_tokens: int) -> Tuple[float, float, float]: + tokens = max(0, int(num_tokens)) if tokens == 0: return 0.0, 0.0, 0.0 @@ -489,16 +484,6 @@ class SchedulerMetricsReporter: ) return flops, read_bytes, write_bytes - def _prefill_sol_suffix( - self, batch: Optional[ScheduleBatch], elapsed_s: float - ) -> str: - """Hook for model-specific speed-of-light metrics on prefill log lines.""" - return "" - - def _decode_sol_suffix(self, batch: ScheduleBatch, elapsed_s: float) -> str: - """Hook for model-specific speed-of-light metrics on decode log lines.""" - return "" - def reset_metrics(self): self.forward_ct_decode = 0 self.num_generated_tokens = 0 @@ -568,13 +553,9 @@ class SchedulerMetricsReporter: msg += f"input throughput (token/s): {self.last_input_throughput:.2f}" if self.enable_mfu_metrics and gap_latency > 0: - 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}" + 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}" if ENABLE_METRICS_DEVICE_TIMER: msg += f", fwd occupancy: {self.fwd_occupancy:.2f}%" @@ -591,7 +572,9 @@ class SchedulerMetricsReporter: dp_cooperation_info=dp_cooperation_info, ) if self.enable_mfu_metrics: - flops, read_bytes, write_bytes = self._estimate_prefill_perf(batch) + flops, read_bytes, write_bytes = self._estimate_prefill_perf( + prefill_stats.log_input_tokens + ) self.metrics_collector.increment_estimated_perf( num_flops_per_gpu=flops, num_read_bytes_per_gpu=read_bytes, @@ -782,10 +765,6 @@ 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 eed4279b8..55d4bf42a 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -3051,16 +3051,11 @@ class ModelRunner(ModelRunnerKVCacheMixin): ): # Prefill cuda graph (piecewise). kwargs = self._extend_forward_kwargs(forward_batch, pp_proxy_tensors) - category = ( - "target_verify" - if forward_batch.forward_mode.is_target_verify() - else "extend" - ) # 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": category}) + self.device_timer.wrap(metadata={"category": "extend"}) 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 0034cceca..8a61636ec 100644 --- a/python/sglang/srt/model_executor/runner/eager_runner.py +++ b/python/sglang/srt/model_executor/runner/eager_runner.py @@ -250,11 +250,6 @@ class EagerRunner(BaseRunner): ) -> Union[LogitsProcessorOutput, PPProxyTensors, EmbeddingPoolerOutput]: model_runner = self.model_runner kwargs = model_runner._extend_forward_kwargs(forward_batch, pp_proxy_tensors) - category = ( - "target_verify" - if forward_batch.forward_mode.is_target_verify() - else "extend" - ) if not model_runner.server_args.enable_pdmux: forward_batch = self.load_batch(forward_batch, pp_proxy_tensors) @@ -283,7 +278,7 @@ class EagerRunner(BaseRunner): forward_positions = sharded_positions ctx = ( - model_runner.device_timer.wrap(metadata={"category": category}) + model_runner.device_timer.wrap(metadata={"category": "extend"}) if model_runner.device_timer else contextlib.nullcontext() )