Tiny simplify evcition metrics collector (#12983)

python/sglang/srt/environ.py

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

python/sglang/srt/managers/scheduler.py

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

python/sglang/srt/mem_cache/base_prefix_cache.py

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

python/sglang/srt/mem_cache/hiradix_cache.py

@@ -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.update_eviction_metrics(num_evicted, start_time)
-            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

python/sglang/srt/mem_cache/radix_cache.py

@@ -508,11 +508,7 @@ class RadixCache(BasePrefixCache):
 
             self._record_remove_event(x)
 
-        if num_evicted > 0 and self.metrics_collector is not None:
+        self.update_eviction_metrics(num_evicted, start_time)
-            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:

python/sglang/srt/mem_cache/radix_cache_cpp.py

@@ -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.device = token_to_kv_pool_allocator.device
-        self.token_to_kv_pool = token_to_kv_pool
+        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)
+            self.token_to_kv_pool_allocator.free(indice)
-        if evicted_device_indices and self.metrics_collector is not None:
+
-            self.metrics_collector.observe_eviction_duration(
+        # FIXME: not sure about the real evict length here
-                time.perf_counter() - start_time
+        self.update_eviction_metrics(num_tokens, start_time)
-            )
-            self.metrics_collector.increment_eviction_num_tokens(num_tokens)
 
     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

python/sglang/srt/mem_cache/swa_radix_cache.py

@@ -674,15 +674,7 @@ class SWARadixCache(BasePrefixCache):
 
                 x = x_next
 
-        if (
+        self.update_eviction_metrics(full_num_evicted + swa_num_evicted, start_time)
-            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]:
         """

test/srt/test_cpp_radix_cache.py

@@ -2,6 +2,7 @@ import os
 import unittest
 from types import SimpleNamespace
 
+from sglang.srt.environ import envs
 from sglang.srt.utils import kill_process_tree
 from sglang.test.run_eval import run_eval
 from sglang.test.test_utils import (
@@ -16,7 +17,7 @@ from sglang.test.test_utils import (
 class TestCppRadixCache(CustomTestCase):
     @classmethod
     def setUpClass(cls):
-        os.environ["SGLANG_EXPERIMENTAL_CPP_RADIX_TREE"] = "1"
+        envs.SGLANG_EXPERIMENTAL_CPP_RADIX_TREE.set(True)
         cls.model = DEFAULT_MODEL_NAME_FOR_TEST
         cls.base_url = DEFAULT_URL_FOR_TEST
         cls.process = popen_launch_server(