From 4b6f7768894b5707d0ab19d6b727950096f15f75 Mon Sep 17 00:00:00 2001 From: Yan Ru Pei Date: Mon, 11 May 2026 10:31:45 -0700 Subject: [PATCH] feat(kv-events): publish SWA radix cache events (#24718) Signed-off-by: PeaBrane --- .../sglang/srt/mem_cache/swa_radix_cache.py | 23 ++++++- .../unit/mem_cache/test_swa_unittest.py | 63 +++++++++++++++++++ 2 files changed, 84 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/mem_cache/swa_radix_cache.py b/python/sglang/srt/mem_cache/swa_radix_cache.py index d22631479..ed436b7e1 100644 --- a/python/sglang/srt/mem_cache/swa_radix_cache.py +++ b/python/sglang/srt/mem_cache/swa_radix_cache.py @@ -41,9 +41,10 @@ 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.radix_cache import RadixKey from sglang.srt.mem_cache.swa_memory_pool import SWATokenToKVPoolAllocator -from sglang.srt.mem_cache.utils import convert_to_bigram_key +from sglang.srt.mem_cache.utils import convert_to_bigram_key, split_node_hash_value if TYPE_CHECKING: from sglang.srt.managers.schedule_batch import Req @@ -77,6 +78,8 @@ class TreeNode: self.hit_count = 0 # store the host indices of KV cache self.host_value = None + # store hash values of each page + self.hash_value: Optional[List[str]] = None # for lru list, invariant: # 1. prev has greater last_access_time @@ -334,7 +337,7 @@ class LRUList: raise Exception(msg) -class SWARadixCache(BasePrefixCache): +class SWARadixCache(KVCacheEventMixin, 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 @@ -342,6 +345,8 @@ class SWARadixCache(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 = [] if self.token_to_kv_pool_allocator: self.device = self.token_to_kv_pool_allocator.device @@ -371,6 +376,7 @@ class SWARadixCache(BasePrefixCache): self.root_node = TreeNode() self.root_node.key = [] self.root_node.value = [] + self.root_node.hash_value = [] self.root_node.full_lock_ref = 1 self.root_node.swa_lock_ref = 1 self.full_evictable_size_ = 0 @@ -380,6 +386,7 @@ class SWARadixCache(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() def match_prefix(self, params: MatchPrefixParams) -> MatchResult: """Find the matching prefix from the radix tree. @@ -573,6 +580,7 @@ class SWARadixCache(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.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` @@ -637,6 +645,7 @@ class SWARadixCache(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.token_to_kv_pool_allocator.free(x.value) full_num_evicted += len(x.value) swa_num_evicted += len(x.value) @@ -995,6 +1004,11 @@ class SWARadixCache(BasePrefixCache): if child.swa_uuid is not None: node.swa_uuid = child.swa_uuid + if node.hash_value is not None and child.hash_value is not None: + node.hash_value = list(node.hash_value) + list(child.hash_value) + else: + node.hash_value = None + self.full_lru_list.remove_node(child) if not child.swa_tombstone: self.swa_lru_list.remove_node(child) @@ -1062,6 +1076,9 @@ class SWARadixCache(BasePrefixCache): assert len(child.key) > 0, f"child.key should not be empty" 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 @@ -1219,6 +1236,7 @@ class SWARadixCache(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) return new_node def _iteratively_delete_tombstone_leaf( @@ -1236,6 +1254,7 @@ class SWARadixCache(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.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/test/registered/unit/mem_cache/test_swa_unittest.py b/test/registered/unit/mem_cache/test_swa_unittest.py index 154a869a5..417b798b6 100644 --- a/test/registered/unit/mem_cache/test_swa_unittest.py +++ b/test/registered/unit/mem_cache/test_swa_unittest.py @@ -2,6 +2,7 @@ import unittest import torch +from sglang.srt.disaggregation.kv_events import BlockRemoved, BlockStored from sglang.srt.environ import envs from sglang.srt.mem_cache.base_prefix_cache import ( DecLockRefParams, @@ -41,6 +42,7 @@ def _build_swa_tree( kv_size: int = 64, kv_size_swa: int = 32, sliding_window_size: int = 4, + enable_kv_cache_events: bool = False, ): head_num = 8 head_dim = 128 @@ -89,6 +91,7 @@ def _build_swa_tree( disable=False, is_eagle=is_eagle, sliding_window_size=sliding_window_size, + enable_kv_cache_events=enable_kv_cache_events, ), ) return tree, allocator, req_to_token_pool @@ -133,6 +136,66 @@ class TestSWA(unittest.TestCase): def tearDownClass(cls): pass + def test_swa_radix_cache_kv_events(self): + tree, allocator, _ = _build_swa_tree( + is_eagle=False, enable_kv_cache_events=True + ) + tree.take_events() # Clear the reset event. + + _insert(tree, allocator, [1, 2, 3, 4]) + first_insert_events = [ + e for e in tree.take_events() if isinstance(e, BlockStored) + ] + self.assertEqual(len(first_insert_events), 4) + self.assertEqual([e.token_ids[0] for e in first_insert_events], [1, 2, 3, 4]) + + _insert(tree, allocator, [1, 2, 3, 4, 5, 6]) + second_insert_events = [ + e for e in tree.take_events() if isinstance(e, BlockStored) + ] + self.assertEqual(len(second_insert_events), 2) + self.assertEqual([e.token_ids[0] for e in second_insert_events], [5, 6]) + + stored_hashes = [ + e.block_hashes[0] for e in first_insert_events + second_insert_events + ] + + # Evicting only SWA tokens tombstones nodes but keeps full KV blocks. + result = tree.evict(EvictParams(num_tokens=0, swa_num_tokens=1)) + self.assertEqual(result.num_tokens_evicted, 0) + self.assertGreaterEqual(result.swa_num_tokens_evicted, 1) + self.assertEqual( + [e for e in tree.take_events() if isinstance(e, BlockRemoved)], [] + ) + + result = tree.evict(EvictParams(num_tokens=1, swa_num_tokens=0)) + self.assertGreaterEqual(result.num_tokens_evicted, 1) + removed_hashes = [ + e.block_hashes[0] for e in tree.take_events() if isinstance(e, BlockRemoved) + ] + self.assertCountEqual(removed_hashes, stored_hashes) + + def test_swa_radix_cache_kv_events_split_hash(self): + tree, allocator, _ = _build_swa_tree( + is_eagle=False, enable_kv_cache_events=True + ) + tree.take_events() # Clear the reset event. + + _insert(tree, allocator, [1, 2, 3, 4]) + 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] + + _insert(tree, allocator, [1, 2, 5, 6]) + 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 test_swa_memory_pool(self): size = 16 size_swa = 16