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

feat(metrics): add scheduler and hiradix cache metrics (#10218) (#10225)

Browse Source 
Co-authored-by: Zhiqiang Xie <xiezhq@stanford.edu>
This commit is contained in:
ShawnKung
2025-11-10 18:14:47 +08:00
committed by GitHub
parent 6f08488042
commit afee2843d5
10 changed files with 210 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -413,6 +413,10 @@ class Scheduler(
f"{'available_cpu_mem' if self.device == 'cpu' else 'available_gpu_mem'}={avail_mem:.2f} GB"
)
# Init metrics stats
self.init_metrics(tp_rank, pp_rank, dp_rank)
self.init_kv_events(server_args.kv_events_config)
# Init memory pool and cache
self.init_memory_pool_and_cache()
@@ -508,9 +512,6 @@ class Scheduler(
# Init disaggregation
self.init_disaggregation()
# Init metrics stats
self.init_metrics(tp_rank, pp_rank, dp_rank)
if self.enable_kv_cache_events:
self.init_kv_events(server_args.kv_events_config)
@@ -725,6 +726,7 @@ class Scheduler(
hicache_ratio=server_args.hicache_ratio,
hicache_size=server_args.hicache_size,
hicache_write_policy=server_args.hicache_write_policy,
enable_metrics=self.enable_metrics,
enable_kv_cache_events=self.enable_kv_cache_events,
)
elif self.enable_hierarchical_cache:
@@ -761,6 +763,7 @@ class Scheduler(
page_size=self.page_size,
disable=server_args.disable_radix_cache,
is_eagle=self.spec_algorithm.is_eagle(),
enable_metrics=self.enable_metrics,
)
elif self.is_hybrid_gdn:
self.tree_cache = MambaRadixCache(
@@ -768,6 +771,7 @@ class Scheduler(
token_to_kv_pool_allocator=self.token_to_kv_pool_allocator,
page_size=self.page_size,
disable=server_args.disable_radix_cache,
enable_metrics=self.enable_metrics,
)
elif server_args.enable_lmcache:
from sglang.srt.mem_cache.storage.lmcache.lmc_radix_cache import (
@@ -779,6 +783,7 @@ class Scheduler(
token_to_kv_pool_allocator=self.token_to_kv_pool_allocator,
page_size=self.page_size,
disable=server_args.disable_radix_cache,
enable_metrics=self.enable_metrics,
model_config=self.model_config,
tp_size=self.tp_size,
rank=self.tp_rank,
@@ -791,6 +796,7 @@ class Scheduler(
token_to_kv_pool_allocator=self.token_to_kv_pool_allocator,
page_size=self.page_size,
disable=server_args.disable_radix_cache,
enable_metrics=self.enable_metrics,
enable_kv_cache_events=self.enable_kv_cache_events,
eviction_policy=server_args.radix_eviction_policy,
is_eagle=self.spec_algorithm.is_eagle(),

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

@@ -132,6 +132,7 @@ class SchedulerMetricsMixin:
num_used, token_usage, _, _ = self._get_token_info()
token_usage_msg = f"token usage: {token_usage:.2f}, "
self.stats.new_token_ratio = adder.new_token_ratio
iter_msg = f" [{self.forward_ct + 1}]" if LOG_FORWARD_ITERS else ""
f = (

10
python/sglang/srt/mem_cache/base_prefix_cache.py
View File

@@ -15,6 +15,7 @@ import torch
from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
from sglang.srt.metrics.collector import RadixCacheMetricsCollector
if TYPE_CHECKING:
from sglang.srt.managers.schedule_batch import Req
@@ -50,6 +51,15 @@ class MatchResult(NamedTuple):
class BasePrefixCache(ABC, PrefixCacheTrait):
"""Cache can be indexed by either rid or key."""
metrics_collector: Optional[RadixCacheMetricsCollector] = (
None # metrics collector for the cache
)
def init_metrics_collector(self):
self.metrics_collector = RadixCacheMetricsCollector(
labels={"cache_type": self.__class__.__name__}
)
@abstractmethod
def reset(self):
pass

25
python/sglang/srt/mem_cache/hiradix_cache.py
View File

@@ -117,7 +117,7 @@ class HiRadixCache(RadixCache):
"tp_rank": self.cache_controller.tp_rank,
"dp_rank": self.cache_controller.dp_rank,
}
self.metrics_collector = StorageMetricsCollector(labels=labels)
self.storage_metrics_collector = StorageMetricsCollector(labels=labels)
# record the nodes with ongoing write through
self.ongoing_write_through = {}
@@ -139,6 +139,7 @@ class HiRadixCache(RadixCache):
disable=False,
eviction_policy=eviction_policy,
is_eagle=is_eagle,
enable_metrics=enable_metrics,
)
def _parse_storage_backend_extra_config(
@@ -334,6 +335,7 @@ class HiRadixCache(RadixCache):
return self.evictable_size_
def evict(self, num_tokens: int):
start_time = time.perf_counter()
leaves = self._collect_leaves_device()
eviction_heap = [
(self.eviction_strategy.get_priority(node), node) for node in leaves
@@ -374,6 +376,12 @@ class HiRadixCache(RadixCache):
assert node.backuped
self._evict_backuped(node)
if num_evicted > 0 and self.metrics_collector is not None:
self.metrics_collector.observe_eviction_duration(
time.perf_counter() - start_time
)
self.metrics_collector.increment_eviction_num_tokens(num_evicted)
def _evict_backuped(self, node: TreeNode):
# evict a node already written to host
num_evicted = self.cache_controller.evict_device(node.value)
@@ -425,6 +433,7 @@ class HiRadixCache(RadixCache):
) -> Optional[torch.Tensor]:
# todo: more loading policies
start_time = time.perf_counter()
last_hit_node = node
nodes_to_load = []
while node.evicted:
@@ -469,6 +478,12 @@ class HiRadixCache(RadixCache):
self.evictable_size_ += len(device_indices)
self.inc_lock_ref(last_hit_node)
if self.metrics_collector is not None:
self.metrics_collector.observe_load_back_duration(
time.perf_counter() - start_time
)
self.metrics_collector.increment_load_back_num_tokens(len(device_indices))
return device_indices
def init_load_back(
@@ -507,7 +522,7 @@ class HiRadixCache(RadixCache):
if self.enable_storage:
self.drain_storage_control_queues()
if self.enable_storage_metrics:
self.metrics_collector.log_storage_metrics(
self.storage_metrics_collector.log_storage_metrics(
self.cache_controller.storage_backend.get_stats()
)
@@ -551,7 +566,9 @@ class HiRadixCache(RadixCache):
if entry is not None:
entry.release_host()
if self.enable_storage_metrics:
self.metrics_collector.log_backuped_tokens(operation.completed_tokens)
self.storage_metrics_collector.log_backuped_tokens(
operation.completed_tokens
)
# release host memory
host_indices_list = []
@@ -664,7 +681,7 @@ class HiRadixCache(RadixCache):
self.cache_controller.prefetch_tokens_occupied -= len(token_ids)
if self.enable_storage_metrics:
self.metrics_collector.log_prefetched_tokens(
self.storage_metrics_collector.log_prefetched_tokens(
min_completed_tokens - matched_length
)

4
python/sglang/srt/mem_cache/mamba_radix_cache.py
View File

@@ -326,6 +326,7 @@ class MambaRadixCache(BasePrefixCache):
token_to_kv_pool_allocator: TokenToKVPoolAllocator,
page_size: int,
disable: bool = False,
enable_metrics: bool = False,
):
assert isinstance(token_to_kv_pool_allocator, TokenToKVPoolAllocator)
self.req_to_token_pool = req_to_token_pool
@@ -340,6 +341,9 @@ class MambaRadixCache(BasePrefixCache):
else:
self.device = torch.device("cpu")
if enable_metrics:
self.init_metrics_collector()
self.key_match_fn = _key_match_page_size1
self.get_child_key_fn = get_child_key
self.reset()

11
python/sglang/srt/mem_cache/radix_cache.py
View File

@@ -191,6 +191,7 @@ class RadixCache(BasePrefixCache):
token_to_kv_pool_allocator: BaseTokenToKVPoolAllocator,
page_size: int,
disable: bool = False,
enable_metrics: bool = False,
enable_kv_cache_events: bool = False,
eviction_policy: str = "lru",
is_eagle: bool = False,
@@ -203,6 +204,9 @@ class RadixCache(BasePrefixCache):
self.kv_event_queue = []
self.is_eagle = is_eagle
if enable_metrics:
self.init_metrics_collector()
if self.token_to_kv_pool_allocator:
self.device = self.token_to_kv_pool_allocator.device
else:
@@ -483,6 +487,7 @@ class RadixCache(BasePrefixCache):
if self.disable:
return
start_time = time.perf_counter()
leaves = self._collect_leaves()
eviction_heap = [
(self.eviction_strategy.get_priority(node), node) for node in leaves
@@ -503,6 +508,12 @@ class RadixCache(BasePrefixCache):
self._record_remove_event(x)
if num_evicted > 0 and self.metrics_collector is not None:
self.metrics_collector.observe_eviction_duration(
time.perf_counter() - start_time
)
self.metrics_collector.increment_eviction_num_tokens(num_evicted)
def inc_lock_ref(self, node: TreeNode):
if self.disable:
return 0

11
python/sglang/srt/mem_cache/radix_cache_cpp.py
View File

@@ -1,6 +1,7 @@
from __future__ import annotations
import logging
import time
from typing import TYPE_CHECKING, List, Set
import torch
@@ -49,6 +50,7 @@ class RadixCacheCpp(BasePrefixCache):
hicache_ratio: float,
hicache_size: int,
hicache_write_policy: str,
enable_metrics: bool = False,
enable_kv_cache_events: bool = False,
hicache_oracle: bool = False,
enable_write_cancel: bool = False,
@@ -76,6 +78,9 @@ class RadixCacheCpp(BasePrefixCache):
self.tp_group = tp_cache_group
if enable_metrics:
self.init_metrics_collector()
if not use_hicache:
self.tree = RadixTreeCpp(
disabled=self.disable,
@@ -138,9 +143,15 @@ class RadixCacheCpp(BasePrefixCache):
self.tree.lock_ref(node, True)
def evict(self, num_tokens: int):
start_time = time.perf_counter()
evicted_device_indices = self.tree.evict(num_tokens)
for indice in evicted_device_indices:
self.token_to_kv_pool.free(indice)
if evicted_device_indices and self.metrics_collector is not None:
self.metrics_collector.observe_eviction_duration(
time.perf_counter() - start_time
)
self.metrics_collector.increment_eviction_num_tokens(num_tokens)
def evictable_size(self):
return self.tree.evictable_size()

2
python/sglang/srt/mem_cache/storage/lmcache/lmc_radix_cache.py
View File

@@ -73,6 +73,7 @@ class LMCRadixCache(RadixCache):
token_to_kv_pool_allocator: BaseTokenToKVPoolAllocator,
page_size: int,
disable: bool = False,
enable_metrics: bool = False,
enable_kv_cache_events: bool = False,
model_config: Optional["ModelConfig"] = None,
tp_size: int = 1,
@@ -85,6 +86,7 @@ class LMCRadixCache(RadixCache):
token_to_kv_pool_allocator=token_to_kv_pool_allocator,
page_size=page_size,
disable=disable,
enable_metrics=enable_metrics,
enable_kv_cache_events=enable_kv_cache_events,
eviction_policy=eviction_policy,
)

17
python/sglang/srt/mem_cache/swa_radix_cache.py
View File

@@ -20,6 +20,7 @@ The radix tree data structure for managing the hybrid (full and SWA) KV cache.
"""
import heapq
import time
from collections import defaultdict
from functools import partial
from typing import TYPE_CHECKING, List, Optional, Tuple
@@ -336,6 +337,7 @@ class SWARadixCache(BasePrefixCache):
page_size: int,
disable: bool = False,
is_eagle: bool = False,
enable_metrics: bool = False,
):
assert isinstance(token_to_kv_pool_allocator, SWATokenToKVPoolAllocator)
self.req_to_token_pool = req_to_token_pool
@@ -361,6 +363,9 @@ class SWARadixCache(BasePrefixCache):
else:
self.key_convert_fn = lambda key: key
if enable_metrics:
self.init_metrics_collector()
self.sliding_window_size = sliding_window_size
self.reset()
@@ -589,7 +594,7 @@ class SWARadixCache(BasePrefixCache):
def evict(self, full_num_tokens: int, swa_num_tokens: int = 0) -> None:
if self.disable:
return
start_time = time.perf_counter()
full_num_evicted = 0
swa_num_evicted = 0
if full_num_tokens > 0:
@@ -669,6 +674,16 @@ class SWARadixCache(BasePrefixCache):
x = x_next
if (
full_num_evicted > 0 or swa_num_evicted > 0
) and self.metrics_collector is not None:
self.metrics_collector.observe_eviction_duration(
time.perf_counter() - start_time
)
self.metrics_collector.increment_eviction_num_tokens(
full_num_evicted + swa_num_evicted
)
def inc_lock_ref(self, node: TreeNode) -> Optional[int]:
"""
Increment the lock reference count for the node. Returns the swa_uuid_for_lock, which needs

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

@@ -12,6 +12,7 @@
# limitations under the License.
# ==============================================================================
"""Utilities for Prometheus Metrics Collection."""
import os
import time
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Union
@@ -24,6 +25,20 @@ from sglang.srt.utils import get_bool_env_var
SGLANG_TEST_REQUEST_TIME_STATS = get_bool_env_var("SGLANG_TEST_REQUEST_TIME_STATS")
def get_histogram_conf_from_env(env_var_name: str) -> Optional[List[float]]:
"""
Get the histogram configuration from the environment variable.
env value should be like "0.1,0.2,0.5,1,2"
"""
if env_var_name not in os.environ:
return None
# if the env var is not set or empty, return None
env_var_value = os.environ[env_var_name]
if not env_var_value:
return None
return [float(x) for x in env_var_value.split(",")]
@dataclass
class TimeStats:
"""
@@ -184,6 +199,7 @@ class SchedulerStats:
# Engine startup
engine_startup_time: float = 0.0
engine_load_weights_time: float = 0.0
new_token_ratio: float = 0.0
# CUDA graph
is_cuda_graph: float = 0.0
@@ -191,7 +207,10 @@ class SchedulerStats:
class SchedulerMetricsCollector:
def __init__(self, labels: Dict[str, str]) -> None:
def __init__(
self,
labels: Dict[str, str],
) -> None:
# We need to import prometheus_client after setting the env variable `PROMETHEUS_MULTIPROC_DIR`
from prometheus_client import Counter, Gauge, Histogram
@@ -562,6 +581,13 @@ class SchedulerMetricsCollector:
multiprocess_mode="mostrecent",
)
self.new_token_ratio = Gauge(
name="sglang:new_token_ratio",
documentation="The new token ratio.",
labelnames=labels.keys(),
multiprocess_mode="mostrecent",
)
def _log_gauge(self, gauge, data: Union[int, float]) -> None:
# Convenience function for logging to gauge.
gauge.labels(**self.labels).set(data)
@@ -639,6 +665,7 @@ class SchedulerMetricsCollector:
self._log_gauge(
self.engine_load_weights_time, stats.engine_load_weights_time
)
self._log_gauge(self.new_token_ratio, stats.new_token_ratio)
# CUDA graph
self._log_gauge(self.is_cuda_graph, stats.is_cuda_graph)
@@ -1068,3 +1095,100 @@ class ExpertDispatchCollector:
labelnames={"layer"},
buckets=ep_size_buckets,
)
class RadixCacheMetricsCollector:
def __init__(
self,
labels: Dict[str, str],
) -> None:
# We need to import prometheus_client after setting the env variable `PROMETHEUS_MULTIPROC_DIR`
from prometheus_client import Counter, Histogram
self.labels = labels
bucket_eviction_duration = get_histogram_conf_from_env(
"SGLANG_BUCKET_EVICTION_DURATION"
)
if bucket_eviction_duration is None:
bucket_eviction_duration = [
0.001,
0.002,
0.003,
0.004,
0.005,
0.006,
0.007,
0.008,
0.009,
0.01,
0.02,
0.03,
0.04,
0.05,
0.1,
0.2,
0.5,
1.0,
]
bucket_load_back_duration = get_histogram_conf_from_env(
"SGLANG_BUCKET_LOAD_BACK_DURATION"
)
if bucket_load_back_duration is None:
bucket_load_back_duration = [
0.001,
0.002,
0.003,
0.004,
0.005,
0.006,
0.007,
0.008,
0.009,
0.01,
0.02,
0.03,
0.04,
0.05,
0.1,
0.2,
0.5,
1.0,
]
self.eviction_duration_seconds = Histogram(
name="sglang:eviction_duration_seconds",
documentation="Time taken to evict memory from GPU to CPU in seconds.",
labelnames=labels.keys(),
buckets=bucket_eviction_duration,
)
self.eviction_num_tokens = Counter(
name="sglang:evicted_tokens_total",
documentation="The number of tokens evicted from GPU to CPU.",
labelnames=labels.keys(),
)
self.load_back_duration_seconds = Histogram(
name="sglang:load_back_duration_seconds",
documentation="Time taken to load memory from CPU to GPU in seconds.",
labelnames=labels.keys(),
buckets=bucket_load_back_duration,
)
self.load_back_num_tokens = Counter(
name="sglang:load_back_tokens_total",
documentation="The number of tokens loaded from CPU to GPU.",
labelnames=labels.keys(),
)
def increment_eviction_num_tokens(self, num_tokens: int) -> None:
self.eviction_num_tokens.labels(**self.labels).inc(num_tokens)
def increment_load_back_num_tokens(self, num_tokens: int) -> None:
self.load_back_num_tokens.labels(**self.labels).inc(num_tokens)
def observe_eviction_duration(self, duration_seconds: float) -> None:
self.eviction_duration_seconds.labels(**self.labels).observe(duration_seconds)
def observe_load_back_duration(self, duration_seconds: float) -> None:
self.load_back_duration_seconds.labels(**self.labels).observe(duration_seconds)