diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index 1a089e8ff..6db94de40 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -572,6 +572,7 @@ class SchedulerDisaggregationPrefillMixin: can_run_cuda_graph = getattr(result, "can_run_cuda_graph", False) self.report_prefill_stats( + batch=batch, prefill_stats=batch.prefill_stats, can_run_cuda_graph=can_run_cuda_graph, dp_cooperation_info=batch.dp_cooperation_info, diff --git a/python/sglang/srt/dllm/mixin/scheduler.py b/python/sglang/srt/dllm/mixin/scheduler.py index 3524d40aa..157ab2192 100644 --- a/python/sglang/srt/dllm/mixin/scheduler.py +++ b/python/sglang/srt/dllm/mixin/scheduler.py @@ -94,6 +94,7 @@ class SchedulerDllmMixin: can_run_cuda_graph = getattr(result, "can_run_cuda_graph", False) self.report_prefill_stats( + batch=batch, prefill_stats=batch.prefill_stats, can_run_cuda_graph=can_run_cuda_graph, dp_cooperation_info=batch.dp_cooperation_info, diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 5aa8bd298..40e2e0dd7 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -1500,6 +1500,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): # Metrics dp_cooperation_info: Optional[DPCooperationInfo] = None prefill_stats: Optional[PrefillStats] = None + forward_iter: Optional[int] = None # HiSparse hisparse_coordinator: Optional[HiSparseCoordinator] = None @@ -2625,6 +2626,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): mamba_track_seqlens=self.mamba_track_seqlens, dp_cooperation_info=self.dp_cooperation_info, prefill_stats=self.prefill_stats, + forward_iter=self.forward_iter, ) def maybe_evict_swa(self): diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 39917179d..9c694321c 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -2949,6 +2949,7 @@ class Scheduler( ) -> Union[GenerationBatchResult, EmbeddingBatchResult]: """Run a batch.""" self.forward_ct += 1 + batch.forward_iter = self.forward_ct # Whether to run the profiler self._profile_batch_predicate(batch) diff --git a/python/sglang/srt/managers/scheduler_output_processor_mixin.py b/python/sglang/srt/managers/scheduler_output_processor_mixin.py index bad8f1b22..670c2de89 100644 --- a/python/sglang/srt/managers/scheduler_output_processor_mixin.py +++ b/python/sglang/srt/managers/scheduler_output_processor_mixin.py @@ -368,6 +368,7 @@ class SchedulerOutputProcessorMixin: can_run_cuda_graph = getattr(result, "can_run_cuda_graph", False) self.report_prefill_stats( + batch=batch, prefill_stats=batch.prefill_stats, can_run_cuda_graph=can_run_cuda_graph, dp_cooperation_info=batch.dp_cooperation_info, diff --git a/python/sglang/srt/observability/scheduler_metrics_mixin.py b/python/sglang/srt/observability/scheduler_metrics_mixin.py index 06b1596b8..5703434dd 100644 --- a/python/sglang/srt/observability/scheduler_metrics_mixin.py +++ b/python/sglang/srt/observability/scheduler_metrics_mixin.py @@ -352,6 +352,7 @@ class SchedulerMetricsMixin: def report_prefill_stats( self: Scheduler, + batch: Optional[ScheduleBatch], prefill_stats: PrefillStats, can_run_cuda_graph: bool, dp_cooperation_info: Optional[DPCooperationInfo] = None, @@ -373,7 +374,12 @@ class SchedulerMetricsMixin: token_usage_msg = ", ".join(pool_stats.get_prefill_usage_msg_parts()) + ", " self.stats.new_token_ratio = prefill_stats.new_token_ratio - iter_msg = f" [{self.forward_ct + 1}]" if LOG_FORWARD_ITERS else "" + batch_iter = ( + batch.forward_iter + if batch is not None and batch.forward_iter is not None + else self.forward_ct + ) + iter_msg = f" [{batch_iter}]" if LOG_FORWARD_ITERS else "" msg = ( f"Prefill batch{iter_msg}, " @@ -533,7 +539,12 @@ class SchedulerMetricsMixin: gap_latency / self.server_args.decode_log_interval ) - iter_msg = f" [{self.forward_ct}]" if LOG_FORWARD_ITERS else "" + batch_iter = ( + batch.forward_iter + if batch is not None and batch.forward_iter is not None + else self.forward_ct + ) + iter_msg = f" [{batch_iter}]" if LOG_FORWARD_ITERS else "" msg = f"Decode batch{iter_msg}, #running-req: {num_running_reqs}, {token_usage_msg}" if self.spec_algorithm.is_none():