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
@@ -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):
|
||||
|
||||
Reference in New Issue
Block a user