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

Co-authored-by: Zhiqiang Xie <xiezhq@stanford.edu>

python/sglang/srt/managers/scheduler.py

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

python/sglang/srt/managers/scheduler_metrics_mixin.py

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

python/sglang/srt/mem_cache/base_prefix_cache.py

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

python/sglang/srt/mem_cache/hiradix_cache.py

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

python/sglang/srt/mem_cache/mamba_radix_cache.py

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

python/sglang/srt/mem_cache/radix_cache.py

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

python/sglang/srt/mem_cache/radix_cache_cpp.py

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

python/sglang/srt/mem_cache/storage/lmcache/lmc_radix_cache.py

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

python/sglang/srt/mem_cache/swa_radix_cache.py

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

python/sglang/srt/metrics/collector.py

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