feat: Session-reference-aware Unified Radix Cache for agentic multi-turn workloads (#29173)
Co-authored-by: Ishan Dhanani <ishandhanani@gmail.com> Co-authored-by: hzh0425 <hzh0425@apache.org> Co-authored-by: ispobock <ispobaoke@gmail.com>
This commit is contained in:
Mingjun Zhang
co-authored by
Ishan Dhanani
hzh0425
ispobock
Ishan Dhanani
hzh0425
ispobock
parent
131bd51b01
commit
056474cdb0
@@ -1,174 +0,0 @@
|
||||
"""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,
|
||||
enable_session_radix_cache=True,
|
||||
)
|
||||
)
|
||||
|
||||
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_disabled_cache_does_not_tag_session_kv(self):
|
||||
self.cache.enable_session_radix_cache = False
|
||||
self._insert([1, 2, 3, 4])
|
||||
self._tag([1, 2, 3, 4], "S")
|
||||
self.assertIsNone(getattr(self._leaf([1, 2, 3, 4]), "session_ids", None))
|
||||
|
||||
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_radix_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_radix_session("A"), 0) # B still holds
|
||||
self.assertEqual(self._cached([1, 2, 3, 4]), 4)
|
||||
self.assertEqual(self.cache.release_radix_session("B"), 1) # last holder frees
|
||||
self.assertEqual(self._cached([1, 2, 3, 4]), 0)
|
||||
|
||||
def test_legacy_release_does_not_release_radix_session(self):
|
||||
self._insert([1, 2, 3, 4])
|
||||
self._tag([1, 2, 3, 4], "S")
|
||||
self.cache.release_session("S")
|
||||
self.assertEqual(self._cached([1, 2, 3, 4]), 4)
|
||||
self.assertEqual(self.cache.release_radix_session("S"), 1)
|
||||
|
||||
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_radix_session("S"), 0
|
||||
) # late close is a no-op
|
||||
|
||||
def test_close_tombstone_blocks_late_finish(self):
|
||||
self._insert([1, 2, 3, 4])
|
||||
self._tag([1, 2, 3, 4], "S")
|
||||
self.assertEqual(self.cache.release_radix_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))
|
||||
|
||||
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_radix_session("B"), 0)
|
||||
self.assertEqual(getattr(self._leaf([1, 2, 3, 4]), "session_ids", None), {"A"})
|
||||
self.assertEqual(self.cache.release_radix_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()
|
||||
@@ -0,0 +1,222 @@
|
||||
"""Tests for session references on UnifiedRadixCache."""
|
||||
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
|
||||
|
||||
import ast
|
||||
import unittest
|
||||
from array import array
|
||||
from pathlib import Path
|
||||
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
|
||||
from sglang.srt.mem_cache.unified_cache.components import ComponentType
|
||||
from sglang.srt.mem_cache.unified_radix_cache import UnifiedRadixCache
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
REPO_ROOT = Path(__file__).resolve().parents[4]
|
||||
MEM_CACHE_ROOT = REPO_ROOT / "python/sglang/srt/mem_cache"
|
||||
|
||||
|
||||
def class_bases(path: Path, class_name: str) -> set[str]:
|
||||
tree = ast.parse(path.read_text())
|
||||
class_node = next(
|
||||
node
|
||||
for node in tree.body
|
||||
if isinstance(node, ast.ClassDef) and node.name == class_name
|
||||
)
|
||||
return {
|
||||
base.id if isinstance(base, ast.Name) else ast.unparse(base)
|
||||
for base in class_node.bases
|
||||
}
|
||||
|
||||
|
||||
class TestSessionCacheOwnership(CustomTestCase):
|
||||
def test_only_unified_radix_cache_owns_session_ref_tracker(self):
|
||||
ordinary_mixin = MEM_CACHE_ROOT / "session_radix_cache.py"
|
||||
radix_cache = MEM_CACHE_ROOT / "radix_cache.py"
|
||||
hiradix_cache = MEM_CACHE_ROOT / "hiradix_cache.py"
|
||||
evict_policy = MEM_CACHE_ROOT / "evict_policy.py"
|
||||
unified_cache = MEM_CACHE_ROOT / "unified_radix_cache.py"
|
||||
session_ref_tracker = (
|
||||
MEM_CACHE_ROOT / "unified_cache" / "session_ref_tracker.py"
|
||||
)
|
||||
|
||||
self.assertFalse(ordinary_mixin.exists())
|
||||
ordinary_source = "\n".join(
|
||||
path.read_text() for path in (radix_cache, hiradix_cache, evict_policy)
|
||||
)
|
||||
for removed_symbol in (
|
||||
"SessionRadixCacheMixin",
|
||||
"SessionAwareEvictionStrategy",
|
||||
"session_ref",
|
||||
"_session_on_",
|
||||
"_session_forget_node",
|
||||
"_account_new_evictable_node",
|
||||
"_supports_session_radix_cache",
|
||||
"enable_session_radix_cache",
|
||||
):
|
||||
self.assertNotIn(removed_symbol, ordinary_source)
|
||||
self.assertNotIn(
|
||||
"SessionRadixCacheMixin", class_bases(radix_cache, "RadixCache")
|
||||
)
|
||||
# Session behavior is composed, not mixed in (general-code-style rule).
|
||||
self.assertEqual(
|
||||
class_bases(unified_cache, "UnifiedRadixCache"), {"BasePrefixCache"}
|
||||
)
|
||||
self.assertIn("UnifiedSessionRefTracker", session_ref_tracker.read_text())
|
||||
self.assertNotIn("SessionUnifiedRadixCacheMixin", unified_cache.read_text())
|
||||
|
||||
for component in (
|
||||
"full_component.py",
|
||||
"swa_component.py",
|
||||
"mamba_component.py",
|
||||
):
|
||||
self.assertIn(
|
||||
"session_ref",
|
||||
(
|
||||
MEM_CACHE_ROOT / "unified_cache" / "components" / component
|
||||
).read_text(),
|
||||
)
|
||||
|
||||
registry = MEM_CACHE_ROOT / "registry.py"
|
||||
self.assertIn(
|
||||
"--enable-session-radix-cache requires UnifiedRadixCache",
|
||||
registry.read_text(),
|
||||
)
|
||||
|
||||
|
||||
def make_params(enable_session: bool) -> CacheInitParams:
|
||||
dtype = torch.float16
|
||||
kv_pool = 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_pool,
|
||||
need_sort=False,
|
||||
)
|
||||
req_pool = ReqToTokenPool(
|
||||
size=8,
|
||||
max_context_len=128,
|
||||
device="cpu",
|
||||
enable_memory_saver=False,
|
||||
)
|
||||
return CacheInitParams(
|
||||
disable=False,
|
||||
req_to_token_pool=req_pool,
|
||||
token_to_kv_pool_allocator=allocator,
|
||||
page_size=1,
|
||||
eviction_policy="lru",
|
||||
enable_session_radix_cache=enable_session,
|
||||
tree_components=(ComponentType.FULL,),
|
||||
)
|
||||
|
||||
|
||||
def insert(cache, token_ids):
|
||||
"""Insert and return the tail node; the cache boundary hands back a NodeId."""
|
||||
indices = cache.token_to_kv_pool_allocator.alloc(len(token_ids))
|
||||
node_id = cache.insert(
|
||||
InsertParams(
|
||||
key=RadixKey(array("q", token_ids)),
|
||||
value=indices.to(torch.int64),
|
||||
)
|
||||
).last_device_node
|
||||
return cache.tree_core.node_by_id(node_id)
|
||||
|
||||
|
||||
def match_len(cache, token_ids) -> int:
|
||||
return len(
|
||||
cache.match_prefix(
|
||||
MatchPrefixParams(key=RadixKey(array("q", token_ids)))
|
||||
).device_indices
|
||||
)
|
||||
|
||||
|
||||
def register(cache, token_ids, session_id, generation=None):
|
||||
if generation is None:
|
||||
generation = cache.ensure_session_generation(session_id)
|
||||
cache.session_refs.register_session_ref(
|
||||
SimpleNamespace(
|
||||
session_id=session_id,
|
||||
session_generation=generation,
|
||||
session=None,
|
||||
last_node=cache.match_prefix(
|
||||
MatchPrefixParams(key=RadixKey(array("q", token_ids)))
|
||||
).last_device_node,
|
||||
origin_input_ids=array("q", token_ids),
|
||||
output_ids=array("q"),
|
||||
kv_committed_len=len(token_ids),
|
||||
extra_key=None,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
class TestRadixCacheSessionRemoval(CustomTestCase):
|
||||
def test_plain_radix_cache_does_not_enable_session_references(self):
|
||||
cache = RadixCache(make_params(enable_session=True))
|
||||
|
||||
self.assertFalse(hasattr(cache, "enable_session_radix_cache"))
|
||||
self.assertFalse(hasattr(cache, "register_session_ref"))
|
||||
self.assertFalse(hasattr(cache, "open_radix_session"))
|
||||
|
||||
|
||||
class TestSessionUnifiedRadixCache(CustomTestCase):
|
||||
def setUp(self):
|
||||
self.cache = UnifiedRadixCache(make_params(enable_session=True))
|
||||
self.full = self.cache.components[ComponentType.FULL]
|
||||
|
||||
def test_register_and_release_update_full_component_reference(self):
|
||||
leaf = insert(self.cache, [1, 2, 3, 4])
|
||||
generation = self.cache.open_radix_session("s1")
|
||||
|
||||
register(self.cache, [1, 2, 3, 4], "s1", generation)
|
||||
self.assertEqual(self.full.session_ref(leaf), 1)
|
||||
|
||||
self.cache.release_radix_session("s1")
|
||||
self.assertEqual(self.full.session_ref(leaf), 0)
|
||||
|
||||
def test_reopen_rejects_stale_generation(self):
|
||||
leaf = insert(self.cache, [1, 2, 3, 4])
|
||||
old_generation = self.cache.open_radix_session("s1")
|
||||
self.cache.release_radix_session("s1")
|
||||
self.cache.open_radix_session("s1")
|
||||
|
||||
register(self.cache, [1, 2, 3, 4], "s1", old_generation)
|
||||
|
||||
self.assertEqual(self.full.session_ref(leaf), 0)
|
||||
|
||||
def test_eviction_prefers_unreferenced_full_kv(self):
|
||||
referenced = insert(self.cache, [1, 2, 3, 4])
|
||||
insert(self.cache, [7, 8, 9])
|
||||
register(self.cache, [1, 2, 3, 4], "s1")
|
||||
|
||||
self.cache.evict(EvictParams(num_tokens=3))
|
||||
|
||||
self.assertEqual(match_len(self.cache, [7, 8, 9]), 0)
|
||||
self.assertEqual(match_len(self.cache, [1, 2, 3, 4]), 4)
|
||||
self.assertEqual(self.full.session_ref(referenced), 1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -65,6 +65,15 @@ class _StubFullComponent(TreeComponent):
|
||||
def _evict_device_end(self) -> None:
|
||||
pass
|
||||
|
||||
def _dec_session_coverage(self, session_id, leaf) -> None:
|
||||
pass
|
||||
|
||||
def _advance_session_coverage(self, session_id, leaf, old_ancestor) -> None:
|
||||
pass
|
||||
|
||||
def _recede_session_coverage(self, session_id, leaf, fallback) -> None:
|
||||
pass
|
||||
|
||||
|
||||
class _StubMambaComponent(_StubFullComponent):
|
||||
component_type = ComponentType.MAMBA
|
||||
|
||||
@@ -188,6 +188,15 @@ class _FakeFullComponent(TreeComponent):
|
||||
def _evict_device_end(self) -> None:
|
||||
pass
|
||||
|
||||
def _dec_session_coverage(self, session_id, leaf) -> None:
|
||||
pass
|
||||
|
||||
def _advance_session_coverage(self, session_id, leaf, old_ancestor) -> None:
|
||||
pass
|
||||
|
||||
def _recede_session_coverage(self, session_id, leaf, fallback) -> None:
|
||||
pass
|
||||
|
||||
|
||||
class TestUnifiedRadixComponentRegistryOverride(CustomTestCase):
|
||||
def test_component_registry_override_is_instance_local(self):
|
||||
|
||||
@@ -1274,17 +1274,6 @@ class TestPrefillOnlyDisableKvCache(unittest.TestCase):
|
||||
self._validate_prefill_only_args(kv_cache_dtype=kv_cache_dtype)
|
||||
|
||||
|
||||
class TestSessionRadixCacheServerArgs(unittest.TestCase):
|
||||
def test_requires_priority_radix_eviction_policy(self):
|
||||
server_args = ServerArgs(
|
||||
model_path="dummy",
|
||||
enable_session_radix_cache=True,
|
||||
radix_eviction_policy="lru",
|
||||
)
|
||||
with self.assertRaisesRegex(ValueError, "--radix-eviction-policy priority"):
|
||||
server_args._handle_cache_compatibility()
|
||||
|
||||
|
||||
class TestCudaGraphConfigDataclassAccess(CustomTestCase):
|
||||
@patch(
|
||||
"sglang.srt.model_executor.runner_backend."
|
||||
|
||||
Reference in New Issue
Block a user