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)