Revert "Improve MFU metrics for prefill and verify timing" (#29079)

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,

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(
+    def _estimate_prefill_perf(self, num_tokens: int) -> Tuple[float, float, float]:
-        self, batch: Optional[ScheduleBatch]
+        tokens = max(0, int(num_tokens))
-    ) -> 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
 
@@ -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)
+            flops, _, _ = self._estimate_prefill_perf(prefill_stats.log_input_tokens)
-            if sol_suffix:
+            tflops_per_s = flops / gap_latency / 1e12
-                msg += sol_suffix
+            msg += f", est. prefill TFLOPS/s (per GPU): {tflops_per_s:.2f}"
-            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}%"
@@ -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

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()
                 )

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()
         )