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

Add metrics for having prefill and decode in different ranks (#15752)

Browse Source
This commit is contained in:
fzyzcjy
2025-12-24 21:35:35 +08:00
committed by GitHub
parent b3b818fd86
commit fd4a558e71
9 changed files with 238 additions and 64 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
python/sglang/srt/environ.py
View File

@@ -386,6 +386,7 @@ class Envs:
# Metrics
SGLANG_ENABLE_METRICS_DEVICE_TIMER = EnvBool(False)
SGLANG_ENABLE_METRICS_DP_ATTENTION = EnvBool(False)
# fmt: on

10
python/sglang/srt/managers/schedule_batch.py
View File

@@ -74,7 +74,11 @@ from sglang.srt.mem_cache.mamba_radix_cache import MambaRadixCache
from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
from sglang.srt.mem_cache.radix_cache import RadixKey
from sglang.srt.mem_cache.swa_radix_cache import SWARadixCache
from sglang.srt.metrics.collector import SchedulerMetricsCollector, TimeStats
from sglang.srt.metrics.collector import (
DPCooperationInfo,
SchedulerMetricsCollector,
TimeStats,
)
from sglang.srt.model_executor.forward_batch_info import (
CaptureHiddenMode,
ForwardBatch,
@@ -1249,6 +1253,9 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
# Diffusion LLM
dllm_config: Optional[DllmConfig] = None
# Metrics
dp_cooperation_info: Optional[DPCooperationInfo] = None
@classmethod
def init_new(
cls,
@@ -2161,6 +2168,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
mamba_track_indices=self.mamba_track_indices,
mamba_track_mask=self.mamba_track_mask,
mamba_track_seqlens=self.mamba_track_seqlens,
dp_cooperation_info=self.dp_cooperation_info,
)
def _is_available_size_sufficient(self, num_tokens: int) -> bool:

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

@@ -716,6 +716,7 @@ class Scheduler(
self.last_batch: Optional[ScheduleBatch] = None
self.forward_ct = 0
self.last_prefill_tokens = 0
self.last_prefill_cache_tokens = 0
self.return_health_check_ct = 0
self.num_retracted_reqs: int = 0
self.num_paused_reqs: int = 0
@@ -1830,6 +1831,8 @@ class Scheduler(
if ret:
trace_event_batch("schedule", ret.reqs)
self.log_prefill_stats_late(ret)
return ret
def get_num_allocatable_reqs(self, running_bs):

12
python/sglang/srt/managers/scheduler_dp_attn_mixin.py
View File

@@ -1,13 +1,14 @@
from __future__ import annotations
from dataclasses import dataclass
from typing import TYPE_CHECKING, Callable
from typing import TYPE_CHECKING, Callable, Optional
import torch
from sglang.srt.batch_overlap.two_batch_overlap import TboDPAttentionPreparer
from sglang.srt.environ import envs
from sglang.srt.managers.schedule_batch import ScheduleBatch
from sglang.srt.metrics.collector import DPCooperationInfo
from sglang.srt.utils.common import require_mlp_tp_gather
if TYPE_CHECKING:
@@ -15,6 +16,9 @@ if TYPE_CHECKING:
from sglang.srt.managers.scheduler import Scheduler
_ENABLE_METRICS_DP_ATTENTION = envs.SGLANG_ENABLE_METRICS_DP_ATTENTION.get()
@dataclass
class MLPSyncBatchInfo:
dp_size: int
@@ -33,6 +37,7 @@ class MLPSyncBatchInfo:
global_num_tokens_for_logprob: list[int] = None
tbo_split_seq_index: torch.Tensor = None
global_forward_mode: int = None
dp_cooperation_info: Optional[DPCooperationInfo] = None
def _get_local_tensor(self, device, dtype=torch.int64) -> torch.Tensor:
return torch.tensor(
@@ -68,6 +73,8 @@ class MLPSyncBatchInfo:
self.global_num_tokens_for_logprob = tp0_info[:, 1].tolist()
self.can_cuda_graph = bool(tp0_info[:, 2].min().item())
self.is_extend_in_batch = bool(tp0_info[:, 3].max().item())
if _ENABLE_METRICS_DP_ATTENTION:
self.dp_cooperation_info = DPCooperationInfo.create(tp0_info[:, 5].tolist())
def _update_gather_batch(
@@ -180,6 +187,9 @@ def prepare_mlp_sync_batch_raw(
batch_to_gather, mlp_sync_info, require_mlp_tp_gather, skip_all_gather
)
if _ENABLE_METRICS_DP_ATTENTION and local_batch is not None:
local_batch.dp_cooperation_info = mlp_sync_info.dp_cooperation_info
return local_batch

38
python/sglang/srt/managers/scheduler_metrics_mixin.py
View File

@@ -88,7 +88,7 @@ class SchedulerMetricsMixin:
if ENABLE_METRICS_DEVICE_TIMER:
self.forward_pass_device_timer = DeviceTimer(
reporter=self.metrics_collector.increment_gpu_execution_seconds
reporter=self.metrics_collector.increment_gpu_execution_seconds,
)
if self.enable_kv_cache_events:
@@ -124,6 +124,7 @@ class SchedulerMetricsMixin:
self.last_prefill_stats_tic = time.perf_counter()
self.last_input_throughput = self.last_prefill_tokens / gap_latency
self.last_prefill_tokens = adder.log_input_tokens
self.last_prefill_cache_tokens = adder.log_hit_tokens
# TODO: generalize this for various memory pools
if self.is_hybrid_swa:
@@ -231,23 +232,26 @@ class SchedulerMetricsMixin:
self.disagg_decode_transfer_queue.queue
)
self.metrics_collector.increment_realtime_tokens(
prefill_compute_tokens=adder.log_input_tokens,
prefill_cache_tokens=adder.log_hit_tokens,
)
# Others
self.calculate_utilization()
self.metrics_collector.log_stats(self.stats)
self._emit_kv_metrics()
self._publish_kv_events()
def log_prefill_stats_late(self: Scheduler, batch: Optional[ScheduleBatch]):
"""This should be called after `batch` has gathered enough metadata."""
if self.enable_metrics and batch is not None:
self.metrics_collector.increment_realtime_tokens(
prefill_compute_tokens=self.last_prefill_tokens,
prefill_cache_tokens=self.last_prefill_cache_tokens,
dp_cooperation_info=batch.dp_cooperation_info,
)
def log_decode_stats(
self: Scheduler, can_run_cuda_graph: bool, running_batch: ScheduleBatch = None
):
batch = running_batch or self.running_batch
last_num_generated_tokens = self.num_generated_tokens
gap_latency = time.perf_counter() - self.last_decode_stats_tic
self.last_decode_stats_tic = time.perf_counter()
self.last_gen_throughput = self.num_generated_tokens / gap_latency
@@ -388,16 +392,21 @@ class SchedulerMetricsMixin:
self.disagg_decode_transfer_queue.queue
)
self.metrics_collector.increment_realtime_tokens(
decode_tokens=last_num_generated_tokens
)
# Others
self.calculate_utilization()
self.metrics_collector.log_stats(self.stats)
self._emit_kv_metrics()
self._publish_kv_events()
def log_decode_stats_every_iteration(
self: Scheduler, batch: ScheduleBatch, num_accepted_tokens: int
):
self.metrics_collector.increment_realtime_tokens(
# TODO unify this w/ the bumping logic in `Scheduler.num_generated_tokens` accumulator
decode_tokens=batch.batch_size() + num_accepted_tokens,
dp_cooperation_info=batch.dp_cooperation_info,
)
def log_batch_result_stats(
self: Scheduler,
batch: ScheduleBatch,
@@ -491,5 +500,10 @@ class SchedulerMetricsMixin:
return
category = "forward_" + batch.forward_mode.name.lower()
with self.forward_pass_device_timer.wrap(category=category):
with self.forward_pass_device_timer.wrap(
metadata=dict(
category=category,
dp_cooperation_info=batch.dp_cooperation_info,
),
):
yield

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

@@ -445,6 +445,10 @@ class SchedulerOutputProcessorMixin:
and self.forward_ct_decode % self.server_args.decode_log_interval == 0
):
self.log_decode_stats(can_run_cuda_graph, running_batch=batch)
if self.enable_metrics:
self.log_decode_stats_every_iteration(
batch, num_accepted_tokens=result.num_accepted_tokens
)
def _mamba_prefix_cache_update(
self, req: Req, batch: ScheduleBatch, result: GenerationBatchResult, i: int

56
python/sglang/srt/metrics/collector.py
View File

@@ -12,6 +12,7 @@
# limitations under the License.
# ==============================================================================
"""Utilities for Prometheus Metrics Collection."""
import dataclasses
import logging
import os
import time
@@ -21,6 +22,7 @@ from typing import Dict, List, Optional, Union
from sglang.srt.disaggregation.utils import DisaggregationMode
from sglang.srt.environ import envs
from sglang.srt.metrics.utils import exponential_buckets, generate_buckets
from sglang.srt.model_executor.forward_batch_info import ForwardMode
from sglang.srt.server_args import ServerArgs
from sglang.srt.utils import get_bool_env_var
@@ -240,6 +242,24 @@ class SchedulerStats:
is_cuda_graph: float = 0.0
@dataclass
class DPCooperationInfo:
# Users can derive that, except for cases with idle, num_decode_ranks=world_size-num_prefill_ranks
# We do not provide `num_decode_ranks` to avoid cardinality explosion.
num_prefill_ranks: int
@staticmethod
def create(forward_modes: List[int]):
return DPCooperationInfo(
num_prefill_ranks=sum(
1 for mode in forward_modes if mode == ForwardMode.EXTEND.value
),
)
def to_labels(self):
return dataclasses.asdict(self)
class SchedulerMetricsCollector:
def __init__(
@@ -680,6 +700,17 @@ class SchedulerMetricsCollector:
labelnames=list(labels.keys()) + ["category"],
)
self.dp_cooperation_realtime_tokens_total = Counter(
name="sglang:dp_cooperation_realtime_tokens_total",
documentation="Total number of tokens processed with labels about DP cooperation.",
labelnames=list(labels.keys()) + ["mode", "num_prefill_ranks"],
)
self.dp_cooperation_gpu_execution_seconds_total = Counter(
name="sglang:dp_cooperation_gpu_execution_seconds_total",
documentation="Total time that GPU is busy executing a workload with labels about DP cooperation.",
labelnames=list(labels.keys()) + ["category", "num_prefill_ranks"],
)
def _log_gauge(self, gauge, data: Union[int, float]) -> None:
# Convenience function for logging to gauge.
gauge.labels(**self.labels).set(data)
@@ -716,7 +747,11 @@ class SchedulerMetricsCollector:
)
def increment_realtime_tokens(
self, prefill_compute_tokens=0, prefill_cache_tokens=0, decode_tokens=0
self,
dp_cooperation_info: Optional[DPCooperationInfo],
prefill_compute_tokens=0,
prefill_cache_tokens=0,
decode_tokens=0,
):
for mode, delta in [
("prefill_compute", prefill_compute_tokens),
@@ -724,10 +759,27 @@ class SchedulerMetricsCollector:
("decode", decode_tokens),
]:
self.realtime_tokens_total.labels(**self.labels, mode=mode).inc(delta)
if dp_cooperation_info is not None:
self.dp_cooperation_realtime_tokens_total.labels(
**self.labels,
mode=mode,
**dp_cooperation_info.to_labels(),
).inc(delta)
def increment_gpu_execution_seconds(self, category: str, t: float):
def increment_gpu_execution_seconds(
self,
category: str,
t: float,
dp_cooperation_info: Optional[DPCooperationInfo],
):
logger.debug(f"GPU execution seconds: {category=} {t=:.3f}")
self.gpu_execution_seconds_total.labels(**self.labels, category=category).inc(t)
if dp_cooperation_info is not None:
self.dp_cooperation_gpu_execution_seconds_total.labels(
**self.labels,
category=category,
**dp_cooperation_info.to_labels(),
).inc(t)
def log_stats(self, stats: SchedulerStats) -> None:
self._log_gauge(self.num_running_reqs, stats.num_running_reqs)

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

@@ -1,23 +1,23 @@
from collections import deque
from contextlib import contextmanager
from dataclasses import dataclass
from typing import Callable, Deque, Optional
from typing import Callable, Deque, Dict, Optional
import torch
class DeviceTimer:
def __init__(self, reporter: Callable[[str, float], None]):
def __init__(self, reporter: Callable):
self._intervals: Deque[_TimingInterval] = deque()
self._reporter = reporter
@contextmanager
def wrap(self, category: str):
def wrap(self, metadata: Dict):
self._intervals.append(_TimingInterval.create())
try:
yield
finally:
self._intervals[-1].end(category=category)
self._intervals[-1].end(metadata=metadata)
self._report()
def _report(self):
@@ -27,14 +27,14 @@ class DeviceTimer:
break
self._intervals.popleft()
self._reporter(interval.category, interval.elapsed_time() / 1000.0)
self._reporter(t=interval.elapsed_time() / 1000.0, **interval.metadata)
@dataclass
class _TimingInterval:
start_event: torch.cuda.Event
end_event: Optional[torch.cuda.Event] = None
category: Optional[str] = None
metadata: Optional[Dict] = None
@staticmethod
def create():
@@ -42,13 +42,13 @@ class _TimingInterval:
start_event.record()
return _TimingInterval(start_event=start_event)
def end(self, category: str):
def end(self, metadata: Dict):
end_event = torch.cuda.Event(enable_timing=True)
end_event.record()
assert self.end_event is None
self.end_event = end_event
self.category = category
self.metadata = metadata
def elapsed_time(self) -> float:
return self.start_event.elapsed_time(self.end_event)

162
test/srt/test_metrics.py
View File

@@ -1,27 +1,79 @@
import unittest
from typing import Dict, List
import requests
from prometheus_client.parser import text_string_to_metric_families
from prometheus_client.samples import Sample
from sglang.srt.environ import envs
from sglang.srt.utils import kill_process_tree
from sglang.test.test_utils import (
DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
CustomTestCase,
is_in_ci,
popen_launch_server,
)
_MODEL_NAME = "Qwen/Qwen3-0.6B"
class TestEnableMetrics(CustomTestCase):
def test_metrics_enabled(self):
def test_metrics_1gpu(self):
"""Test that metrics endpoint returns data when enabled"""
process = popen_launch_server(
DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
DEFAULT_URL_FOR_TEST,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=["--enable-metrics", "--cuda-graph-max-bs", 2],
self._execute_core(
other_args=[],
verify_metrics_extra=None,
)
def test_metrics_2gpu(self):
# TODO enable when we have 2-gpu runner in nightly CI
if is_in_ci():
print("Skip test_metrics_2gpu since in 1-gpu CI")
return
def _verify_metrics_extra(metrics):
metrics_to_check = [
(
"sglang:dp_cooperation_realtime_tokens_total",
{"mode": "prefill_compute"},
),
("sglang:dp_cooperation_realtime_tokens_total", {"mode": "decode"}),
(
"sglang:dp_cooperation_gpu_execution_seconds_total",
{"category": "forward_prefill"},
),
(
"sglang:dp_cooperation_gpu_execution_seconds_total",
{"category": "forward_decode"},
),
]
_check_metrics_positive(self, metrics, metrics_to_check)
num_prefill_ranks_values = {
s.labels["num_prefill_ranks"]
for s in metrics["sglang:dp_cooperation_realtime_tokens_total"]
}
self.assertIn("0", num_prefill_ranks_values)
self.assertIn("1", num_prefill_ranks_values)
self._execute_core(
other_args=["--tp", "2", "--dp", "2", "--enable-dp-attention"],
verify_metrics_extra=_verify_metrics_extra,
)
def _execute_core(self, other_args, verify_metrics_extra):
with (
envs.SGLANG_ENABLE_METRICS_DP_ATTENTION.override(True),
envs.SGLANG_ENABLE_METRICS_DEVICE_TIMER.override(True),
):
process = popen_launch_server(
_MODEL_NAME,
DEFAULT_URL_FOR_TEST,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=["--enable-metrics", "--cuda-graph-max-bs", 2, *other_args],
)
try:
# Make some requests to generate some metrics
response = requests.get(f"{DEFAULT_URL_FOR_TEST}/health_generate")
@@ -45,44 +97,74 @@ class TestEnableMetrics(CustomTestCase):
# Get metrics
metrics_response = requests.get(f"{DEFAULT_URL_FOR_TEST}/metrics")
self.assertEqual(metrics_response.status_code, 200)
metrics_content = metrics_response.text
metrics_text = metrics_response.text
print(f"metrics_content=\n{metrics_content}")
# Verify essential metrics are present
essential_metrics = [
"sglang:num_running_reqs",
"sglang:num_used_tokens",
"sglang:token_usage",
"sglang:gen_throughput",
"sglang:num_queue_reqs",
"sglang:num_grammar_queue_reqs",
"sglang:cache_hit_rate",
"sglang:spec_accept_length",
"sglang:prompt_tokens_total",
"sglang:generation_tokens_total",
"sglang:cached_tokens_total",
"sglang:num_requests_total",
"sglang:time_to_first_token_seconds",
"sglang:inter_token_latency_seconds",
"sglang:e2e_request_latency_seconds",
]
for metric in essential_metrics:
self.assertIn(metric, metrics_content, f"Missing metric: {metric}")
# Verify model name label is present and correct
expected_model_name = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
self.assertIn(f'model_name="{expected_model_name}"', metrics_content)
# Verify metrics have values (not empty)
self.assertIn("_sum{", metrics_content)
self.assertIn("_count{", metrics_content)
self.assertIn("_bucket{", metrics_content)
print(f"metrics_text=\n{metrics_text}")
metrics = _parse_prometheus_metrics(metrics_text)
self._verify_metrics_common(metrics_text, metrics)
if verify_metrics_extra is not None:
verify_metrics_extra(metrics)
finally:
kill_process_tree(process.pid)
def _verify_metrics_common(self, metrics_text, metrics):
essential_metrics = [
"sglang:num_running_reqs",
"sglang:num_used_tokens",
"sglang:token_usage",
"sglang:gen_throughput",
"sglang:num_queue_reqs",
"sglang:num_grammar_queue_reqs",
"sglang:cache_hit_rate",
"sglang:spec_accept_length",
"sglang:prompt_tokens_total",
"sglang:generation_tokens_total",
"sglang:cached_tokens_total",
"sglang:num_requests_total",
"sglang:time_to_first_token_seconds",
"sglang:inter_token_latency_seconds",
"sglang:e2e_request_latency_seconds",
]
for metric in essential_metrics:
self.assertIn(metric, metrics_text, f"Missing metric: {metric}")
self.assertIn(f'model_name="{_MODEL_NAME}"', metrics_text)
self.assertIn("_sum{", metrics_text)
self.assertIn("_count{", metrics_text)
self.assertIn("_bucket{", metrics_text)
metrics_to_check = [
("sglang:realtime_tokens_total", {"mode": "prefill_compute"}),
("sglang:realtime_tokens_total", {"mode": "decode"}),
("sglang:gpu_execution_seconds_total", {"category": "forward_extend"}),
("sglang:gpu_execution_seconds_total", {"category": "forward_decode"}),
]
_check_metrics_positive(self, metrics, metrics_to_check)
def _parse_prometheus_metrics(metrics_text: str) -> Dict[str, List[Sample]]:
result = {}
for family in text_string_to_metric_families(metrics_text):
for sample in family.samples:
if sample.name not in result:
result[sample.name] = []
result[sample.name].append(sample)
return result
def _get_sample_value_by_labels(samples: List[Sample], labels: Dict[str, str]) -> float:
for sample in samples:
if all(sample.labels.get(k) == v for k, v in labels.items()):
return sample.value
raise KeyError(f"No sample found with labels {labels}")
def _check_metrics_positive(test_case, metrics, metrics_to_check):
for metric_name, labels in metrics_to_check:
value = _get_sample_value_by_labels(metrics[metric_name], labels)
test_case.assertGreater(value, 0, f"{metric_name} {labels}")
if __name__ == "__main__":
unittest.main()