diff --git a/python/sglang/srt/mem_cache/events.py b/python/sglang/srt/mem_cache/events.py new file mode 100644 index 000000000..354ea1daa --- /dev/null +++ b/python/sglang/srt/mem_cache/events.py @@ -0,0 +1,127 @@ +# Copyright 2025 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""KV cache placement event emission mixin. + +The mixin produces the ``BlockStored`` / ``BlockRemoved`` / ``AllBlocksCleared`` +events consumed by KV-aware routers (e.g. dynamo). +""" + +from typing import Any + +from sglang.srt.disaggregation.kv_events import ( + AllBlocksCleared, + BlockRemoved, + BlockStored, + StorageMedium, +) +from sglang.srt.mem_cache.utils import ( + compute_node_hash_values, + hash_str_to_int64, +) + + +class KVCacheEventMixin: + def _record_store_event(self, node: Any, medium=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 + + # 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) + + # 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.parent.hash_value is not None + and len(node.parent.hash_value) > 0 + ): + parent_block_hash = hash_str_to_int64(node.parent.hash_value[-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 = raw[start:end] + + block_hash = hash_str_to_int64(node.hash_value[page_index]) + + self.kv_event_queue.append( + BlockStored( + block_hashes=[block_hash], + parent_block_hash=parent_block_hash, + token_ids=page_tokens, + block_size=len(page_tokens), + lora_id=None, + medium=medium, + ) + ) + + parent_block_hash = block_hash + page_index += 1 + + def _record_remove_event(self, node: Any, medium=None): + # One BlockRemoved per chunk. + # ``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 + + # 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) + + page_index = 0 + logical_len = len(node.key) + for start in range(0, logical_len, self.page_size): + end = min(start + self.page_size, logical_len) + if end <= start: + continue + + block_hash = hash_str_to_int64(node.hash_value[page_index]) + + self.kv_event_queue.append( + BlockRemoved(block_hashes=[block_hash], medium=medium) + ) + + page_index += 1 + + def _record_all_cleared_event(self): + if self.enable_kv_cache_events: + self.kv_event_queue.append(AllBlocksCleared()) + + def take_events(self): + """Atomically takes all events and clears the queue. + + Returns: + A list of KV cache events. + """ + if not self.enable_kv_cache_events: + return [] + events = self.kv_event_queue + self.kv_event_queue = [] + return events diff --git a/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py b/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py index 0a3ae9160..5ec990081 100644 --- a/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py +++ b/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py @@ -12,6 +12,7 @@ from typing import TYPE_CHECKING, Dict, List, Optional, Tuple import torch +from sglang.srt.disaggregation.kv_events import StorageMedium from sglang.srt.mem_cache.base_prefix_cache import ( DecLockRefParams, DecLockRefResult, @@ -38,9 +39,8 @@ from sglang.srt.mem_cache.mamba_radix_cache import ( from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool, HybridReqToTokenPool from sglang.srt.mem_cache.radix_cache import ( RadixKey, - compute_node_hash_values, - split_node_hash_value, ) +from sglang.srt.mem_cache.utils import compute_node_hash_values, split_node_hash_value from sglang.srt.observability.metrics_collector import StorageMetricsCollector if TYPE_CHECKING: @@ -319,6 +319,7 @@ class HiMambaRadixCache(MambaRadixCache): n.value = full_device_indices[offset : offset + n_len].clone() offset += n_len + self._record_store_event(n, medium=StorageMedium.GPU) self.full_lru_list.insert_mru(n) self.full_evictable_size_ += n_len self._update_leaf_status(n) @@ -379,6 +380,9 @@ class HiMambaRadixCache(MambaRadixCache): finish_event.synchronize() for ack_id in ack_list: backuped_node = self.ongoing_write_through.pop(ack_id) + self._record_store_event( + backuped_node, medium=StorageMedium.CPU + ) if self.enable_storage: self.write_backup_storage(backuped_node) self.cache_controller.ack_write_queue.clear() @@ -408,6 +412,7 @@ class HiMambaRadixCache(MambaRadixCache): finish_event.synchronize() for ack_id in ack_list: backuped_node = self.ongoing_write_through.pop(ack_id) + self._record_store_event(backuped_node, medium=StorageMedium.CPU) self.dec_lock_ref(backuped_node) if self.enable_storage: self.write_backup_storage(backuped_node) @@ -514,6 +519,7 @@ class HiMambaRadixCache(MambaRadixCache): num_full = len(node.value) + self._record_remove_event(node, medium=StorageMedium.GPU) self.cache_controller.evict_device(node.value) self.full_evictable_size_ -= num_full if self.full_lru_list.in_list(node): @@ -534,6 +540,7 @@ class HiMambaRadixCache(MambaRadixCache): full_num_evicted = len(node.value) + self._record_remove_event(node, medium=StorageMedium.GPU) self.cache_controller.evict_device(node.value) self.full_evictable_size_ -= full_num_evicted if self.full_lru_list.in_list(node): @@ -576,6 +583,7 @@ class HiMambaRadixCache(MambaRadixCache): node.host_mamba_ref_counter == 0 ), f"host mamba in use, {node.id=} {node.host_mamba_ref_counter=}" + self._record_remove_event(node, medium=StorageMedium.CPU) full_num_evicted = self.cache_controller.evict_host(node.host_value) node.host_value = None @@ -612,6 +620,7 @@ class HiMambaRadixCache(MambaRadixCache): and node.host_ref_counter == 0 and node.host_mamba_ref_counter == 0 ): + self._record_remove_event(node, medium=StorageMedium.CPU) self.cache_controller.evict_host(node.host_value) node.host_value = None @@ -641,6 +650,7 @@ class HiMambaRadixCache(MambaRadixCache): parent = node.parent if not parent.evicted: + self._record_remove_event(parent, medium=StorageMedium.GPU) full_num_evicted += len(parent.value) self.full_evictable_size_ -= len(parent.value) self.cache_controller.evict_device(parent.value) @@ -808,6 +818,7 @@ class HiMambaRadixCache(MambaRadixCache): node.value = fresh_value.clone() self.full_lru_list.insert_mru(node) self.full_evictable_size_ += n + self._record_store_event(node, medium=StorageMedium.GPU) self._update_leaf_status(node) if node.parent is not None: @@ -903,8 +914,9 @@ class HiMambaRadixCache(MambaRadixCache): parent.children[child_key] = new_node self.full_evictable_size_ += len(value) self.mamba_evictable_size_ += len(mamba_value) - if self.enable_storage: + if self.enable_storage or self.enable_kv_cache_events: new_node.hash_value = compute_node_hash_values(new_node, self.page_size) + self._record_store_event(new_node, medium=StorageMedium.GPU) self._update_full_device_leaf_status(new_node) self._update_full_device_leaf_status(parent) return new_node @@ -1067,10 +1079,6 @@ class HiMambaRadixCache(MambaRadixCache): new_node.host_value = child.host_value[:split_len].clone() child.host_value = child.host_value[split_len:].clone() - new_node.hash_value, child.hash_value = split_node_hash_value( - child.hash_value, split_len, self.page_size - ) - self._update_leaf_status(new_node) self._update_leaf_status(child) @@ -1879,6 +1887,7 @@ class HiMambaRadixCache(MambaRadixCache): leaf_node = new_node self._update_full_host_leaf_status(new_node) self._update_full_host_leaf_status(node) + self._record_store_event(new_node, medium=StorageMedium.CPU) # Attach mamba state to the new leaf if leaf_node is not None and mamba_host_value is not None and mamba_loaded: diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index e1f6b2b14..febe60ae1 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -50,6 +50,8 @@ from sglang.srt.mem_cache.radix_cache import ( RadixCache, RadixKey, TreeNode, +) +from sglang.srt.mem_cache.utils import ( compute_node_hash_values, split_node_hash_value, ) diff --git a/python/sglang/srt/mem_cache/mamba_radix_cache.py b/python/sglang/srt/mem_cache/mamba_radix_cache.py index 1f2c2017e..55ee7983e 100644 --- a/python/sglang/srt/mem_cache/mamba_radix_cache.py +++ b/python/sglang/srt/mem_cache/mamba_radix_cache.py @@ -45,8 +45,10 @@ from sglang.srt.mem_cache.base_prefix_cache import ( MatchPrefixParams, MatchResult, ) +from sglang.srt.mem_cache.events import KVCacheEventMixin from sglang.srt.mem_cache.memory_pool import HybridReqToTokenPool from sglang.srt.mem_cache.radix_cache import RadixKey +from sglang.srt.mem_cache.utils import split_node_hash_value from sglang.srt.server_args import get_global_server_args if TYPE_CHECKING: @@ -415,7 +417,7 @@ class LRUList: raise Exception(msg) -class MambaRadixCache(BasePrefixCache): +class MambaRadixCache(KVCacheEventMixin, BasePrefixCache): def __init__(self, params: CacheInitParams): assert isinstance( params.token_to_kv_pool_allocator, TokenToKVPoolAllocator @@ -425,7 +427,9 @@ class MambaRadixCache(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.kv_event_queue = [] if not self.enable_mamba_extra_buffer: assert ( @@ -451,6 +455,7 @@ class MambaRadixCache(BasePrefixCache): self.root_node = TreeNode() self.root_node.key = RadixKey([], None) self.root_node.value = [] + self.root_node.hash_value = [] self.root_node.full_lock_ref = 1 self.root_node.mamba_lock_ref = 1 self.full_evictable_size_ = 0 @@ -460,6 +465,7 @@ class MambaRadixCache(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() def match_prefix(self, params: MatchPrefixParams) -> MatchResult: """Find the matching prefix from the radix tree. @@ -725,6 +731,7 @@ class MambaRadixCache(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.token_to_kv_pool_allocator.free(x.value) full_num_evicted = len(x.value) self.req_to_token_pool.mamba_pool.free(x.mamba_value) @@ -1089,6 +1096,9 @@ class MambaRadixCache(BasePrefixCache): child.key = child.key[split_len:] child.value = child.value[split_len:].clone() new_node.parent.children[key.child_key(self.page_size)] = new_node + new_node.hash_value, child.hash_value = split_node_hash_value( + child.hash_value, split_len, self.page_size + ) # insert the new node and child into the lru lists, insert # parent first so that parent is after child in the lru list @@ -1156,6 +1166,7 @@ class MambaRadixCache(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) elif node.mamba_value is None: # add for mamba tombstone node.mamba_value = mamba_value self.full_lru_list.reset_node_mru(node) @@ -1185,6 +1196,7 @@ class MambaRadixCache(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.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/radix_cache.py b/python/sglang/srt/mem_cache/radix_cache.py index 7f9eca81f..8a24c5e15 100644 --- a/python/sglang/srt/mem_cache/radix_cache.py +++ b/python/sglang/srt/mem_cache/radix_cache.py @@ -34,12 +34,6 @@ import torch logger = logging.getLogger(__name__) -from sglang.srt.disaggregation.kv_events import ( - AllBlocksCleared, - BlockRemoved, - BlockStored, - StorageMedium, -) from sglang.srt.mem_cache.base_prefix_cache import ( BasePrefixCache, DecLockRefParams, @@ -52,6 +46,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.evict_policy import ( EvictionStrategy, FIFOStrategy, @@ -62,7 +57,7 @@ from sglang.srt.mem_cache.evict_policy import ( PriorityStrategy, SLRUStrategy, ) -from sglang.srt.mem_cache.utils import hash_str_to_int64 +from sglang.srt.mem_cache.utils import split_node_hash_value if TYPE_CHECKING: from sglang.srt.managers.schedule_batch import Req @@ -271,54 +266,7 @@ class TreeNode: return self.last_access_time < other.last_access_time -def compute_node_hash_values(node: "TreeNode", page_size: int) -> List[str]: - """Compute SHA256-based hash values for position-aware identification.""" - hash_values = [] - - parent_hash = None - if node.parent is not None and node.parent.hash_value is not None: - if len(node.parent.key) > 0 and len(node.parent.hash_value) > 0: - parent_hash = node.parent.hash_value[-1] - - logical_len = len(node.key) - for start in range(0, logical_len, page_size): - end = min(start + page_size, logical_len) - if end <= start: - continue - hash_val = node.key.hash_page(start, end, parent_hash) - hash_values.append(hash_val) - parent_hash = hash_val - return hash_values - - -def split_node_hash_value( - child_hash_value: Optional[List[str]], split_len: int, page_size: int -) -> tuple[Optional[List[str]], Optional[List[str]]]: - """Split hash_value between parent and child nodes during node splitting. - - Args: - child_hash_value: The hash_value list from the child node being split - split_len: The length at which to split (in tokens) - page_size: The page size for calculating number of pages - - Returns: - Tuple of (new_node_hash_value, updated_child_hash_value) - """ - if child_hash_value is None: - return None, None - - if page_size == 1: - split_pages = split_len - else: - split_pages = split_len // page_size - - new_node_hash = child_hash_value[:split_pages] - child_hash = child_hash_value[split_pages:] - - return new_node_hash, child_hash - - -class RadixCache(BasePrefixCache): +class RadixCache(KVCacheEventMixin, BasePrefixCache): def __init__(self, params: CacheInitParams): self.disable = params.disable self.req_to_token_pool = params.req_to_token_pool @@ -857,100 +805,6 @@ class RadixCache(BasePrefixCache): stack.append(child) return total_size - def _record_store_event(self, node: TreeNode, medium=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 - - # 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) - - # 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.parent.hash_value is not None - and len(node.parent.hash_value) > 0 - ): - parent_block_hash = hash_str_to_int64(node.parent.hash_value[-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 = raw[start:end] - - block_hash = hash_str_to_int64(node.hash_value[page_index]) - - self.kv_event_queue.append( - BlockStored( - block_hashes=[block_hash], - parent_block_hash=parent_block_hash, - token_ids=page_tokens, - block_size=len(page_tokens), - lora_id=None, - medium=medium, - ) - ) - - parent_block_hash = block_hash - page_index += 1 - - def _record_remove_event(self, node: TreeNode, medium=None): - # One BlockRemoved per chunk. - # ``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 - - # 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) - - page_index = 0 - logical_len = len(node.key) - for start in range(0, logical_len, self.page_size): - end = min(start + self.page_size, logical_len) - if end <= start: - continue - - block_hash = hash_str_to_int64(node.hash_value[page_index]) - - self.kv_event_queue.append( - BlockRemoved(block_hashes=[block_hash], medium=medium) - ) - - page_index += 1 - - def _record_all_cleared_event(self): - if self.enable_kv_cache_events: - self.kv_event_queue.append(AllBlocksCleared()) - - def take_events(self): - """Atomically takes all events and clears the queue. - - Returns: - A list of KV cache events. - """ - if not self.enable_kv_cache_events: - return [] - events = self.kv_event_queue - self.kv_event_queue = [] - return events - if __name__ == "__main__": tree = RadixCache.create_simulated() diff --git a/python/sglang/srt/mem_cache/utils.py b/python/sglang/srt/mem_cache/utils.py index 66012c4e7..65b7b165c 100644 --- a/python/sglang/srt/mem_cache/utils.py +++ b/python/sglang/srt/mem_cache/utils.py @@ -389,3 +389,50 @@ def hash_str_to_int64(hash_str: str) -> int: if uint64_val >= 2**63: return uint64_val - 2**64 return uint64_val + + +def compute_node_hash_values(node: Any, page_size: int) -> List[str]: + """Compute SHA256-based hash values for position-aware KV block IDs.""" + hash_values = [] + + parent_hash = None + if node.parent is not None and node.parent.hash_value is not None: + if len(node.parent.key) > 0 and len(node.parent.hash_value) > 0: + parent_hash = node.parent.hash_value[-1] + + logical_len = len(node.key) + for start in range(0, logical_len, page_size): + end = min(start + page_size, logical_len) + if end <= start: + continue + hash_val = node.key.hash_page(start, end, parent_hash) + hash_values.append(hash_val) + parent_hash = hash_val + return hash_values + + +def split_node_hash_value( + child_hash_value: Optional[List[str]], split_len: int, page_size: int +) -> tuple[Optional[List[str]], Optional[List[str]]]: + """Split hash_value between parent and child nodes during node splitting. + + Args: + child_hash_value: The hash_value list from the child node being split + split_len: The length at which to split (in tokens) + page_size: The page size for calculating number of pages + + Returns: + Tuple of (new_node_hash_value, updated_child_hash_value) + """ + if child_hash_value is None: + return None, None + + if page_size == 1: + split_pages = split_len + else: + split_pages = split_len // page_size + + new_node_hash = child_hash_value[:split_pages] + child_hash = child_hash_value[split_pages:] + + return new_node_hash, child_hash diff --git a/test/registered/4-gpu-models/test_qwen35_hicache.py b/test/registered/4-gpu-models/test_qwen35_hicache.py index 66b6cd9f3..5dc05062e 100644 --- a/test/registered/4-gpu-models/test_qwen35_hicache.py +++ b/test/registered/4-gpu-models/test_qwen35_hicache.py @@ -1,10 +1,14 @@ import shutil import tempfile +import time import unittest from types import SimpleNamespace import requests +import zmq +from msgspec.msgpack import Decoder +from sglang.srt.disaggregation.kv_events import BlockStored, KVEventBatch from sglang.srt.utils import kill_process_tree from sglang.test.ci.ci_register import register_cuda_ci @@ -72,6 +76,8 @@ class TestQwen35WithHiCache(CustomTestCase): "file", "--hicache-storage-prefetch-policy", "wait_complete", + "--kv-events-config", + '{"publisher": "zmq", "topic": "kv-events"}', ], ) @@ -125,6 +131,42 @@ class TestQwen35WithHiCache(CustomTestCase): f"first={first_metrics['score']}, second={second_metrics['score']}", ) + def test_kv_events_smoke(self): + decoder = Decoder(type=KVEventBatch) + context = zmq.Context() + sub = context.socket(zmq.SUB) + sub.connect("tcp://localhost:5557") + sub.setsockopt_string(zmq.SUBSCRIBE, "kv-events") + + try: + time.sleep(1.0) + res = requests.post( + f"{self.base_url}/generate", + json={ + "text": "HiCache KV event compatibility check. " * 64, + "sampling_params": {"temperature": 0, "max_new_tokens": 1}, + }, + timeout=120, + ) + res.raise_for_status() + + events = [] + deadline = time.time() + 10 + while time.time() < deadline and not any( + isinstance(event, BlockStored) for event in events + ): + if sub.poll(timeout=100): + _, _, payload = sub.recv_multipart() + events.extend(decoder.decode(payload).events) + + self.assertTrue( + any(isinstance(event, BlockStored) for event in events), + "Expected at least one BlockStored event from Qwen3.5 HiCache server", + ) + finally: + sub.close() + context.term() + if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/mem_cache/test_mamba_unittest.py b/test/registered/unit/mem_cache/test_mamba_unittest.py index 4d4450ea5..8ea1db755 100755 --- a/test/registered/unit/mem_cache/test_mamba_unittest.py +++ b/test/registered/unit/mem_cache/test_mamba_unittest.py @@ -3,6 +3,7 @@ import unittest import torch from sglang.srt.configs.mamba_utils import Mamba2CacheParams, Mamba2StateShape +from sglang.srt.disaggregation.kv_events import BlockRemoved, BlockStored from sglang.srt.environ import envs from sglang.srt.managers.schedule_batch import Req from sglang.srt.mem_cache.allocator import TokenToKVPoolAllocator @@ -305,7 +306,98 @@ class TestMamba(unittest.TestCase): print(available_and_evictable_str(tree)) tree.sanity_check() - def _setup_tree_and_allocator(self): + def test_mamba_radix_cache_kv_events(self): + tree, allocator, _, make_dummy_req = self._setup_tree_and_allocator( + enable_kv_cache_events=True + ) + tree.take_events() # Clear the reset event. + + stored_hashes = [] + + req1 = make_dummy_req() + key1 = RadixKey([1, 2, 3]) + tree.insert( + InsertParams( + key=key1, + value=allocator.alloc(3)[: len(key1)], + mamba_value=req1.mamba_pool_idx.unsqueeze(0), + ) + ) + events = tree.take_events() + stored_events = [e for e in events if isinstance(e, BlockStored)] + self.assertEqual(len(stored_events), 3) + self.assertEqual([e.token_ids[0] for e in stored_events], [1, 2, 3]) + stored_hashes.extend(e.block_hashes[0] for e in stored_events) + + req2 = make_dummy_req() + key2 = RadixKey([1, 2, 3, 4, 5]) + tree.insert( + InsertParams( + key=key2, + value=allocator.alloc(5)[: len(key2)], + mamba_value=req2.mamba_pool_idx.unsqueeze(0), + ) + ) + events = tree.take_events() + stored_events = [e for e in events if isinstance(e, BlockStored)] + self.assertEqual(len(stored_events), 2) + self.assertEqual([e.token_ids[0] for e in stored_events], [4, 5]) + stored_hashes.extend(e.block_hashes[0] for e in stored_events) + + # Evicting an internal mamba state creates a tombstone but does not + # remove full-attention KV blocks, so it must not emit BlockRemoved. + result = tree.evict(EvictParams(num_tokens=0, mamba_num=1)) + self.assertEqual(result.num_tokens_evicted, 0) + self.assertEqual(result.mamba_num_evicted, 1) + events = tree.take_events() + self.assertEqual([e for e in events if isinstance(e, BlockRemoved)], []) + + result = tree.evict(EvictParams(num_tokens=1)) + self.assertGreaterEqual(result.num_tokens_evicted, 1) + events = tree.take_events() + removed_hashes = [ + e.block_hashes[0] for e in events if isinstance(e, BlockRemoved) + ] + self.assertCountEqual(removed_hashes, stored_hashes) + + def test_mamba_radix_cache_kv_events_split_hash(self): + tree, allocator, _, make_dummy_req = self._setup_tree_and_allocator( + enable_kv_cache_events=True + ) + tree.take_events() # Clear the reset event. + + req1 = make_dummy_req() + key1 = RadixKey([1, 2, 3, 4]) + tree.insert( + InsertParams( + key=key1, + value=allocator.alloc(4)[: len(key1)], + mamba_value=req1.mamba_pool_idx.unsqueeze(0), + ) + ) + first_insert_events = [ + e for e in tree.take_events() if isinstance(e, BlockStored) + ] + self.assertEqual(len(first_insert_events), 4) + split_parent_hash = first_insert_events[1].block_hashes[0] + + req2 = make_dummy_req() + key2 = RadixKey([1, 2, 5, 6]) + tree.insert( + InsertParams( + key=key2, + value=allocator.alloc(4)[: len(key2)], + mamba_value=req2.mamba_pool_idx.unsqueeze(0), + ) + ) + second_insert_events = [ + e for e in tree.take_events() if isinstance(e, BlockStored) + ] + self.assertEqual(len(second_insert_events), 2) + self.assertEqual(second_insert_events[0].token_ids, [5]) + self.assertEqual(second_insert_events[0].parent_block_hash, split_parent_hash) + + def _setup_tree_and_allocator(self, enable_kv_cache_events=False): """Helper to create a MambaRadixCache with allocator for testing.""" set_global_server_args_for_scheduler( ServerArgs(model_path="dummy", page_size=1) @@ -374,6 +466,7 @@ class TestMamba(unittest.TestCase): token_to_kv_pool_allocator=allocator, page_size=1, disable=False, + enable_kv_cache_events=enable_kv_cache_events, ) tree = MambaRadixCache(params=params)