minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

feat(kv-events): publish SWA radix cache events (#24718)

Browse Source 
Signed-off-by: PeaBrane <yanrpei@gmail.com>
This commit is contained in:
Yan Ru Pei
2026-05-11 10:31:45 -07:00
committed by GitHub
parent 95985f983d
commit 4b6f776889
2 changed files with 84 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/mem_cache/swa_radix_cache.py
+21 -2
View File
@@ -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)
test/registered/unit/mem_cache/test_swa_unittest.py
+63
View File
@@ -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