diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index f3f8e644d..25da56d87 100644 --- a/python/sglang/srt/environ.py +++ b/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) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 077bc8543..a36ea17f5 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/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, diff --git a/python/sglang/srt/mem_cache/base_prefix_cache.py b/python/sglang/srt/mem_cache/base_prefix_cache.py index 111507462..9c16b639b 100644 --- a/python/sglang/srt/mem_cache/base_prefix_cache.py +++ b/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 diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index 312732e0d..11665d666 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/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.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 diff --git a/python/sglang/srt/mem_cache/radix_cache.py b/python/sglang/srt/mem_cache/radix_cache.py index 12118e536..24a2eaf8b 100644 --- a/python/sglang/srt/mem_cache/radix_cache.py +++ b/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.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: diff --git a/python/sglang/srt/mem_cache/radix_cache_cpp.py b/python/sglang/srt/mem_cache/radix_cache_cpp.py index b84980d6a..fbe6afc41 100644 --- a/python/sglang/srt/mem_cache/radix_cache_cpp.py +++ b/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.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 diff --git a/python/sglang/srt/mem_cache/swa_radix_cache.py b/python/sglang/srt/mem_cache/swa_radix_cache.py index 19e14af81..e19f9320c 100644 --- a/python/sglang/srt/mem_cache/swa_radix_cache.py +++ b/python/sglang/srt/mem_cache/swa_radix_cache.py @@ -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]: """ diff --git a/test/srt/test_cpp_radix_cache.py b/test/srt/test_cpp_radix_cache.py index cb2822b88..2e6baa653 100644 --- a/test/srt/test_cpp_radix_cache.py +++ b/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(