diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index aaa6a4ce4..02ecb538c 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/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(), diff --git a/python/sglang/srt/managers/scheduler_metrics_mixin.py b/python/sglang/srt/managers/scheduler_metrics_mixin.py index e251936b8..53714c9ed 100644 --- a/python/sglang/srt/managers/scheduler_metrics_mixin.py +++ b/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 = ( diff --git a/python/sglang/srt/mem_cache/base_prefix_cache.py b/python/sglang/srt/mem_cache/base_prefix_cache.py index be6968e8b..111507462 100644 --- a/python/sglang/srt/mem_cache/base_prefix_cache.py +++ b/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 diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index 6ea4e1ba9..312732e0d 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/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 ) diff --git a/python/sglang/srt/mem_cache/mamba_radix_cache.py b/python/sglang/srt/mem_cache/mamba_radix_cache.py index 7ad5746f9..286f702ca 100644 --- a/python/sglang/srt/mem_cache/mamba_radix_cache.py +++ b/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() diff --git a/python/sglang/srt/mem_cache/radix_cache.py b/python/sglang/srt/mem_cache/radix_cache.py index 145e4dac3..c96e83ef2 100644 --- a/python/sglang/srt/mem_cache/radix_cache.py +++ b/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 diff --git a/python/sglang/srt/mem_cache/radix_cache_cpp.py b/python/sglang/srt/mem_cache/radix_cache_cpp.py index 7994e372b..5342e9dc2 100644 --- a/python/sglang/srt/mem_cache/radix_cache_cpp.py +++ b/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() diff --git a/python/sglang/srt/mem_cache/storage/lmcache/lmc_radix_cache.py b/python/sglang/srt/mem_cache/storage/lmcache/lmc_radix_cache.py index 68f1d658e..cf45675d9 100644 --- a/python/sglang/srt/mem_cache/storage/lmcache/lmc_radix_cache.py +++ b/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, ) diff --git a/python/sglang/srt/mem_cache/swa_radix_cache.py b/python/sglang/srt/mem_cache/swa_radix_cache.py index 43e816125..adf17a668 100644 --- a/python/sglang/srt/mem_cache/swa_radix_cache.py +++ b/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 diff --git a/python/sglang/srt/metrics/collector.py b/python/sglang/srt/metrics/collector.py index a6e94c89a..205a3cdbc 100644 --- a/python/sglang/srt/metrics/collector.py +++ b/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)