diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index abba37441..64a569edf 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -123,6 +123,7 @@ from sglang.srt.managers.io_struct import ( LoadLoRAAdapterReqInput, LoadLoRAAdapterReqOutput, OpenSessionReqInput, + OpenSessionReqOutput, PauseGenerationReqInput, ProfileReq, ReleaseMemoryOccupationReqInput, @@ -583,7 +584,7 @@ class Scheduler( from sglang.srt.hardware_backend.npu.utils import init_zbal if self.ps.pp_size > 1: - logger.error(f"only zbal mix mode support pp_size > 1!") + logger.error("only zbal mix mode support pp_size > 1!") init_zbal( self.ps.tp_size, self.ps.gpu_id, self.ps.tp_rank ) # only switch allocator if is mix mode @@ -2003,9 +2004,39 @@ class Scheduler( session_id = ( recv_req.session_params.id if recv_req.session_params is not None else None ) + # Radix-native session: session_id is just a tag; KV bulk-freed on close. + radix_native_session = ( + session_id is not None and self.server_args.enable_session_radix_cache + ) + if radix_native_session: + sp = recv_req.session_params + if ( + sp.rid is not None + or sp.offset is not None + or sp.replace is not None + or sp.drop_previous_output is not None + ): + error_msg = ( + "Invalid request: radix-native sessions do not support " + "session_params rid/offset/replace/drop_previous_output; " + "send full context each turn." + ) + req = Req( + recv_req.rid, + recv_req.input_text, + recv_req.input_ids, + recv_req.sampling_params, + vocab_size=self.model_config.vocab_size, + http_worker_ipc=recv_req.http_worker_ipc, + ) + req.tokenizer = self.tokenizer + req.set_finish_with_abort(error_msg) + self.init_req_max_new_tokens(req) + self._add_request_to_queue(req) + return - if session_id is None: - # Normal non-session request + if session_id is None or radix_native_session: + # Normal non-session request, or a radix-native session request if recv_req.input_embeds is not None: # Generate fake input_ids based on the length of input_embeds seq_length = len(recv_req.input_embeds) @@ -2056,6 +2087,8 @@ class Scheduler( multi_item_delimiter_indices=recv_req.multi_item_delimiter_indices, ) req.tokenizer = self.tokenizer + if radix_native_session: + req.session_id = session_id if self.disaggregation_mode != DisaggregationMode.NULL: # Invalid request for disaggregated mode @@ -4020,13 +4053,23 @@ class Scheduler( return ExpertDistributionReqOutput() def open_session(self, recv_req: OpenSessionReqInput): - output = self.session_controller.open(recv_req) + if self.server_args.enable_session_radix_cache: + # Radix-native: open is implicit; explicit open only permits id reuse. + session_id = recv_req.session_id + self.tree_cache.register_session(session_id) + output = OpenSessionReqOutput(session_id, session_id is not None) + else: + output = self.session_controller.open(recv_req) if self.ps.pp_rank == 0 and self.ps.tp_rank == 0 and self.ps.attn_cp_rank == 0: return output return None def close_session(self, recv_req: CloseSessionReqInput): - self.session_controller.close(recv_req) + if self.server_args.enable_session_radix_cache: + # "Close" just triggers eviction of the session's tagged KV. + self.tree_cache.release_session(recv_req.session_id) + else: + self.session_controller.close(recv_req) def maybe_sleep_on_idle(self): if self.idle_sleeper is not None: diff --git a/python/sglang/srt/mem_cache/base_prefix_cache.py b/python/sglang/srt/mem_cache/base_prefix_cache.py index 4ef141822..36685f793 100644 --- a/python/sglang/srt/mem_cache/base_prefix_cache.py +++ b/python/sglang/srt/mem_cache/base_prefix_cache.py @@ -75,6 +75,7 @@ class InsertResult: prefix_len: int total_len: int = 0 + last_device_node: Any = None mamba_exist: bool = False inserted_host_node: Any = None diff --git a/python/sglang/srt/mem_cache/radix_cache.py b/python/sglang/srt/mem_cache/radix_cache.py index b4365b55c..2a10850a9 100644 --- a/python/sglang/srt/mem_cache/radix_cache.py +++ b/python/sglang/srt/mem_cache/radix_cache.py @@ -47,6 +47,7 @@ from sglang.srt.mem_cache.base_prefix_cache import ( MatchResult, ) from sglang.srt.mem_cache.events import KVCacheEventMixin +from sglang.srt.mem_cache.session_radix_cache import SessionRadixCacheMixin from sglang.srt.mem_cache.utils import get_eviction_strategy, split_node_hash_value if TYPE_CHECKING: @@ -282,7 +283,7 @@ class TreeNode: return self.last_access_time < other.last_access_time -class RadixCache(KVCacheEventMixin, BasePrefixCache): +class RadixCache(SessionRadixCacheMixin, KVCacheEventMixin, BasePrefixCache): def __init__(self, params: CacheInitParams): self.disable = params.disable self.req_to_token_pool = params.req_to_token_pool @@ -343,6 +344,7 @@ class RadixCache(KVCacheEventMixin, BasePrefixCache): self.evictable_size_ = 0 self.protected_size_ = 0 self.evictable_leaves.clear() + self._reset_session_radix_state() self._empty_match_result = MatchResult( device_indices=torch.empty( (0,), @@ -432,8 +434,10 @@ class RadixCache(KVCacheEventMixin, BasePrefixCache): # Debug/test fallback: use token ids themselves as values. value = torch.tensor(key.token_ids[: len(key)], dtype=torch.int64) - prefix_len = self._insert_helper(self.root_node, key, value, priority, chunked) - return InsertResult(prefix_len=prefix_len) + prefix_len, last_node = self._insert_helper( + self.root_node, key, value, priority, chunked + ) + return InsertResult(prefix_len=prefix_len, last_device_node=last_node) def cache_finished_req(self, req: Req, is_insert: bool = True): """Cache request when it finishes.""" @@ -466,11 +470,13 @@ class RadixCache(KVCacheEventMixin, BasePrefixCache): result = self.insert( InsertParams(key=radix_key, value=values, priority=priority) ) + session_leaf = result.last_device_node # Free the duplicates that were already in the tree self.token_to_kv_pool_allocator.free( kv_indices[req.cache_protected_len : result.prefix_len] ) else: + session_leaf = None self.token_to_kv_pool_allocator.free( kv_indices[req.cache_protected_len : key_len] ) @@ -478,6 +484,8 @@ class RadixCache(KVCacheEventMixin, BasePrefixCache): # free the unaligned tail self.token_to_kv_pool_allocator.free(kv_indices[key_len:]) + self._tag_session_leaf(req, radix_key, node=session_leaf) + # Remove req slot release the cache lock if req.last_node is not None: self.dec_lock_ref(req.last_node) @@ -548,6 +556,8 @@ class RadixCache(KVCacheEventMixin, BasePrefixCache): req.last_node = new_last_node + self._tag_session_leaf(req, radix_key, node=new_last_node) + def pretty_print(self): self._print_helper(self.root_node, 0) print(f"#tokens: {self.total_size()}") @@ -616,7 +626,7 @@ class RadixCache(KVCacheEventMixin, BasePrefixCache): if node.parent is None: assert ( node is self.root_node - ), f"This request holds the node from another tree" + ), "This request holds the node from another tree" node = node.parent return DecLockRefResult(delta=delta) @@ -712,7 +722,7 @@ class RadixCache(KVCacheEventMixin, BasePrefixCache): # Update priority along the path (take max to propagate higher priority) node.priority = max(node.priority, priority) if len(key) == 0: - return 0 + return 0, node child_key = key.child_key(self.page_size) @@ -748,7 +758,8 @@ class RadixCache(KVCacheEventMixin, BasePrefixCache): self._update_leaf_status(new_node) # Hash will be computed lazily during event emission self._record_store_event(new_node) - return total_prefix_length + node = new_node + return total_prefix_length, node def _print_helper(self, node: TreeNode, indent: int): """Prints the radix tree in a human-readable format.""" @@ -773,6 +784,7 @@ class RadixCache(KVCacheEventMixin, BasePrefixCache): v = node.parent.children.pop(key, None) assert v == node, f"parent does not have child key, {key}" + self._discard_session_leaf(node) self.evictable_size_ -= len(node.key) if node in self.evictable_leaves: self.evictable_leaves.remove(node) diff --git a/python/sglang/srt/mem_cache/session_radix_cache.py b/python/sglang/srt/mem_cache/session_radix_cache.py new file mode 100644 index 000000000..7583ba040 --- /dev/null +++ b/python/sglang/srt/mem_cache/session_radix_cache.py @@ -0,0 +1,126 @@ +"""Session radix cache (``--enable-session-radix-cache``): tag each request's KV +by session_id; ``release_session`` (close) frees a session's tagged KV.""" + +from __future__ import annotations + +import logging +from collections import OrderedDict, defaultdict +from typing import TYPE_CHECKING + +from sglang.srt.mem_cache.base_prefix_cache import MatchPrefixParams + +if TYPE_CHECKING: + from sglang.srt.managers.schedule_batch import Req + +logger = logging.getLogger(__name__) + +# Bounded guard against a request finishing after close. If a session id falls +# out of this LRU after 8192 later closes, an extremely late finish can tag +# again; explicit register_session clears the tombstone for intentional id reuse. +_CLOSED_SESSION_TOMBSTONE_LIMIT = 8192 + + +class SessionRadixCacheMixin: + """Tags radix KV by session id; ``release_session`` (close) frees a session's + tagged chains. A node holds the set of sessions on it, so a node shared by + several sessions is freed only when its last holder closes. Tagged KV is + ordinary LRU radix -- no pinning, no open. Mixed into RadixCache.""" + + def _reset_session_radix_state(self) -> None: + self._session_leaves = defaultdict(set) + self._closed_session_ids = OrderedDict() + + def _ensure_session_radix_state(self) -> None: + if not hasattr(self, "_session_leaves"): + self._reset_session_radix_state() + + def register_session(self, session_id: str) -> None: + self._ensure_session_radix_state() + if session_id is None: + return + self._closed_session_ids.pop(session_id, None) + self._session_leaves.setdefault(session_id, set()) + + def _remember_closed_session(self, session_id: str) -> None: + self._closed_session_ids[session_id] = None + self._closed_session_ids.move_to_end(session_id) + while len(self._closed_session_ids) > _CLOSED_SESSION_TOMBSTONE_LIMIT: + self._closed_session_ids.popitem(last=False) + + def _discard_session_leaf(self, node) -> None: + session_ids = getattr(node, "session_ids", None) + if not session_ids or not hasattr(self, "_session_leaves"): + return + for sid in tuple(session_ids): + leaves = self._session_leaves.get(sid) + if leaves is not None: + leaves.discard(node) + if not leaves and sid not in self._closed_session_ids: + self._session_leaves.pop(sid, None) + if hasattr(node, "session_ids"): + delattr(node, "session_ids") + + def _tag_session_leaf(self, req: Req, radix_key, node=None) -> None: + """Add this request's session id to its leaf's holder set; no-op for non-session reqs.""" + self._ensure_session_radix_state() + sid = getattr(req, "session_id", None) + if sid is None or sid in self._closed_session_ids: + return + if node is None: + logger.warning( + "_tag_session_leaf called without node; falling back to match_prefix" + ) + node = self.match_prefix(MatchPrefixParams(key=radix_key)).last_device_node + if node is not None and node is not self.root_node: + session_ids = getattr(node, "session_ids", None) + if session_ids is None: + session_ids = set() + node.session_ids = session_ids + session_ids.add(sid) + self._session_leaves[sid].add(node) + logger.debug( + "tag session %s: node=%d holders=%d indexed=%d", + sid, + node.id, + len(session_ids), + len(self._session_leaves[sid]), + ) + + def release_session(self, session_id: str) -> int: + """Close: drop this session from each of its tagged leaves, freeing a node + only once no other session still holds it (last holder). Shared + prefixes/leaves kept.""" + self._ensure_session_radix_state() + self._remember_closed_session(session_id) + indexed = self._session_leaves.pop(session_id, set()) + freed = 0 + for leaf in indexed: + if session_id not in getattr(leaf, "session_ids", set()): + continue + node = leaf + while True: + session_ids = getattr(node, "session_ids", None) + if session_ids is not None: + session_ids.discard(session_id) + if not session_ids: + delattr(node, "session_ids") + if ( + node is self.root_node + or node.lock_ref != 0 + or len(node.children) != 0 + or node not in self.evictable_leaves + or getattr(node, "session_ids", None) + ): + break + parent = node.parent + self.token_to_kv_pool_allocator.free(node.value) + self._delete_leaf(node) + freed += 1 + node = parent + logger.info( + "release_session %s: indexed %d leaves, freed %d nodes", + session_id, + len(indexed), + freed, + ) + return freed diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 74659c53a..f92d625c0 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -970,6 +970,10 @@ class ServerArgs: bool, "Enable streaming session mode and StreamingSession wrapper.", ] = False + enable_session_radix_cache: A[ + bool, + "Hold per-session KV as ordinary evictable radix entries, tagged by session id and bulk-evicted on close. Requires --radix-eviction-policy priority.", + ] = False # ------------------------------------------------------------------------- # Constrained decoding @@ -2476,6 +2480,10 @@ class ServerArgs: ) handle_pd_disaggregation(self) + if self.enable_session_radix_cache and self.radix_eviction_policy != "priority": + raise ValueError( + "--enable-session-radix-cache requires --radix-eviction-policy priority" + ) # Normalize deprecated CP aliases before validations or model-specific # defaults inspect enable_prefill_cp/cp_strategy. diff --git a/test/manual/core/test_session_radix_cache.py b/test/manual/core/test_session_radix_cache.py new file mode 100644 index 000000000..25b89abc3 --- /dev/null +++ b/test/manual/core/test_session_radix_cache.py @@ -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() diff --git a/test/registered/unit/server_args/test_server_args.py b/test/registered/unit/server_args/test_server_args.py index be2215977..f1e7643d3 100644 --- a/test/registered/unit/server_args/test_server_args.py +++ b/test/registered/unit/server_args/test_server_args.py @@ -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."