From 6ffc74efd7d23627c0245efb61af7f62cde4a970 Mon Sep 17 00:00:00 2001 From: Yilong Zhao <74357408+happierpig@users.noreply.github.com> Date: Fri, 6 Mar 2026 17:41:27 -0800 Subject: [PATCH] [Metrics] Add overlap bubble timing, full KV usage gauge, and prefill cuda graph tracking (#19982) --- python/sglang/srt/managers/scheduler.py | 6 +- .../scheduler_output_processor_mixin.py | 8 +- .../srt/observability/metrics_collector.py | 33 +++- .../observability/scheduler_metrics_mixin.py | 183 +++++++++++------- python/sglang/srt/utils/device_timer.py | 29 +++ 5 files changed, 181 insertions(+), 78 deletions(-) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index c18539f19..a7d3c4550 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -1158,6 +1158,7 @@ class Scheduler( recv_reqs = self.recv_requests() self.process_input_requests(recv_reqs) if self._engine_paused: + self.cancel_bubble_timer() continue # Get the next batch to run @@ -1212,6 +1213,7 @@ class Scheduler( self.result_queue.append((batch.copy(), batch_result)) else: batch_result = None + self.cancel_bubble_timer() # Process the last batch if self.last_batch: @@ -2363,7 +2365,7 @@ class Scheduler( bs = len(model_worker_batch.seq_lens) future_indices = self.future_map.alloc_future_indices(bs) - with self.forward_stream_ctx: + with self.forward_stream_ctx, self.record_bubble_metrics(batch): self.forward_stream.wait_stream(self.schedule_stream) self.future_map.resolve_future(model_worker_batch) with self.record_forward_metrics(batch): @@ -2439,7 +2441,7 @@ class Scheduler( if self.enable_overlap: self.record_batch_in_overlap(model_worker_batch) - with self.forward_stream_ctx: + with self.forward_stream_ctx, self.record_bubble_metrics(batch): self.forward_stream.wait_stream(self.schedule_stream) embeddings = self.tp_worker.forward_batch_embedding( model_worker_batch diff --git a/python/sglang/srt/managers/scheduler_output_processor_mixin.py b/python/sglang/srt/managers/scheduler_output_processor_mixin.py index 6d64ffd9a..f3b2d01ee 100644 --- a/python/sglang/srt/managers/scheduler_output_processor_mixin.py +++ b/python/sglang/srt/managers/scheduler_output_processor_mixin.py @@ -308,6 +308,10 @@ class SchedulerOutputProcessorMixin: if self.current_scheduler_metrics_enabled: can_run_cuda_graph = getattr(result, "can_run_cuda_graph", False) + if self.enable_metrics: + self.metrics_collector.increment_prefill_cuda_graph_pass( + value=can_run_cuda_graph + ) self.log_prefill_stats( prefill_stats=batch.prefill_stats, can_run_cuda_graph=can_run_cuda_graph, @@ -379,7 +383,9 @@ class SchedulerOutputProcessorMixin: if not batch.spec_algorithm.is_none(): self.update_spec_metrics(batch.batch_size(), result.num_accepted_tokens) if self.enable_metrics: - self.metrics_collector.increment_cuda_graph_pass(value=can_run_cuda_graph) + self.metrics_collector.increment_decode_cuda_graph_pass( + value=can_run_cuda_graph + ) self.token_to_kv_pool_allocator.free_group_begin() diff --git a/python/sglang/srt/observability/metrics_collector.py b/python/sglang/srt/observability/metrics_collector.py index 2ff61f694..16da400d5 100644 --- a/python/sglang/srt/observability/metrics_collector.py +++ b/python/sglang/srt/observability/metrics_collector.py @@ -80,6 +80,7 @@ class SchedulerStats: num_running_reqs: QueueCount = field(default_factory=QueueCount) num_used_tokens: int = 0 token_usage: float = 0.0 + full_token_usage: float = 0.0 pending_prealloc_token_usage: float = 0.0 swa_token_usage: float = 0.0 mamba_usage: float = 0.0 @@ -202,6 +203,12 @@ class SchedulerMetricsCollector: labelnames=labels.keys(), multiprocess_mode="mostrecent", ) + self.full_token_usage = Gauge( + name="sglang:full_token_usage", + documentation="The token usage for full attention layers.", + labelnames=labels.keys(), + multiprocess_mode="mostrecent", + ) self.pending_prealloc_token_usage = Gauge( name="sglang:pending_prealloc_token_usage", documentation="The token usage for pending preallocated tokens (not preallocated yet).", @@ -687,6 +694,14 @@ class SchedulerMetricsCollector: ), labelnames=list(labels.keys()) + ["category"], ) + self.gpu_overlap_wait_seconds_total = Counter( + name="sglang:gpu_overlap_wait_seconds_total", + documentation=( + "Total time that GPU forward stream was idle waiting for " + "the CPU schedule stream (overlap bubble)." + ), + labelnames=list(labels.keys()) + ["category"], + ) self.dp_cooperation_realtime_tokens_total = Counter( name="sglang:dp_cooperation_realtime_tokens_total", @@ -851,11 +866,14 @@ class SchedulerMetricsCollector: num_retracted_output_tokens ) - def increment_cuda_graph_pass(self, value: bool) -> None: - # leave room for piecewise cuda graph, etc + def increment_decode_cuda_graph_pass(self, value: bool) -> None: mode = "decode_cuda_graph" if value else "decode_none" self.cuda_graph_passes_total.labels(**self.labels, mode=mode).inc(1) + def increment_prefill_cuda_graph_pass(self, value: bool) -> None: + mode = "prefill_cuda_graph" if value else "prefill_none" + self.cuda_graph_passes_total.labels(**self.labels, mode=mode).inc(1) + def increment_eplb_balancedness( self, forward_mode: str, balancedness: float ) -> None: @@ -885,6 +903,16 @@ class SchedulerMetricsCollector: **dp_cooperation_info.to_labels(), ).inc(delta) + def increment_gpu_overlap_wait_seconds( + self, + category: str, + t: float, + dp_cooperation_info: Optional[DPCooperationInfo], + ): + self.gpu_overlap_wait_seconds_total.labels( + **self.labels, category=category + ).inc(t) + def increment_gpu_execution_seconds( self, category: str, @@ -904,6 +932,7 @@ class SchedulerMetricsCollector: self._log_gauge_queue_count(self.num_running_reqs, stats.num_running_reqs) self._log_gauge(self.num_used_tokens, stats.num_used_tokens) self._log_gauge(self.token_usage, stats.token_usage) + self._log_gauge(self.full_token_usage, stats.full_token_usage) self._log_gauge( self.pending_prealloc_token_usage, stats.pending_prealloc_token_usage ) diff --git a/python/sglang/srt/observability/scheduler_metrics_mixin.py b/python/sglang/srt/observability/scheduler_metrics_mixin.py index 170753aa2..33ce9cc9a 100644 --- a/python/sglang/srt/observability/scheduler_metrics_mixin.py +++ b/python/sglang/srt/observability/scheduler_metrics_mixin.py @@ -31,7 +31,7 @@ from sglang.srt.observability.metrics_collector import ( compute_routing_key_stats, ) from sglang.srt.utils import get_bool_env_var -from sglang.srt.utils.device_timer import DeviceTimer +from sglang.srt.utils.device_timer import DeviceTimer, GapTimer from sglang.srt.utils.scheduler_status_logger import SchedulerStatusLogger if TYPE_CHECKING: @@ -117,7 +117,6 @@ class SchedulerMetricsMixin: self.current_scheduler_metrics_enabled = ( self.attn_tp_rank == 0 or self.enable_metrics_for_all_schedulers ) - if self.enable_metrics: if self.server_args.disaggregation_mode == DisaggregationMode.PREFILL.value: engine_type = "prefill" @@ -152,6 +151,9 @@ class SchedulerMetricsMixin: self.forward_pass_device_timer = DeviceTimer( reporter=self.metrics_collector.increment_gpu_execution_seconds, ) + self.bubble_timer = GapTimer( + reporter=self.metrics_collector.increment_gpu_overlap_wait_seconds, + ) if self.enable_kv_cache_events: self.init_kv_events(self.server_args.kv_events_config) @@ -189,43 +191,46 @@ class SchedulerMetricsMixin: self.last_prefill_tokens = prefill_stats.log_input_tokens # TODO: generalize this for various memory pools + msg_parts = [] + num_used = token_usage = full_token_usage = None + if self.is_hybrid_swa: - ( - full_num_used, - swa_num_used, - full_token_usage, - swa_token_usage, - _, - _, - _, - _, - ) = self._get_swa_token_info() + full_num_used, swa_num_used, full_tok, swa_token_usage, *_ = ( + self._get_swa_token_info() + ) num_used = max(full_num_used, swa_num_used) - token_usage = max(full_token_usage, swa_token_usage) - token_usage_msg = ( - f"full token usage: {full_token_usage:.2f}, " - f"swa token usage: {swa_token_usage:.2f}, " + token_usage = max(full_tok, swa_token_usage) + full_token_usage = full_tok + msg_parts += [ + f"full token usage: {full_tok:.2f}", + f"swa token usage: {swa_token_usage:.2f}", + ] + + if self.is_hybrid_ssm: + num_used_m, _, full_tok_m, mamba_usage, *_ = self._get_mamba_token_info() + num_used = max(num_used, num_used_m) if num_used is not None else num_used_m + token_usage = ( + max(token_usage, mamba_usage) + if token_usage is not None + else max(full_tok_m, mamba_usage) ) - elif self.is_hybrid_ssm: - ( - full_num_used, - _, - full_token_usage, - mamba_usage, - _, - _, - _, - _, - ) = self._get_mamba_token_info() - num_used = full_num_used - token_usage = full_token_usage - token_usage_msg = ( - f"full token usage: {full_token_usage:.2f}, " - f"mamba usage: {mamba_usage:.2f}, " - ) - else: - num_used, token_usage, _, _ = self._get_token_info() - token_usage_msg = f"token usage: {token_usage:.2f}, " + if full_token_usage is None: + full_token_usage = full_tok_m + msg_parts.append(f"full token usage: {full_tok_m:.2f}") + msg_parts.append(f"mamba usage: {mamba_usage:.2f}") + + if full_token_usage is None: + num_used, tok, _, _ = self._get_token_info() + full_token_usage = tok + token_usage = tok + msg_parts.append(f"token usage: {tok:.2f}") + + assert ( + num_used is not None + and token_usage is not None + and full_token_usage is not None + ) + token_usage_msg = ", ".join(msg_parts) + ", " self.stats.new_token_ratio = prefill_stats.new_token_ratio iter_msg = f" [{self.forward_ct + 1}]" if LOG_FORWARD_ITERS else "" @@ -281,6 +286,7 @@ class SchedulerMetricsMixin: self.stats.num_running_reqs_offline_batch = 0 self.stats.num_used_tokens = num_used self.stats.token_usage = token_usage + self.stats.full_token_usage = full_token_usage if self.is_hybrid_swa: self.stats.swa_token_usage = swa_token_usage if self.is_hybrid_ssm: @@ -340,47 +346,56 @@ class SchedulerMetricsMixin: num_running_reqs_offline_batch = 0 # TODO: generalize this for various memory pools + msg_parts = [] + num_used = token_usage = full_token_usage = None + if self.is_hybrid_swa: - ( - full_num_used, - swa_num_used, - full_token_usage, - swa_token_usage, - _, - _, - _, - _, - ) = self._get_swa_token_info() + full_num_used, swa_num_used, full_tok, swa_token_usage, *_ = ( + self._get_swa_token_info() + ) num_used = max(full_num_used, swa_num_used) - token_usage = max(full_token_usage, swa_token_usage) - token_usage_msg = ( - f"#full token: {full_num_used}, " - f"full token usage: {full_token_usage:.2f}, " - f"#swa token: {swa_num_used}, " - f"swa token usage: {swa_token_usage:.2f}, " + token_usage = max(full_tok, swa_token_usage) + full_token_usage = full_tok + msg_parts += [ + f"#full token: {full_num_used}", + f"full token usage: {full_tok:.2f}", + f"#swa token: {swa_num_used}", + f"swa token usage: {swa_token_usage:.2f}", + ] + + if self.is_hybrid_ssm: + num_used_m, mamba_num, full_tok_m, mamba_usage, *_ = ( + self._get_mamba_token_info() ) - elif self.is_hybrid_ssm: - ( - full_num_used, - mamba_used, - full_token_usage, - mamba_usage, - _, - _, - _, - _, - ) = self._get_mamba_token_info() - num_used = full_num_used - token_usage = full_token_usage - token_usage_msg = ( - f"#full token: {full_num_used}, " - f"full token usage: {full_token_usage:.2f}, " - f"mamba num: {mamba_used}, " - f"mamba usage: {mamba_usage:.2f}, " + num_used = max(num_used, num_used_m) if num_used is not None else num_used_m + token_usage = ( + max(token_usage, mamba_usage) + if token_usage is not None + else max(full_tok_m, mamba_usage) ) - else: - num_used, token_usage, _, _ = self._get_token_info() - token_usage_msg = f"#token: {num_used}, token usage: {token_usage:.2f}, " + if full_token_usage is None: + full_token_usage = full_tok_m + msg_parts += [ + f"#full token: {num_used_m}", + f"full token usage: {full_tok_m:.2f}", + ] + msg_parts += [ + f"mamba num: {mamba_num}", + f"mamba usage: {mamba_usage:.2f}", + ] + + if full_token_usage is None: + num_used, tok, _, _ = self._get_token_info() + full_token_usage = tok + token_usage = tok + msg_parts.append(f"#token: {num_used}, token usage: {tok:.2f}") + + assert ( + num_used is not None + and token_usage is not None + and full_token_usage is not None + ) + token_usage_msg = ", ".join(msg_parts) + ", " if RECORD_STEP_TIME: self.step_time_dict[num_running_reqs].append( @@ -449,7 +464,10 @@ class SchedulerMetricsMixin: ) self.stats.num_running_reqs_offline_batch = num_running_reqs_offline_batch self.stats.num_used_tokens = num_used + # maximum usage of all pools self.stats.token_usage = token_usage + # usage of full attention + self.stats.full_token_usage = full_token_usage if self.is_hybrid_swa: self.stats.swa_token_usage = swa_token_usage if self.is_hybrid_ssm: @@ -816,3 +834,22 @@ class SchedulerMetricsMixin: ), ): yield + + @contextmanager + def record_bubble_metrics(self: Scheduler, batch: ScheduleBatch): + if not (self.enable_metrics and ENABLE_METRICS_DEVICE_TIMER): + yield + return + + category = "forward_" + batch.forward_mode.name.lower() + with self.bubble_timer.wrap( + metadata=dict( + category=category, + dp_cooperation_info=batch.dp_cooperation_info, + ), + ): + yield + + def cancel_bubble_timer(self: Scheduler): + if self.enable_metrics and ENABLE_METRICS_DEVICE_TIMER: + self.bubble_timer.cancel() diff --git a/python/sglang/srt/utils/device_timer.py b/python/sglang/srt/utils/device_timer.py index 686c8d3d9..4cd5817ad 100644 --- a/python/sglang/srt/utils/device_timer.py +++ b/python/sglang/srt/utils/device_timer.py @@ -30,6 +30,35 @@ class DeviceTimer: self._reporter(t=interval.elapsed_time() / 1000.0, **interval.metadata) +class GapTimer(DeviceTimer): + """Measures GPU idle gaps between consecutive uses of a stream. + + Where DeviceTimer.wrap() measures the duration *inside* a block, + GapTimer.wrap() measures the time *between* consecutive blocks + (gap = next_block_start - last_block_end). + """ + + def __init__(self, reporter: Callable): + super().__init__(reporter) + self._pending: Optional[_TimingInterval] = None + + @contextmanager + def wrap(self, metadata: Dict): + if self._pending is not None: + self._pending.end(metadata=metadata) + self._intervals.append(self._pending) + self._pending = None + self._report() + try: + yield + finally: + self._pending = _TimingInterval.create() + + def cancel(self): + """Discard a pending gap (e.g. server went idle).""" + self._pending = None + + @dataclass class _TimingInterval: start_event: torch.cuda.Event