Add scheduler metrics extension hooks (#29207)

Co-authored-by: Yinghai Lu <yinghai@meta.com>

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,

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

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

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