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
python/sglang/srt/managers/scheduler.py
+48 -5
View File
@@ -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:
python/sglang/srt/mem_cache/base_prefix_cache.py
+1
View File
@@ -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
+18 -6
View File
@@ -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)
python/sglang/srt/mem_cache/session_radix_cache.py
+126
View File
@@ -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
+8
View File
@@ -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
+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."