diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index f1dc81d17..7ee1a1728 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/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, diff --git a/python/sglang/srt/managers/scheduler_components/metrics_reporter.py b/python/sglang/srt/managers/scheduler_components/metrics_reporter.py index d7bfcbc02..295ca3949 100644 --- a/python/sglang/srt/managers/scheduler_components/metrics_reporter.py +++ b/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 diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 1cff5c983..0bf828dbf 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/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() ) diff --git a/python/sglang/srt/model_executor/runner/eager_runner.py b/python/sglang/srt/model_executor/runner/eager_runner.py index b5788833b..691fd28e5 100644 --- a/python/sglang/srt/model_executor/runner/eager_runner.py +++ b/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() )