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

feat: session radix cache (#27058)

Browse Source 
Signed-off-by: Ishan Dhanani <ishandhanani@gmail.com>
This commit is contained in:
ishandhanani
2026-06-23 01:51:41 -07:00
committed by GitHub
parent 349a6af6b8
commit e67b228d4c
7 changed files with 373 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/manual/core/test_session_radix_cache.py
+162
View File
@@ -0,0 +1,162 @@
"""Manual test for the session radix cache (--enable-session-radix-cache).
Run directly: python test/manual/core/test_session_radix_cache.py
"""
import unittest
from array import array
from types import SimpleNamespace
import torch
from sglang.srt.mem_cache.allocator import TokenToKVPoolAllocator
from sglang.srt.mem_cache.base_prefix_cache import (
EvictParams,
InsertParams,
MatchPrefixParams,
)
from sglang.srt.mem_cache.cache_init_params import CacheInitParams
from sglang.srt.mem_cache.memory_pool import MHATokenToKVPool, ReqToTokenPool
from sglang.srt.mem_cache.radix_cache import RadixCache, RadixKey
class TestSessionRadixCache(unittest.TestCase):
def setUp(self):
dtype = torch.float16
kv = MHATokenToKVPool(
size=64,
page_size=1,
dtype=dtype,
head_num=2,
head_dim=8,
layer_num=1,
device="cpu",
enable_memory_saver=False,
)
allocator = TokenToKVPoolAllocator(
size=64, dtype=dtype, device="cpu", kvcache=kv, need_sort=False
)
req_to_token_pool = ReqToTokenPool(
size=8, max_context_len=1024, device="cpu", enable_memory_saver=False
)
self.cache = RadixCache(
CacheInitParams(
disable=False,
req_to_token_pool=req_to_token_pool,
token_to_kv_pool_allocator=allocator,
page_size=1,
eviction_policy="lru",
enable_kv_cache_events=False,
)
)
def _insert(self, toks):
idx = self.cache.token_to_kv_pool_allocator.alloc(len(toks))
self.cache.insert(
InsertParams(key=RadixKey(array("q", toks)), value=idx.to(torch.int64))
)
def _tag(self, toks, sid):
self.cache._tag_session_leaf(
SimpleNamespace(session_id=sid),
RadixKey(array("q", toks)),
node=self._leaf(toks),
)
def _cached(self, toks):
return int(
self.cache.match_prefix(
MatchPrefixParams(key=RadixKey(array("q", toks)))
).device_indices.numel()
)
def _leaf(self, toks):
return self.cache.match_prefix(
MatchPrefixParams(key=RadixKey(array("q", toks)))
).last_device_node
def test_tag_with_known_node_skips_match_prefix(self):
self._insert([1, 2, 3, 4])
leaf = self._leaf([1, 2, 3, 4])
orig_match_prefix = self.cache.match_prefix
def fail_match_prefix(_params):
raise AssertionError("match_prefix should not run when node is supplied")
self.cache.match_prefix = fail_match_prefix
try:
self.cache._tag_session_leaf(
SimpleNamespace(session_id="S"),
RadixKey(array("q", [1, 2, 3, 4])),
node=leaf,
)
finally:
self.cache.match_prefix = orig_match_prefix
self.assertEqual(getattr(leaf, "session_ids", None), {"S"})
self.assertIn(leaf, self.cache._session_leaves["S"])
def test_shared_prefix_frees_only_unique_tail(self):
# A/B share prefix [1,2]; close(A) frees only A's tail, B + shared stay.
self._insert([1, 2, 3, 4])
self._tag([1, 2, 3, 4], "A")
self._insert([1, 2, 5, 6])
self._tag([1, 2, 5, 6], "B")
self.assertGreater(self.cache.release_session("A"), 0)
self.assertEqual(self._cached([1, 2, 3, 4]), 2) # only shared [1,2] left
self.assertEqual(self._cached([1, 2, 5, 6]), 4) # B intact
def test_same_leaf_freed_only_on_last_holder(self):
# Identical content -> one leaf held by {A,B}; freed only on last close.
self._insert([1, 2, 3, 4])
self._tag([1, 2, 3, 4], "A")
self._insert([1, 2, 3, 4])
self._tag([1, 2, 3, 4], "B")
self.assertEqual(
getattr(self._leaf([1, 2, 3, 4]), "session_ids", None), {"A", "B"}
)
self.assertEqual(self.cache.release_session("A"), 0) # B still holds
self.assertEqual(self._cached([1, 2, 3, 4]), 4)
self.assertEqual(self.cache.release_session("B"), 1) # last holder frees
self.assertEqual(self._cached([1, 2, 3, 4]), 0)
def test_tag_is_lru_neutral_not_pinned(self):
# The tag must add no lock/pin: a tagged, never-closed node is evictable.
self._insert([1, 2, 3, 4])
self._tag([1, 2, 3, 4], "S")
leaf = self._leaf([1, 2, 3, 4])
self.assertEqual(leaf.lock_ref, 0)
self.assertEqual(self.cache.protected_size(), 0)
self.assertIn(leaf, self.cache.evictable_leaves)
self.cache.evict(EvictParams(num_tokens=4)) # LRU reclaims it while open
self.assertEqual(self._cached([1, 2, 3, 4]), 0)
self.assertNotIn("S", self.cache._session_leaves)
self.assertEqual(self.cache.release_session("S"), 0) # late close is a no-op
def test_close_tombstone_blocks_late_finish_until_reopen(self):
self._insert([1, 2, 3, 4])
self._tag([1, 2, 3, 4], "S")
self.assertEqual(self.cache.release_session("S"), 1)
self._insert([5, 6, 7, 8])
self._tag([5, 6, 7, 8], "S") # simulates a finish racing after close
self.assertIsNone(getattr(self._leaf([5, 6, 7, 8]), "session_ids", None))
self.cache.register_session("S")
self._tag([5, 6, 7, 8], "S")
self.assertEqual(getattr(self._leaf([5, 6, 7, 8]), "session_ids", None), {"S"})
def test_tombstoned_shared_holder_cannot_retag_after_last_holder_close(self):
self._insert([1, 2, 3, 4])
self._tag([1, 2, 3, 4], "A")
self._tag([1, 2, 3, 4], "B")
self.assertEqual(self.cache.release_session("B"), 0)
self.assertEqual(getattr(self._leaf([1, 2, 3, 4]), "session_ids", None), {"A"})
self.assertEqual(self.cache.release_session("A"), 1)
self._insert([5, 6, 7, 8])
self._tag([5, 6, 7, 8], "B")
self.assertIsNone(getattr(self._leaf([5, 6, 7, 8]), "session_ids", None))
if __name__ == "__main__":
unittest.main()
test/registered/unit/server_args/test_server_args.py
+10
View File
@@ -1008,6 +1008,16 @@ class TestPrefillOnlyDisableKvCache(unittest.TestCase):
ServerArgs(**self._base_kwargs(kv_cache_dtype="fp4_e2m1"))
class TestSessionRadixCacheServerArgs(unittest.TestCase):
def test_requires_priority_radix_eviction_policy(self):
with self.assertRaisesRegex(ValueError, "--radix-eviction-policy priority"):
ServerArgs(
model_path="dummy",
enable_session_radix_cache=True,
radix_eviction_policy="lru",
)
class TestCudaGraphConfigDataclassAccess(CustomTestCase):
@patch(
"sglang.srt.model_executor.runner_backend."