chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Yilong Zhao
2026-03-06 17:41:27 -08:00
committed by GitHub
parent 50bbdcf8e9
commit 6ffc74efd7
5 changed files with 181 additions and 78 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
python/sglang/srt/managers/scheduler.py
View File

@@ -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

8
python/sglang/srt/managers/scheduler_output_processor_mixin.py
View File

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

33
python/sglang/srt/observability/metrics_collector.py
View File

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

183
python/sglang/srt/observability/scheduler_metrics_mixin.py
View File

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

29
python/sglang/srt/utils/device_timer.py
View File

@@ -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