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

Tiny simplify evcition metrics collector (#12983)

Browse Source
This commit is contained in:
Liangsheng Yin
2025-11-11 23:23:48 +08:00
committed by GitHub
parent e38994dd71
commit f09eee036d
8 changed files with 34 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -162,6 +162,8 @@ class Envs:
# Scheduler: others:
SGLANG_EMPTY_CACHE_INTERVAL = EnvFloat(-1) # in seconds. Set if you observe high memory accumulation over a long serving period.
SGLANG_EXPERIMENTAL_CPP_RADIX_TREE = EnvBool(False)
# Test: pd-disaggregation
SGLANG_TEST_PD_DISAGG_BACKEND = EnvStr("mooncake")
SGLANG_TEST_PD_DISAGG_DEVICES = EnvStr(None)

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

@@ -713,15 +713,16 @@ class Scheduler(
page_size=self.page_size,
)
else:
if os.environ.get("SGLANG_EXPERIMENTAL_CPP_RADIX_TREE") == "1":
if envs.SGLANG_EXPERIMENTAL_CPP_RADIX_TREE.get():
# lazy import to avoid JIT overhead
from sglang.srt.mem_cache.radix_cache_cpp import RadixCacheCpp
logger.info("Using experimental C++ radix tree implementation.")
self.tree_cache = RadixCacheCpp(
disable=False,
use_hicache=self.enable_hierarchical_cache,
req_to_token_pool=self.req_to_token_pool,
token_to_kv_pool=self.token_to_kv_pool_allocator,
token_to_kv_pool_allocator=self.token_to_kv_pool_allocator,
tp_cache_group=self.tp_cpu_group,
page_size=self.page_size,
hicache_ratio=server_args.hicache_ratio,

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

@@ -1,5 +1,6 @@
from __future__ import annotations
import time
from abc import ABC, abstractmethod
from typing import (
TYPE_CHECKING,
@@ -60,6 +61,13 @@ class BasePrefixCache(ABC, PrefixCacheTrait):
labels={"cache_type": self.__class__.__name__}
)
def update_eviction_metrics(self, num_evicted: int, start_time: float):
if self.metrics_collector is not None and num_evicted > 0:
self.metrics_collector.observe_eviction_duration(
time.perf_counter() - start_time
)
self.metrics_collector.increment_eviction_num_tokens(num_evicted)
@abstractmethod
def reset(self):
pass

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

@@ -376,11 +376,7 @@ 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)
self.update_eviction_metrics(num_evicted, start_time)
def _evict_backuped(self, node: TreeNode):
# evict a node already written to host

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

@@ -508,11 +508,7 @@ 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)
self.update_eviction_metrics(num_evicted, start_time)
def inc_lock_ref(self, node: TreeNode):
if self.disable:

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

@@ -44,7 +44,7 @@ class RadixCacheCpp(BasePrefixCache):
disable: bool,
use_hicache: bool,
req_to_token_pool: ReqToTokenPool,
token_to_kv_pool: BaseTokenToKVPoolAllocator,
token_to_kv_pool_allocator: BaseTokenToKVPoolAllocator,
tp_cache_group: torch.distributed.ProcessGroup,
page_size: int,
hicache_ratio: float,
@@ -61,7 +61,7 @@ class RadixCacheCpp(BasePrefixCache):
assert (
enable_kv_cache_events is False
), "HiRadixCache does not support kv cache events yet"
self.kv_cache = token_to_kv_pool.get_kvcache()
self.kv_cache = token_to_kv_pool_allocator.get_kvcache()
# record the nodes with ongoing write through
self.ongoing_write_through: Set[IOHandle] = set()
@@ -71,8 +71,8 @@ class RadixCacheCpp(BasePrefixCache):
self.write_through_threshold = (
1 if hicache_write_policy == "write_through" else 2
)
self.device = token_to_kv_pool.device
self.token_to_kv_pool = token_to_kv_pool
self.device = token_to_kv_pool_allocator.device
self.token_to_kv_pool_allocator = token_to_kv_pool_allocator
self.req_to_token_pool = req_to_token_pool
self.page_size = page_size
@@ -146,12 +146,10 @@ class RadixCacheCpp(BasePrefixCache):
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)
self.token_to_kv_pool_allocator.free(indice)
# FIXME: not sure about the real evict length here
self.update_eviction_metrics(num_tokens, start_time)
def evictable_size(self):
return self.tree.evictable_size()
@@ -184,16 +182,18 @@ class RadixCacheCpp(BasePrefixCache):
assert old_prefix_len <= new_prefix_len, "Wrong prefix indices"
# Free duplicates that were already in the pool
if old_prefix_len < new_prefix_len:
self.token_to_kv_pool.free(kv_indices[old_prefix_len:new_prefix_len])
self.token_to_kv_pool_allocator.free(
kv_indices[old_prefix_len:new_prefix_len]
)
else:
self.token_to_kv_pool.free(
self.token_to_kv_pool_allocator.free(
kv_indices[old_prefix_len:page_aligned_overall_len]
)
# need to free the unaligned part, since it cannot be inserted into the radix tree
if page_aligned_overall_len < kv_committed_len:
# NOTE: sglang PagedAllocator support unaligned free (which will automatically align it)
self.token_to_kv_pool.free(kv_indices[page_aligned_overall_len:])
self.token_to_kv_pool_allocator.free(kv_indices[page_aligned_overall_len:])
# Remove req slot release the cache lock
self.dec_lock_ref(req.last_node)
@@ -225,7 +225,9 @@ class RadixCacheCpp(BasePrefixCache):
# KVCache between old & new is newly generated, but already exists in the pool
# we need to free this newly generated kv indices and reuse the indices in the pool
if old_prefix_len < new_prefix_len:
self.token_to_kv_pool.free(kv_indices[old_prefix_len:new_prefix_len])
self.token_to_kv_pool_allocator.free(
kv_indices[old_prefix_len:new_prefix_len]
)
reused_indices = new_indices[old_prefix_len:new_prefix_len]
self.req_to_token_pool.req_to_token[
req.req_pool_idx, old_prefix_len:new_prefix_len

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

@@ -674,15 +674,7 @@ 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
)
self.update_eviction_metrics(full_num_evicted + swa_num_evicted, start_time)
def inc_lock_ref(self, node: TreeNode) -> Optional[int]:
"""