diff --git a/python/sglang/srt/mem_cache/base_prefix_cache.py b/python/sglang/srt/mem_cache/base_prefix_cache.py index 73341b7f3..abc840cda 100644 --- a/python/sglang/srt/mem_cache/base_prefix_cache.py +++ b/python/sglang/srt/mem_cache/base_prefix_cache.py @@ -17,6 +17,7 @@ from typing import ( import torch from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator +from sglang.srt.mem_cache.events import KVCacheEventRecorder from sglang.srt.mem_cache.memory_pool import ReqToTokenPool from sglang.srt.observability.metrics_collector import ( STAT_LOGGER_ROLE_RADIX_CACHE, @@ -238,6 +239,8 @@ class BasePrefixCache(ABC, PrefixCacheTrait): None # metrics collector for the cache ) cache_controller: Optional[HiCacheController] = None + # Set by caches that publish KV placement events; None means they don't. + kv_events: Optional[KVCacheEventRecorder] = None def init_metrics_collector(self): from sglang.srt.runtime_context import get_server_args @@ -377,7 +380,7 @@ class BasePrefixCache(ABC, PrefixCacheTrait): raise NotImplementedError() def take_events(self): - return [] + return [] if self.kv_events is None else self.kv_events.take() def supports_swa(self) -> bool: return False diff --git a/python/sglang/srt/mem_cache/events.py b/python/sglang/srt/mem_cache/events.py index 4f473d46b..a3cc54b9e 100644 --- a/python/sglang/srt/mem_cache/events.py +++ b/python/sglang/srt/mem_cache/events.py @@ -11,13 +11,14 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""KV cache placement event emission mixin. +"""KV cache placement event recording. -The mixin produces the ``BlockStored`` / ``BlockRemoved`` / ``AllBlocksCleared`` -events consumed by KV-aware routers (e.g. dynamo). +Produces the ``BlockStored`` / ``BlockRemoved`` / ``AllBlocksCleared`` events +consumed by KV-aware routers (e.g. dynamo). A cache holds one recorder and calls +it; the recorder owns the queue and needs nothing back from its owner. """ -from typing import Any +from typing import Any, Optional from sglang.srt.disaggregation.kv_events import ( AllBlocksCleared, @@ -34,16 +35,27 @@ from sglang.srt.mem_cache.utils import ( ) -class KVCacheEventMixin: - def _enqueue_kv_event(self, event): +class KVCacheEventRecorder: + """Collects KV placement events for one cache. + + ``enabled=False`` makes every ``record_*`` call a no-op and ``take`` return an + empty list, so callers never have to guard. + """ + + def __init__(self, *, enabled: bool, page_size: int): + self.enabled = enabled + self.page_size = page_size + self._queue: list = [] + + def enqueue(self, event) -> None: """Append an event, coalescing it with a compatible queue tail. KV event batches already support multiple block hashes. Combining them here avoids emitting one event per page while preserving ordering and the parent-linked store chains consumers use to rebuild the cache tree. """ - if self.kv_event_queue: - tail = self.kv_event_queue[-1] + if self._queue: + tail = self._queue[-1] if isinstance(tail, BlockRemoved) and isinstance(event, BlockRemoved): if tail.medium == event.medium: @@ -71,117 +83,122 @@ class KVCacheEventMixin: tail.token_ids.extend(event.token_ids) return - self.kv_event_queue.append(event) + self._queue.append(event) - def _record_store_event(self, node: Any, medium=None): + def _node_event_hash_values(self, node: Any) -> list: + """Hash values to publish for ``node``, computing them if not yet set.""" + if node.hash_value is None: + node.hash_value = compute_node_hash_values(node, self.page_size) + if node.key.cache_salt is not None: + return compute_node_event_hash_values(node, self.page_size) + return node.hash_value + + def _parent_block_hash(self, node: Any) -> Optional[int]: + """The hash the first page of ``node`` links back to. + + ``None`` when the parent is the tree root: a root carries an empty + ``hash_value`` and no event hash, so it contributes no link. Every other + node on the path has a parent, which is what distinguishes the two. + """ + parent = node.parent + if parent is None or parent.parent is None: + return None + if node.key.cache_salt is not None: + parent_hash_values = parent.event_hash_value + assert parent_hash_values is not None + else: + parent_hash_values = parent.hash_value + if not parent_hash_values: + return None + return hash_str_to_int64(parent_hash_values[-1]) + + def record_store(self, node: Any, medium=None) -> None: # One BlockStored per ``page_size`` chunk. # ``medium`` defaults to StorageMedium.GPU but callers may override # for lower-tier insertions (e.g. StorageMedium.CPU for host/L2 cache). - if self.enable_kv_cache_events: - if medium is None: - medium = StorageMedium.GPU + if not self.enabled: + return + if medium is None: + medium = StorageMedium.GPU - # Compute hash_value lazily if not already set - if node.hash_value is None: - node.hash_value = compute_node_hash_values(node, self.page_size) - event_hash_values = ( - compute_node_event_hash_values(node, self.page_size) - if node.key.cache_salt is not None - else node.hash_value - ) + event_hash_values = self._node_event_hash_values(node) + parent_block_hash = self._parent_block_hash(node) - # Get parent's last hash value for first page - parent_block_hash = None - if node.parent is not None and node.parent != self.root_node: - if node.key.cache_salt is not None: - parent_hash_values = node.parent.event_hash_value - assert parent_hash_values is not None - else: - parent_hash_values = node.parent.hash_value - if parent_hash_values: - parent_block_hash = hash_str_to_int64(parent_hash_values[-1]) + page_index = 0 + logical_len = len(node.key) + is_bigram = node.key.is_bigram + raw = node.key.token_ids + for start in range(0, logical_len, self.page_size): + end = min(start + self.page_size, logical_len) + if end <= start: + continue + # Preserve historical event payload: bigram pages expose tuples. + if is_bigram: + page_tokens = [(raw[j], raw[j + 1]) for j in range(start, end)] + else: + page_tokens = list(raw[start:end]) - page_index = 0 - logical_len = len(node.key) - is_bigram = node.key.is_bigram - raw = node.key.token_ids - for start in range(0, logical_len, self.page_size): - end = min(start + self.page_size, logical_len) - if end <= start: - continue - # Preserve historical event payload: bigram pages expose tuples. - if is_bigram: - page_tokens = [(raw[j], raw[j + 1]) for j in range(start, end)] - else: - page_tokens = list(raw[start:end]) + block_hash = hash_str_to_int64(event_hash_values[page_index]) - block_hash = hash_str_to_int64(event_hash_values[page_index]) + event_args = { + "block_hashes": [block_hash], + "parent_block_hash": parent_block_hash, + "token_ids": page_tokens, + "block_size": len(page_tokens), + "lora_id": None, + "medium": medium, + } + if node.key.cache_salt is None: + event = BlockStored(**event_args) + else: + event = BlockStoredWithMetadata( + **event_args, + metadata=BlockStoredMetadata(cache_salt=node.key.cache_salt), + ) + self.enqueue(event) - event_args = { - "block_hashes": [block_hash], - "parent_block_hash": parent_block_hash, - "token_ids": page_tokens, - "block_size": len(page_tokens), - "lora_id": None, - "medium": medium, - } - if node.key.cache_salt is None: - event = BlockStored(**event_args) - else: - event = BlockStoredWithMetadata( - **event_args, - metadata=BlockStoredMetadata(cache_salt=node.key.cache_salt), - ) - self._enqueue_kv_event(event) + parent_block_hash = block_hash + page_index += 1 - parent_block_hash = block_hash - page_index += 1 - - def _record_remove_event(self, node: Any, medium=None): + def record_remove(self, node: Any, medium=None) -> None: # One BlockRemoved per radix node. # ``medium`` defaults to StorageMedium.GPU but callers may override for # lower-tier removals (e.g. StorageMedium.CPU when evicting from host). - if self.enable_kv_cache_events: - if medium is None: - medium = StorageMedium.GPU + if not self.enabled: + return + if medium is None: + medium = StorageMedium.GPU - # Compute hash_value lazily if not already set (must match what was stored) - if node.hash_value is None: - node.hash_value = compute_node_hash_values(node, self.page_size) - event_hash_values = ( - compute_node_event_hash_values(node, self.page_size) - if node.key.cache_salt is not None - else node.hash_value - ) + # Hash values must match what was stored. + event_hash_values = self._node_event_hash_values(node) - block_hashes = [] - logical_len = len(node.key) - page_index = 0 - for start in range(0, logical_len, self.page_size): - end = min(start + self.page_size, logical_len) - if end <= start: - continue + block_hashes = [] + logical_len = len(node.key) + page_index = 0 + for start in range(0, logical_len, self.page_size): + end = min(start + self.page_size, logical_len) + if end <= start: + continue - block_hashes.append(hash_str_to_int64(event_hash_values[page_index])) - page_index += 1 + block_hashes.append(hash_str_to_int64(event_hash_values[page_index])) + page_index += 1 - if block_hashes: - self._enqueue_kv_event( - BlockRemoved(block_hashes=block_hashes, medium=medium) - ) + if block_hashes: + self.enqueue(BlockRemoved(block_hashes=block_hashes, medium=medium)) - def _record_all_cleared_event(self): - if self.enable_kv_cache_events: - self._enqueue_kv_event(AllBlocksCleared()) + def record_all_cleared(self) -> None: + if not self.enabled: + return + self.enqueue(AllBlocksCleared()) - def take_events(self): + def take(self) -> list: """Atomically takes all events and clears the queue. Returns: A list of KV cache events. """ - if not self.enable_kv_cache_events: + if not self.enabled: return [] - events = self.kv_event_queue - self.kv_event_queue = [] + events = self._queue + self._queue = [] return events diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index cee15cb05..b86485657 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -908,7 +908,7 @@ class HiRadixCache(RadixCache): if node.write_through_pending_id == ack_id: node.write_through_pending_id = None # DMA confirmed -- block is now on host. - self._record_store_event(node, medium=StorageMedium.CPU) + self.kv_events.record_store(node, medium=StorageMedium.CPU) if self.enable_storage: self.write_backup_storage(lock_node, backup_len) if release_lock: @@ -1264,7 +1264,7 @@ class HiRadixCache(RadixCache): def _detach_backuped(self, node: TreeNode) -> int: # detach nodes from tree while keeping device slots, for write-back eviction - self._record_remove_event(node, medium=StorageMedium.GPU) + self.kv_events.record_remove(node, medium=StorageMedium.GPU) num_evicted = len(node.value) assert num_evicted > 0 self.evictable_size_ -= num_evicted @@ -1285,7 +1285,7 @@ class HiRadixCache(RadixCache): # evict a node not initiated write to host -- emit BlockRemoved assert len(node.children) == 0, f"non-leaf, {node.id=}" - self._record_remove_event(node) + self.kv_events.record_remove(node) self.cache_controller.mem_pool_device_allocator.free(node.value) num_evicted = len(node.value) self._delete_leaf(node) @@ -1311,11 +1311,11 @@ class HiRadixCache(RadixCache): freed_device = 0 for n in nodes: if n.host_value is not None: - self._record_remove_event(n, medium=StorageMedium.CPU) + self.kv_events.record_remove(n, medium=StorageMedium.CPU) self.cache_controller.evict_host(n.host_value) n.host_value = None if n.value is not None: - self._record_remove_event(n, medium=StorageMedium.GPU) + self.kv_events.record_remove(n, medium=StorageMedium.GPU) self.cache_controller.mem_pool_device_allocator.free(n.value) freed_device += len(n.value) self.evictable_size_ -= len(n.value) @@ -1357,7 +1357,7 @@ class HiRadixCache(RadixCache): # Block deleted entirely (GPU already evicted, now CPU freed) -- # emit remove(CPU) so the router drops the host-tier entry. - self._record_remove_event(x, medium=StorageMedium.CPU) + self.kv_events.record_remove(x, medium=StorageMedium.CPU) num_evicted += self.cache_controller.evict_host(x.host_value) key = x.key.child_key(self.page_size) @@ -1438,7 +1438,7 @@ class HiRadixCache(RadixCache): offset += len(node.host_value) # Block promoted from host to GPU -- emit store(GPU) so downstream # indexers see it as device-local again. - self._record_store_event(node, medium=StorageMedium.GPU) + self.kv_events.record_store(node, medium=StorageMedium.GPU) self.evictable_size_ += len(device_indices) self.inc_lock_ref(last_hit_node) @@ -1849,7 +1849,7 @@ class HiRadixCache(RadixCache): self._update_host_leaf_status(node) # Publish the newly materialized host suffix immediately so downstream # cache indexers can resolve descendants that extend this L2-only prefix. - self._record_store_event(new_node, medium=StorageMedium.CPU) + self.kv_events.record_store(new_node, medium=StorageMedium.CPU) return matched_length @@ -1987,11 +1987,11 @@ class HiRadixCache(RadixCache): self._update_leaf_status(new_node) # Compute hash_value if storage or kv events are enabled - if self.enable_storage or self.enable_kv_cache_events: + if self.enable_storage or self.kv_events.enabled: new_node.hash_value = compute_node_hash_values(new_node, self.page_size) # Emit BlockStored so the router indexes this block. - self._record_store_event(new_node) + self.kv_events.record_store(new_node) if self.cache_controller.write_policy != "write_back": self._inc_hit_count(new_node, chunked) diff --git a/python/sglang/srt/mem_cache/mamba_radix_cache.py b/python/sglang/srt/mem_cache/mamba_radix_cache.py index a8559db3e..a5afb89d4 100644 --- a/python/sglang/srt/mem_cache/mamba_radix_cache.py +++ b/python/sglang/srt/mem_cache/mamba_radix_cache.py @@ -43,7 +43,7 @@ from sglang.srt.mem_cache.base_prefix_cache import ( MatchPrefixParams, MatchResult, ) -from sglang.srt.mem_cache.events import KVCacheEventMixin +from sglang.srt.mem_cache.events import KVCacheEventRecorder from sglang.srt.mem_cache.memory_pool import HybridReqToTokenPool from sglang.srt.mem_cache.multi_ended_allocator import ( UnifiedMambaTokenToKVPoolAllocator, @@ -441,7 +441,7 @@ class LRUList: raise Exception(msg) -class MambaRadixCache(KVCacheEventMixin, BasePrefixCache): +class MambaRadixCache(BasePrefixCache): def __init__(self, params: CacheInitParams): assert ( isinstance(params.token_to_kv_pool_allocator, TokenToKVPoolAllocator) @@ -458,10 +458,11 @@ class MambaRadixCache(KVCacheEventMixin, BasePrefixCache): self.page_size = params.page_size self.disable = params.disable - self.enable_kv_cache_events = params.enable_kv_cache_events self.enable_mamba_extra_buffer = params.enable_mamba_extra_buffer self.enable_mamba_extra_buffer_lazy = params.enable_mamba_extra_buffer_lazy - self.kv_event_queue = [] + self.kv_events = KVCacheEventRecorder( + enabled=params.enable_kv_cache_events, page_size=self.page_size + ) if not self.enable_mamba_extra_buffer: assert ( @@ -497,7 +498,7 @@ class MambaRadixCache(KVCacheEventMixin, BasePrefixCache): # LRU lists are used to maintain the order of eviction of the nodes in the tree self.full_lru_list = LRUList(mamba=False) self.mamba_lru_list = LRUList(mamba=True) - self._record_all_cleared_event() + self.kv_events.record_all_cleared() def match_prefix(self, params: MatchPrefixParams) -> MatchResult: """Find the matching prefix from the radix tree. @@ -819,7 +820,7 @@ class MambaRadixCache(KVCacheEventMixin, BasePrefixCache): assert x.mamba_value is not None, f"leaf node mamba value is not None, {x.id=}" # 1. a leaf node, free full tokens and mamba - self._record_remove_event(x) + self.kv_events.record_remove(x) # Tree values are page-aligned copies of a kv row: page-exact segment. self.token_to_kv_pool_allocator.free_segment(x.value, start_pos=0) full_num_evicted = len(x.value) @@ -1301,7 +1302,7 @@ class MambaRadixCache(KVCacheEventMixin, BasePrefixCache): node.children[child_key] = new_node self.full_evictable_size_ += len(value) self.mamba_evictable_size_ += len(mamba_value) - self._record_store_event(new_node) + self.kv_events.record_store(new_node) elif node.mamba_value is None: # add for mamba tombstone node.mamba_value = mamba_value self.full_lru_list.reset_node_mru(node) @@ -1330,7 +1331,7 @@ class MambaRadixCache(KVCacheEventMixin, BasePrefixCache): node.parent.mamba_lock_ref == 0 ), f"tombstone mamba_lock_ref should always be 0, {node.parent.full_lock_ref=}, {node.parent.mamba_lock_ref=}, {node.parent.id=}" # delete tombstone node evicts full tokens - self._record_remove_event(node.parent) + self.kv_events.record_remove(node.parent) self.token_to_kv_pool_allocator.free_segment(node.parent.value, start_pos=0) full_num_evicted += len(node.parent.value) self.full_lru_list.remove_node(node.parent) diff --git a/python/sglang/srt/mem_cache/radix_cache.py b/python/sglang/srt/mem_cache/radix_cache.py index f7eebd75f..60ec95137 100644 --- a/python/sglang/srt/mem_cache/radix_cache.py +++ b/python/sglang/srt/mem_cache/radix_cache.py @@ -45,7 +45,7 @@ from sglang.srt.mem_cache.base_prefix_cache import ( MatchPrefixParams, MatchResult, ) -from sglang.srt.mem_cache.events import KVCacheEventMixin +from sglang.srt.mem_cache.events import KVCacheEventRecorder from sglang.srt.mem_cache.utils import ( get_eviction_strategy, get_hash_str, @@ -300,18 +300,19 @@ class TreeNode: return self.last_access_time < other.last_access_time -class RadixCache(KVCacheEventMixin, BasePrefixCache): +class RadixCache(BasePrefixCache): def __init__(self, params: CacheInitParams): self.disable = params.disable self.req_to_token_pool = params.req_to_token_pool self.token_to_kv_pool_allocator = params.token_to_kv_pool_allocator self.page_size = params.page_size - self.enable_kv_cache_events = params.enable_kv_cache_events self.is_eagle = params.is_eagle self.disable_finished_insert = params.disable_finished_insert self.eviction_policy = params.eviction_policy.lower() - self.kv_event_queue = [] + self.kv_events = KVCacheEventRecorder( + enabled=params.enable_kv_cache_events, page_size=self.page_size + ) if params.enable_metrics: self.init_metrics_collector() @@ -371,7 +372,7 @@ class RadixCache(KVCacheEventMixin, BasePrefixCache): last_host_node=self.root_node, best_match_node=self.root_node, ) - self._record_all_cleared_event() + self.kv_events.record_all_cleared() def match_prefix(self, params: MatchPrefixParams) -> MatchResult: """Find the longest cached prefix of ``key`` in the radix tree. @@ -614,7 +615,7 @@ class RadixCache(KVCacheEventMixin, BasePrefixCache): new_priority = self.eviction_strategy.get_priority(x.parent) heapq.heappush(eviction_heap, (new_priority, x.parent)) - self._record_remove_event(x) + self.kv_events.record_remove(x) self.update_eviction_metrics(num_evicted, start_time) return EvictResult(num_tokens_evicted=num_evicted) @@ -785,7 +786,7 @@ class RadixCache(KVCacheEventMixin, BasePrefixCache): self._update_leaf_status(node) self._update_leaf_status(new_node) # Hash will be computed lazily during event emission - self._record_store_event(new_node) + self.kv_events.record_store(new_node) node = new_node return total_prefix_length, node diff --git a/python/sglang/srt/mem_cache/storage/flexkv/flexkv_radix_cache.py b/python/sglang/srt/mem_cache/storage/flexkv/flexkv_radix_cache.py index aeaa90f17..a7a090cd9 100644 --- a/python/sglang/srt/mem_cache/storage/flexkv/flexkv_radix_cache.py +++ b/python/sglang/srt/mem_cache/storage/flexkv/flexkv_radix_cache.py @@ -375,8 +375,8 @@ class FlexKVRadixCache(RadixCache): self._update_leaf_status(last_node) self._update_leaf_status(new_node) - self._record_store_event(new_node.parent) - self._record_store_event(new_node) + self.kv_events.record_store(new_node.parent) + self.kv_events.record_store(new_node) return fetched_slots, new_node 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 f8ce455d0..5a24c5d19 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 @@ -383,8 +383,8 @@ class LMCRadixCache(RadixCache): self._update_leaf_status(last_node) self._update_leaf_status(new_node) - self._record_store_event(new_node.parent) - self._record_store_event(new_node) + self.kv_events.record_store(new_node.parent) + self.kv_events.record_store(new_node) return token_slots[:fetched], new_node diff --git a/python/sglang/srt/mem_cache/swa_radix_cache.py b/python/sglang/srt/mem_cache/swa_radix_cache.py index 01c8814db..4ecbd68b2 100644 --- a/python/sglang/srt/mem_cache/swa_radix_cache.py +++ b/python/sglang/srt/mem_cache/swa_radix_cache.py @@ -42,7 +42,7 @@ from sglang.srt.mem_cache.base_prefix_cache import ( MatchResult, ) from sglang.srt.mem_cache.cache_init_params import CacheInitParams -from sglang.srt.mem_cache.events import KVCacheEventMixin +from sglang.srt.mem_cache.events import KVCacheEventRecorder from sglang.srt.mem_cache.radix_cache import RadixKey from sglang.srt.mem_cache.utils import split_node_hash_value @@ -342,7 +342,7 @@ class LRUList: raise Exception(msg) -class SWARadixCache(KVCacheEventMixin, BasePrefixCache): +class SWARadixCache(BasePrefixCache): def __init__(self, params: CacheInitParams): assert isinstance(params.token_to_kv_pool_allocator, SWATokenToKVPoolAllocator) self.req_to_token_pool = params.req_to_token_pool @@ -350,8 +350,9 @@ class SWARadixCache(KVCacheEventMixin, BasePrefixCache): self.page_size = params.page_size self.disable = params.disable self.is_eagle = params.is_eagle - self.enable_kv_cache_events = params.enable_kv_cache_events - self.kv_event_queue = [] + self.kv_events = KVCacheEventRecorder( + enabled=params.enable_kv_cache_events, page_size=self.page_size + ) if self.token_to_kv_pool_allocator: self.device = self.token_to_kv_pool_allocator.device @@ -407,7 +408,7 @@ class SWARadixCache(KVCacheEventMixin, BasePrefixCache): # LRU lists are used to maintain the order of eviction of the nodes in the tree self.full_lru_list = LRUList(is_swa_list=False) self.swa_lru_list = LRUList(is_swa_list=True) - self._record_all_cleared_event() + self.kv_events.record_all_cleared() def match_prefix(self, params: MatchPrefixParams) -> MatchResult: """Find the matching prefix from the radix tree. @@ -609,7 +610,7 @@ class SWARadixCache(KVCacheEventMixin, BasePrefixCache): assert x.full_lock_ref == 0, f"node is in use, {x.id=}" # 1. free node kv indices, evict full and swa tokens - self._record_remove_event(x) + self.kv_events.record_remove(x) self.token_to_kv_pool_allocator.free(x.value) full_num_evicted += len(x.value) # Tombstoned leaves had their SWA freed earlier in `dec_swa_lock_only` @@ -674,7 +675,7 @@ class SWARadixCache(KVCacheEventMixin, BasePrefixCache): x.full_lock_ref == 0 ), f"leaf node with full lock must also have swa lock, {x.id=}" # 1. a leaf node, free full and swa tokens - self._record_remove_event(x) + self.kv_events.record_remove(x) self.token_to_kv_pool_allocator.free(x.value) full_num_evicted += len(x.value) swa_num_evicted += len(x.value) @@ -1326,7 +1327,7 @@ class SWARadixCache(KVCacheEventMixin, BasePrefixCache): if not swa_tombstone: self.swa_lru_list.insert_mru(new_node) self.swa_evictable_size_ += len(value) - self._record_store_event(new_node) + self.kv_events.record_store(new_node) return new_node def _iteratively_delete_tombstone_leaf( @@ -1344,7 +1345,7 @@ class SWARadixCache(KVCacheEventMixin, BasePrefixCache): node.parent.swa_lock_ref == 0 ), f"tombstone swa_lock_ref should always be 0, {node.parent.full_lock_ref=}, {node.parent.swa_lock_ref=}, {node.parent.id=}" # delete tombstone node evicts full tokens - self._record_remove_event(node.parent) + self.kv_events.record_remove(node.parent) self.token_to_kv_pool_allocator.free(node.parent.value) full_num_evicted += len(node.parent.value) self.full_lru_list.remove_node(node.parent) diff --git a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py index e39a9e01c..16180016e 100644 --- a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py +++ b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py @@ -36,6 +36,7 @@ from sglang.srt.mem_cache.base_prefix_cache import ( MatchPrefixParams, MatchResult, ) +from sglang.srt.mem_cache.events import KVCacheEventRecorder from sglang.srt.mem_cache.hicache_storage import ( PoolName, PoolTransfer, @@ -417,8 +418,9 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface): self.components_by_type.values() ) - self.enable_kv_cache_events = params.enable_kv_cache_events - self.kv_event_queue = [] + self.kv_events = KVCacheEventRecorder( + enabled=params.enable_kv_cache_events, page_size=self.page_size + ) self.reset() @@ -1119,7 +1121,7 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface): self._update_evictable_leaf_sets(new_node) self._update_evictable_leaf_sets(parent) - self._record_store_event(new_node) + self.kv_events.record_store(new_node) return new_node def _unevict_node_on_insert( @@ -1138,7 +1140,7 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface): self._update_duplicate_tracking(node) if node.parent is not None: self._update_evictable_leaf_sets(node.parent) - self._record_store_event(node, medium=StorageMedium.GPU) + self.kv_events.record_store(node, medium=StorageMedium.GPU) def _update_evictable_leaf_sets(self, node: UnifiedTreeNode) -> None: """Update both device and host leaf sets for a node.""" @@ -1303,7 +1305,7 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface): ) -> None: """Free every component layer on the node and detach it from the LRU lists and evictable leaf sets.""" - self._record_remove_event(node, medium=medium) + self.kv_events.record_remove(node, medium=medium) for comp in self.components: self._evict_component_and_detach_lru( node, @@ -1422,7 +1424,7 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface): """Free only the Full host layer; aux host slices stay under their own pools' LRU (a host-only aux slice may be a sole copy).""" assert self._can_reclaim_full_host_duplicate(node) - self._record_remove_event(node, medium=StorageMedium.CPU) + self.kv_events.record_remove(node, medium=StorageMedium.CPU) self._evict_component_and_detach_lru( node, self.components_by_type[BASE_COMPONENT_TYPE], @@ -1447,7 +1449,7 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface): All freed tokens are accumulated into *tracker*.""" assert self._is_host_leaf(node), f"node {node.id} is not an H-leaf" - self._record_remove_event(node, medium=StorageMedium.CPU) + self.kv_events.record_remove(node, medium=StorageMedium.CPU) for comp in self.components: _, hf = self._evict_component_and_detach_lru( node, @@ -1494,7 +1496,7 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface): self._cascade_evict( node, trigger, tracker, device_frees=device_frees, host_frees=host_frees ) - self._record_remove_event(node, medium=StorageMedium.GPU) + self.kv_events.record_remove(node, medium=StorageMedium.GPU) # after device eviction, insert aux components into host LRU. self._for_each_component_lru( @@ -1980,7 +1982,7 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface): cache_actions=cache_actions, ) for nid in kv_xfer.nodes_to_load or (): - self._record_store_event(self.node_by_id(nid), medium=StorageMedium.GPU) + self.kv_events.record_store(self.node_by_id(nid), medium=StorageMedium.GPU) for ct, xfers in comp_xfers.items(): self.components_by_type[ct].commit_hicache_transfer( node, @@ -2023,7 +2025,7 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface): node.write_through_pending_id = None # The backed-up copy becomes a tracked duplicate only now. self._update_duplicate_tracking(node) - self._record_store_event(node, medium=StorageMedium.CPU) + self.kv_events.record_store(node, medium=StorageMedium.CPU) def set_component_device_value( self, node_id: NodeId, component_type: ComponentType, value: torch.Tensor diff --git a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core_interface.py b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core_interface.py index 367c111d7..b7c63fb55 100644 --- a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core_interface.py +++ b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core_interface.py @@ -12,8 +12,6 @@ from typing import TYPE_CHECKING, Optional, Sequence import msgspec -from sglang.srt.mem_cache.events import KVCacheEventMixin - # Tree node id -- the node handle used outside the TreeCore. The concrete tree # node is a TreeCore-internal type. NodeId = int @@ -95,6 +93,7 @@ if TYPE_CHECKING: MatchPrefixParams, MatchResult, ) + from sglang.srt.mem_cache.events import KVCacheEventRecorder from sglang.srt.mem_cache.hicache_storage import PoolTransfer, PoolTransferResult from sglang.srt.mem_cache.radix_cache import RadixKey from sglang.srt.mem_cache.unified_cache.cache_action import ( @@ -112,11 +111,10 @@ if TYPE_CHECKING: ) -class UnifiedTreeCoreInterface(KVCacheEventMixin, ABC): +class UnifiedTreeCoreInterface(ABC): """Methods the Controller invokes on the Tree Core. The Controller treats the Tree Core as opaque behind this surface, which grows as tree operations - migrate onto the TreeCore. Inherits KVCacheEventMixin for the KV-event API - (take_events, _record_* recorders).""" + migrate onto the TreeCore.""" # ==== Tree-owned state the Controller reads (or, via its facade setters, writes) ==== page_size: int @@ -127,9 +125,14 @@ class UnifiedTreeCoreInterface(KVCacheEventMixin, ABC): write_through_threshold: int is_write_back: bool has_swa_host_pool: bool + kv_events: KVCacheEventRecorder # ==== Tree API ==== + def take_events(self) -> list: + """Hand the queued KV placement events to the Controller.""" + return self.kv_events.take() + @abstractmethod def reset(self) -> None: """Drop the entire tree and reinitialize empty state.""" diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index 3db4fa36e..2f0dfb517 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -322,7 +322,7 @@ class UnifiedRadixCache(BasePrefixCache): self.cache_controller.mem_pool_host.clear() self.enable_storage = self.cache_controller.enable_storage - self.tree_core._record_all_cleared_event() + self.tree_core.kv_events.record_all_cleared() def init_hicache(self, server_args: ServerArgs, params: CacheInitParams) -> None: """Initialize HiCache infrastructure.""" diff --git a/test/registered/unit/mem_cache/test_radix_cache_unit.py b/test/registered/unit/mem_cache/test_radix_cache_unit.py index 3a4656f4d..680b86149 100644 --- a/test/registered/unit/mem_cache/test_radix_cache_unit.py +++ b/test/registered/unit/mem_cache/test_radix_cache_unit.py @@ -45,7 +45,7 @@ from sglang.srt.mem_cache.base_prefix_cache import ( InsertParams, MatchPrefixParams, ) -from sglang.srt.mem_cache.events import KVCacheEventMixin +from sglang.srt.mem_cache.events import KVCacheEventRecorder from sglang.srt.mem_cache.mamba_radix_cache import TreeNode as MambaTreeNode from sglang.srt.mem_cache.radix_cache import RadixCache, RadixKey, TreeNode from sglang.srt.utils import get_device @@ -54,12 +54,6 @@ from sglang.srt.utils import get_device DEFAULT_PAGE_SIZE = 4 -class _KVCacheEventQueue(KVCacheEventMixin): - def __init__(self): - self.enable_kv_cache_events = True - self.kv_event_queue = [] - - class TestKVCacheEventQueue(unittest.TestCase): @staticmethod def _store( @@ -87,26 +81,22 @@ class TestKVCacheEventQueue(unittest.TestCase): ) def test_enqueue_coalesces_compatible_stores(self): - queue = _KVCacheEventQueue() - queue._enqueue_kv_event(self._store(1, None)) - queue._enqueue_kv_event(self._store(2, 1)) + queue = KVCacheEventRecorder(enabled=True, page_size=DEFAULT_PAGE_SIZE) + queue.enqueue(self._store(1, None)) + queue.enqueue(self._store(2, 1)) - events = queue.take_events() + events = queue.take() self.assertEqual(len(events), 1) self.assertEqual(events[0].block_hashes, [1, 2]) self.assertEqual(events[0].parent_block_hash, None) self.assertEqual(events[0].token_ids, [1, 2, 2, 3]) def test_enqueue_coalesces_compatible_removes(self): - queue = _KVCacheEventQueue() - queue._enqueue_kv_event( - BlockRemoved(block_hashes=[1], medium=StorageMedium.GPU) - ) - queue._enqueue_kv_event( - BlockRemoved(block_hashes=[2, 3], medium=StorageMedium.GPU) - ) + queue = KVCacheEventRecorder(enabled=True, page_size=DEFAULT_PAGE_SIZE) + queue.enqueue(BlockRemoved(block_hashes=[1], medium=StorageMedium.GPU)) + queue.enqueue(BlockRemoved(block_hashes=[2, 3], medium=StorageMedium.GPU)) - events = queue.take_events() + events = queue.take() self.assertEqual(len(events), 1) self.assertIsInstance(events[0], BlockRemoved) self.assertEqual(events[0].block_hashes, [1, 2, 3]) @@ -119,33 +109,27 @@ class TestKVCacheEventQueue(unittest.TestCase): self._store(5, None), ] for incoming in incompatible_stores: - queue = _KVCacheEventQueue() - queue._enqueue_kv_event(self._store(1, None)) - queue._enqueue_kv_event(incoming) - self.assertEqual(len(queue.take_events()), 2) + queue = KVCacheEventRecorder(enabled=True, page_size=DEFAULT_PAGE_SIZE) + queue.enqueue(self._store(1, None)) + queue.enqueue(incoming) + self.assertEqual(len(queue.take()), 2) - queue = _KVCacheEventQueue() - queue._enqueue_kv_event(self._store(1, None)) - queue._enqueue_kv_event( - BlockRemoved(block_hashes=[1], medium=StorageMedium.GPU) - ) - queue._enqueue_kv_event(AllBlocksCleared()) - queue._enqueue_kv_event(self._store(2, None)) - self.assertEqual(len(queue.take_events()), 4) + queue = KVCacheEventRecorder(enabled=True, page_size=DEFAULT_PAGE_SIZE) + queue.enqueue(self._store(1, None)) + queue.enqueue(BlockRemoved(block_hashes=[1], medium=StorageMedium.GPU)) + queue.enqueue(AllBlocksCleared()) + queue.enqueue(self._store(2, None)) + self.assertEqual(len(queue.take()), 4) - queue = _KVCacheEventQueue() - queue._enqueue_kv_event( - BlockRemoved(block_hashes=[1], medium=StorageMedium.GPU) - ) - queue._enqueue_kv_event( - BlockRemoved(block_hashes=[2], medium=StorageMedium.CPU) - ) - self.assertEqual(len(queue.take_events()), 2) + queue = KVCacheEventRecorder(enabled=True, page_size=DEFAULT_PAGE_SIZE) + queue.enqueue(BlockRemoved(block_hashes=[1], medium=StorageMedium.GPU)) + queue.enqueue(BlockRemoved(block_hashes=[2], medium=StorageMedium.CPU)) + self.assertEqual(len(queue.take()), 2) - queue = _KVCacheEventQueue() - queue._enqueue_kv_event(self._store(1, None, cache_salt="tenant-a")) - queue._enqueue_kv_event(self._store(2, 1, cache_salt="tenant-b")) - self.assertEqual(len(queue.take_events()), 2) + queue = KVCacheEventRecorder(enabled=True, page_size=DEFAULT_PAGE_SIZE) + queue.enqueue(self._store(1, None, cache_salt="tenant-a")) + queue.enqueue(self._store(2, 1, cache_salt="tenant-b")) + self.assertEqual(len(queue.take()), 2) class TestRadixKey(unittest.TestCase): @@ -410,7 +394,7 @@ class TestRadixCache(unittest.TestCase): self.assertEqual(cache.page_size, page_size) self.assertEqual(cache.disable, disable) - self.assertEqual(cache.enable_kv_cache_events, enable_events) + self.assertEqual(cache.kv_events.enabled, enable_events) self.assertEqual(cache.device, torch.device("cpu")) self.assertIsNotNone(cache.root_node) self.assertEqual(len(cache.root_node.key), 0)