feat: session radix cache (#27058)

Signed-off-by: Ishan Dhanani <ishandhanani@gmail.com>

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:
+        if session_id is None or radix_native_session:
-            # Normal non-session request
+            # 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:

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
 

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)
+        prefix_len, last_node = self._insert_helper(
-        return InsertResult(prefix_len=prefix_len)
+            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)

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

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.

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()

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."