Improve MFU metrics for prefill and verify timing (#29000)

Co-authored-by: Pranjal Shankhdhar <pranjal.ssh@gmail.com>

python/sglang/srt/managers/schedule_batch.py

@@ -2822,6 +2822,11 @@ 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,8 +435,13 @@ 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: 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)))
         if tokens == 0:
             return 0.0, 0.0, 0.0
 
@@ -484,6 +489,16 @@ 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
@@ -553,9 +568,13 @@ 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}"
+            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 +591,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 +782,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

python/sglang/srt/model_executor/model_runner.py

@@ -3047,11 +3047,16 @@ 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": "extend"})
+                    self.device_timer.wrap(metadata={"category": category})
                     if self.device_timer
                     else contextlib.nullcontext()
                 )

python/sglang/srt/model_executor/runner/eager_runner.py

@@ -309,6 +309,11 @@ 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)
@@ -337,7 +342,7 @@ class EagerRunner(BaseRunner):
             forward_positions = sharded_positions
 
         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()
         )