[Metrics] Add overlap bubble timing, full KV usage gauge, and prefill cuda graph tracking (#19982)

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

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

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:
+    def increment_decode_cuda_graph_pass(self, value: bool) -> None:
-        # leave room for piecewise cuda graph, etc
         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
         )

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_tok, swa_token_usage, *_ = (
-                full_num_used,
+                self._get_swa_token_info()
-                swa_num_used,
+            )
-                full_token_usage,
-                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 = max(full_tok, swa_token_usage)
-            token_usage_msg = (
+            full_token_usage = full_tok
-                f"full token usage: {full_token_usage:.2f}, "
+            msg_parts += [
-                f"swa token usage: {swa_token_usage:.2f}, "
+                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:
+            if full_token_usage is None:
-            (
+                full_token_usage = full_tok_m
-                full_num_used,
+                msg_parts.append(f"full token usage: {full_tok_m:.2f}")
-                _,
+            msg_parts.append(f"mamba usage: {mamba_usage:.2f}")
-                full_token_usage,
+
-                mamba_usage,
+        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}")
-            ) = self._get_mamba_token_info()
+
-            num_used = full_num_used
+        assert (
-            token_usage = full_token_usage
+            num_used is not None
-            token_usage_msg = (
+            and token_usage is not None
-                f"full token usage: {full_token_usage:.2f}, "
+            and full_token_usage is not None
-                f"mamba usage: {mamba_usage:.2f}, "
         )
-        else:
+        token_usage_msg = ", ".join(msg_parts) + ", "
-            num_used, token_usage, _, _ = self._get_token_info()
-            token_usage_msg = f"token usage: {token_usage:.2f}, "
 
         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_tok, swa_token_usage, *_ = (
-                full_num_used,
+                self._get_swa_token_info()
-                swa_num_used,
+            )
-                full_token_usage,
-                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 = max(full_tok, swa_token_usage)
-            token_usage_msg = (
+            full_token_usage = full_tok
-                f"#full token: {full_num_used}, "
+            msg_parts += [
-                f"full token usage: {full_token_usage:.2f}, "
+                f"#full token: {full_num_used}",
-                f"#swa token: {swa_num_used}, "
+                f"full token usage: {full_tok:.2f}",
-                f"swa token usage: {swa_token_usage:.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:
+            num_used = max(num_used, num_used_m) if num_used is not None else num_used_m
-            (
+            token_usage = (
-                full_num_used,
+                max(token_usage, mamba_usage)
-                mamba_used,
+                if token_usage is not None
-                full_token_usage,
+                else max(full_tok_m, mamba_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}, "
             )
-        else:
+            if full_token_usage is None:
-            num_used, token_usage, _, _ = self._get_token_info()
+                full_token_usage = full_tok_m
-            token_usage_msg = f"#token: {num_used}, token usage: {token_usage:.2f}, "
+                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()

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