minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

feat: Session-reference-aware Unified Radix Cache for agentic multi-turn workloads (#29173)

Browse Source 
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
2026-08-02 13:43:55 +08:00
committed by GitHub
co-authored by Ishan Dhanani hzh0425 ispobock
parent 131bd51b01
commit 056474cdb0
23 changed files with 1056 additions and 373 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/disaggregation/prefill.py
+1 -1
View File
@@ -868,9 +868,9 @@ class SchedulerDisaggregationPrefillMixin:
# todo: set Transferring correctly in backend
undone_reqs.append(req)
elif poll == KVPoll.Success: # transfer done
release_kv_cache(req, self.tree_cache) # unlock the tree
if not isinstance(req.finished_reason, FINISH_ABORT):
req.finished_reason = FINISH_LENGTH(length=0)
release_kv_cache(req, self.tree_cache) # unlock the tree
# FIXME: clean up req's data in transfer engine
req.disagg_kv_sender.clear()
done_reqs.append(req)
python/sglang/srt/managers/schedule_batch.py
+2
View File
@@ -789,6 +789,8 @@ class Req(ReqDllmMixin):
self.session = session
self.session_id = session_id
# Used by the session radix cache to reject registration after a close/reopen.
self.session_generation: Optional[int] = None
self.input_embeds = input_embeds
self.positional_embed_overrides = positional_embed_overrides
self.multi_item_delimiter_indices = multi_item_delimiter_indices
python/sglang/srt/managers/scheduler.py
+17 -4
View File
@@ -410,6 +410,7 @@ class Scheduler(
)
self.page_size = server_args.page_size
self.enable_hierarchical_cache = server_args.enable_hierarchical_cache
self.enable_session_radix_cache = server_args.enable_session_radix_cache
self.enable_hicache_storage = server_args.hicache_storage_backend is not None
self.enable_decode_hicache = (
server_args.disaggregation_decode_enable_radix_cache
@@ -2227,7 +2228,7 @@ class Scheduler(
)
# Radix-native sessions use only the top-level session_id.
radix_native_session = (
recv_req.session_id is not None and get_memory().enable_session_radix_cache
recv_req.session_id is not None and self.enable_session_radix_cache
)
if session_id is None or radix_native_session:
@@ -2286,6 +2287,11 @@ class Scheduler(
)
req.tokenizer = self.tokenizer
if radix_native_session:
req.session_generation = self.tree_cache.ensure_session_generation(
recv_req.session_id
)
if self.disaggregation_mode != DisaggregationMode.NULL:
# Invalid request for disaggregated mode
if (
@@ -2316,6 +2322,10 @@ class Scheduler(
self.model_config.vocab_size,
eos_token_ids=self.model_config.hf_eos_token_id,
)
if self.enable_session_radix_cache:
req.session_generation = self.tree_cache.ensure_session_generation(
session_id
)
# TODO: set trace context
if self.metrics_reporter.enable_metrics:
req.time_stats.set_metrics_collector(self.metrics_collector)
@@ -4605,15 +4615,18 @@ class Scheduler(
def open_session(self, recv_req: OpenSessionReqInput):
output = self.session_controller.open(recv_req)
if output.success and self.enable_session_radix_cache:
self.tree_cache.open_radix_session(recv_req.session_id)
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):
if get_memory().enable_session_radix_cache:
if self.enable_session_radix_cache:
self.tree_cache.release_radix_session(recv_req.session_id)
if recv_req.session_id in self.session_controller or not (
get_memory().enable_session_radix_cache
if (
recv_req.session_id in self.session_controller
or not self.enable_session_radix_cache
):
self.session_controller.close(recv_req)
python/sglang/srt/mem_cache/radix_cache.py
+1 -11
View File
@@ -46,7 +46,6 @@ 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,
get_hash_str,
@@ -277,14 +276,13 @@ class TreeNode:
return self.last_access_time < other.last_access_time
class RadixCache(SessionRadixCacheMixin, KVCacheEventMixin, BasePrefixCache):
class RadixCache(KVCacheEventMixin, BasePrefixCache):
def __init__(self, params: CacheInitParams):
self.disable = params.disable
self.req_to_token_pool = params.req_to_token_pool
self.token_to_kv_pool_allocator = params.token_to_kv_pool_allocator
self.page_size = params.page_size
self.enable_kv_cache_events = params.enable_kv_cache_events
self.enable_session_radix_cache = params.enable_session_radix_cache
self.is_eagle = params.is_eagle
self.disable_finished_insert = params.disable_finished_insert
self.eviction_policy = params.eviction_policy.lower()
@@ -339,7 +337,6 @@ class RadixCache(SessionRadixCacheMixin, 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,),
@@ -469,10 +466,8 @@ class RadixCache(SessionRadixCacheMixin, KVCacheEventMixin, BasePrefixCache):
result = self.insert(
InsertParams(key=radix_key, value=values, priority=priority)
)
session_leaf = result.last_device_node
freed_end = result.prefix_len
else:
session_leaf = None
freed_end = key_len
# duplicates / uninserted range, then the unaligned tail
@@ -486,8 +481,6 @@ class RadixCache(SessionRadixCacheMixin, KVCacheEventMixin, BasePrefixCache):
]
)
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)
@@ -559,8 +552,6 @@ class RadixCache(SessionRadixCacheMixin, 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()}")
@@ -788,7 +779,6 @@ class RadixCache(SessionRadixCacheMixin, 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/registry.py
+10
View File
@@ -210,6 +210,16 @@ def create_tree_cache(ctx: TreeCacheBuildContext) -> BasePrefixCache:
cache = default_radix_cache_factory(ctx)
source = "default"
if ctx.server_args.enable_session_radix_cache and not getattr(
cache, "enable_session_radix_cache", False
):
raise ValueError(
"--enable-session-radix-cache requires UnifiedRadixCache, but "
f"tree_cache is {type(cache).__name__}. Set "
"SGLANG_ENABLE_UNIFIED_RADIX_TREE=1 (or remove "
"--enable-session-radix-cache)."
)
streaming_wrapped = False
if (
ctx.server_args.enable_streaming_session
python/sglang/srt/mem_cache/session_radix_cache.py
-121
View File
@@ -1,121 +0,0 @@
"""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 _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."""
if not self.enable_session_radix_cache:
return
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_radix_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/mem_cache/unified_cache/components/full_component.py
+66 -4
View File
@@ -43,6 +43,60 @@ class FullComponent(TreeComponent):
super().__init__(cache, params)
# HiCache state: set to host KV pool when HiCache enabled
self._full_kv_pool_host = None
# Lazy bind eviction strategy since tree core is initialized after component init.
self.session_ref_eviction_strategy = (
self._session_ref_eviction_strategy
if cache.enable_session_radix_cache
else None
)
def _ensure_eviction_strategy(self) -> None:
if self.session_ref_eviction_strategy is None:
self.session_ref_eviction_strategy = (
self.tree_core.eviction_strategy.get_priority
)
def _dec_session_coverage(self, session_id: str, leaf: UnifiedTreeNode) -> None:
node = leaf
while node is not None and node is not self.tree_core.root_node:
cd = node.component_data[self.component_type]
assert cd.session_ref > 0
cd.session_ref -= 1
node = node.parent
def _advance_session_coverage(
self,
session_id: str,
leaf: UnifiedTreeNode,
old_ancestor: Optional[UnifiedTreeNode],
) -> None:
stop = old_ancestor if old_ancestor is not None else self.tree_core.root_node
node = leaf
while (
node is not None
and node is not stop
and node is not self.tree_core.root_node
):
node.component_data[self.component_type].session_ref += 1
node = node.parent
def _recede_session_coverage(
self,
session_id: str,
leaf: UnifiedTreeNode,
fallback: Optional[UnifiedTreeNode],
) -> None:
stop = fallback if fallback is not None else self.tree_core.root_node
node = leaf
while (
node is not None
and node is not stop
and node is not self.tree_core.root_node
):
cd = node.component_data[self.component_type]
assert cd.session_ref > 0
cd.session_ref -= 1
node = node.parent
def create_match_validator(
self, match_device_only: bool = False
@@ -86,7 +140,9 @@ class FullComponent(TreeComponent):
):
ct = self.component_type
new_parent.component_data[ct].lock_ref = child.component_data[ct].lock_ref
new_parent.component_data[ct].session_ref = child.component_data[ct].session_ref
child_cd = child.component_data[ct]
assert new_parent.component_data[ct].session_ids is None
split_len = len(new_parent.key)
if child_cd.value is not None:
new_parent.component_data[ct].value = child_cd.value[:split_len].clone()
@@ -127,11 +183,16 @@ class FullComponent(TreeComponent):
def eviction_priority(self, is_leaf: bool) -> int:
return 0 if is_leaf else 2
def _session_ref_eviction_strategy(self, node: UnifiedTreeNode):
ref = self.session_ref(node)
return ref > 0, ref, self.tree_core.eviction_strategy.get_priority(node)
def _evict_device_start(self, request_cnt: int) -> None:
self._ensure_eviction_strategy()
self._evict_device_request_cnt = request_cnt
self._evict_device_last_node = None
self._evict_device_heap = [
(self.tree_core.eviction_strategy.get_priority(n), n)
(self.session_ref_eviction_strategy(n), n)
for n in self.tree_core.evictable_device_leaves
]
heapq.heapify(self._evict_device_heap)
@@ -151,7 +212,7 @@ class FullComponent(TreeComponent):
):
heapq.heappush(
self._evict_device_heap,
(self.tree_core.eviction_strategy.get_priority(lv.parent), lv.parent),
(self.session_ref_eviction_strategy(lv.parent), lv.parent),
)
self._evict_device_last_node = None
while tracker[ct] < self._evict_device_request_cnt and self._evict_device_heap:
@@ -174,8 +235,9 @@ class FullComponent(TreeComponent):
host_frees: dict[ComponentType, list[torch.Tensor]],
) -> None:
"""Evict host leaves to free KV host pool space."""
self._ensure_eviction_strategy()
heap = [
(self.tree_core.eviction_strategy.get_priority(n), n)
(self.session_ref_eviction_strategy(n), n)
for n in self.tree_core.evictable_host_leaves
]
heapq.heapify(heap)
@@ -191,7 +253,7 @@ class FullComponent(TreeComponent):
):
heapq.heappush(
heap,
(self.tree_core.eviction_strategy.get_priority(x.parent), x.parent),
(self.session_ref_eviction_strategy(x.parent), x.parent),
)
def acquire_component_lock(
python/sglang/srt/mem_cache/unified_cache/components/mamba_component.py
+79 -15
View File
@@ -70,6 +70,39 @@ class MambaComponent(TreeComponent):
# HiCache state
self._mamba_pool_host = None # set to host mamba pool when HiCache enabled
def _inc_session_coverage(self, session_id: str, leaf: UnifiedTreeNode) -> None:
cd = leaf.component_data[self.component_type]
cd.session_ref += 1
if cd.session_ref == 1:
self._refresh_session_partition(leaf)
def _dec_session_coverage(self, session_id: str, leaf: UnifiedTreeNode) -> None:
cd = leaf.component_data[self.component_type]
assert cd.session_ref > 0
cd.session_ref -= 1
if cd.session_ref == 0:
self._refresh_session_partition(leaf)
def _advance_session_coverage(
self,
session_id: str,
leaf: UnifiedTreeNode,
old_ancestor: Optional[UnifiedTreeNode],
) -> None:
self._inc_session_coverage(session_id, leaf)
if old_ancestor is not None:
self._dec_session_coverage(session_id, old_ancestor)
def _recede_session_coverage(
self,
session_id: str,
leaf: UnifiedTreeNode,
fallback: Optional[UnifiedTreeNode],
) -> None:
self._dec_session_coverage(session_id, leaf)
if fallback is not None:
self._inc_session_coverage(session_id, fallback)
def refresh_lru(
self,
phase: LRURefreshPhase,
@@ -266,6 +299,8 @@ class MambaComponent(TreeComponent):
ct = self.component_type
new_parent.component_data[ct].value = None
new_parent.component_data[ct].lock_ref = 0
new_parent.component_data[ct].session_ref = 0
new_parent.component_data[ct].session_ids = None
# HiCache: mamba host_value stays on child (mamba = leaf-only data)
new_parent.component_data[ct].host_value = None
new_parent.component_data[ct].host_lock_ref = 0
@@ -311,9 +346,14 @@ class MambaComponent(TreeComponent):
def _evict_device_start(self, request_cnt: int) -> None:
"""Begin the device-eviction walk from this component's LRU cursor."""
self._evict_device_request_cnt = request_cnt
self._evict_device_cursor = self.tree_core.lru_lists[
self.component_type
].get_lru_no_lock()
if self.tree_core.enable_session_radix_cache:
lru = self.tree_core.lru_lists[self.component_type]
lru.cursor_begin()
self._evict_device_cursor = lru.cursor_next()
else:
self._evict_device_cursor = self.tree_core.lru_lists[
self.component_type
].get_lru_no_lock()
def _evict_device_next_node(
self,
@@ -322,14 +362,18 @@ class MambaComponent(TreeComponent):
host_frees: dict[ComponentType, list[torch.Tensor]],
) -> Optional[NodeId]:
"""Return the next device-leaf node for the driver to evict, or None.
Internal nodes are tombstoned inline (no IO); the cursor is re-validated
(reset to LRU head) if the previous node's eviction removed it."""
Internal nodes are tombstoned inline (no IO). If the previous node's
eviction removed the cursor, the walk resumes from the partition
sentinel with session refs on, else it restarts at the LRU tail."""
ct = self.component_type
lru = self.tree_core.lru_lists[ct]
enabled = self.tree_core.enable_session_radix_cache
if self._evict_device_cursor is not None and not lru.in_list(
self._evict_device_cursor
):
self._evict_device_cursor = lru.get_lru_no_lock()
self._evict_device_cursor = (
lru.cursor_next() if enabled else lru.get_lru_no_lock()
)
while (
tracker[ct] < self._evict_device_request_cnt
and self._evict_device_cursor is not None
@@ -337,10 +381,15 @@ class MambaComponent(TreeComponent):
):
x = self._evict_device_cursor
assert x.component_data[ct].value is not None
if x in self.tree_core.evictable_device_leaves:
self._evict_device_cursor = lru.get_prev_no_lock(x)
if x in self.tree_core.evictable_device_leaves and (
not enabled or self._can_evict_leaf_atomically(x)
):
self._evict_device_cursor = (
lru.cursor_next() if enabled else lru.get_prev_no_lock(x)
)
return x.id
x_next = lru.get_prev_no_lock(x)
if not enabled:
x_next = lru.get_prev_no_lock(x)
self.tree_core._evict_component_and_detach_lru(
x,
self,
@@ -352,11 +401,13 @@ class MambaComponent(TreeComponent):
self.tree_core._cascade_evict(
x, self, tracker, device_frees=device_frees, host_frees=host_frees
)
self._evict_device_cursor = x_next
self._evict_device_cursor = lru.cursor_next() if enabled else x_next
return None
def _evict_device_end(self) -> None:
"""Clear the device-eviction walk cursor state."""
if self.tree_core.enable_session_radix_cache:
self.tree_core.lru_lists[self.component_type].cursor_end()
self._evict_device_cursor = None
def acquire_component_lock(
@@ -787,15 +838,23 @@ class MambaComponent(TreeComponent):
Host leaves: atomic eviction via _evict_host_leaf."""
ct = self.component_type
host_lru = self.tree_core.host_lru_lists[ct]
x = host_lru.get_lru_no_host_lock()
enabled = self.tree_core.enable_session_radix_cache
if enabled:
host_lru.cursor_begin()
x = host_lru.cursor_next(host_lock=True)
else:
x = host_lru.get_lru_no_host_lock()
while tracker[ct] < num_tokens and x is not None and host_lru.in_list(x):
x_next = host_lru.get_prev_no_host_lock(x)
if not enabled:
x_next = host_lru.get_prev_no_host_lock(x)
cd = x.component_data[ct]
if x in self.tree_core.evictable_host_leaves:
if x in self.tree_core.evictable_host_leaves and (
not enabled or self._can_evict_leaf_atomically(x)
):
# Host leaf: atomic eviction (all components host + delete)
self.tree_core._evict_host_leaf(x, tracker, device_frees, host_frees)
else:
# Internal: tombstone Mamba + cascade
# Internal (or a leaf a session still pins): tombstone Mamba + cascade
assert cd.host_value is not None
self.tree_core._evict_component_and_detach_lru(
x,
@@ -814,7 +873,12 @@ class MambaComponent(TreeComponent):
target=EvictLayer.HOST,
)
self.tree_core._update_evictable_leaf_sets(x)
x = x_next
if enabled:
x = host_lru.cursor_next(host_lock=True)
else:
x = x_next
if enabled:
host_lru.cursor_end()
def free_host_values(self, host_values: list[torch.Tensor]) -> None:
if self._mamba_pool_host is None:
python/sglang/srt/mem_cache/unified_cache/components/swa_component.py
+136 -14
View File
@@ -69,12 +69,101 @@ class SWAComponent(TreeComponent):
params.token_to_kv_pool_allocator, SWATokenToKVPoolAllocator
), f"SWAComponent requires SWATokenToKVPoolAllocator, got {type(params.token_to_kv_pool_allocator)}"
super().__init__(cache, params)
self._session_leaf_covered_len: dict[str, dict[UnifiedTreeNode, int]] = {}
self.sliding_window_size = params.sliding_window_size
# HiCache state: set to host SWA pool when HiCache enabled
self._swa_kv_pool_host = None
component_type = ComponentType.SWA
def reset_session_state(self) -> None:
super().reset_session_state()
self._session_leaf_covered_len = {}
def _walk_session_coverage(
self,
leaf: UnifiedTreeNode,
span: int,
delta: int,
) -> int:
node = leaf
covered = 0
while node is not self.tree_core.root_node and covered < span:
cd = node.component_data[self.component_type]
if delta < 0:
assert cd.session_ref > 0
prev_ref = cd.session_ref
cd.session_ref += delta
if (prev_ref == 0) != (cd.session_ref == 0):
self._refresh_session_partition(node)
covered += len(node.key)
node = node.parent
return covered
def _inc_session_coverage(self, session_id: str, leaf: UnifiedTreeNode) -> None:
covered_by_leaf = self._session_leaf_covered_len.setdefault(session_id, {})
assert leaf not in covered_by_leaf
target_span = self.sliding_window_size + self.tree_core.page_size
covered = self._walk_session_coverage(leaf, target_span, 1)
assert covered > 0
covered_by_leaf[leaf] = covered
def _dec_session_coverage(self, session_id: str, leaf: UnifiedTreeNode) -> None:
covered_by_leaf = self._session_leaf_covered_len.get(session_id)
assert covered_by_leaf is not None and leaf in covered_by_leaf
covered_len = covered_by_leaf.pop(leaf)
if not covered_by_leaf:
self._session_leaf_covered_len.pop(session_id, None)
actual = self._walk_session_coverage(leaf, covered_len, -1)
assert actual == covered_len
def _advance_session_coverage(
self,
session_id: str,
leaf: UnifiedTreeNode,
old_ancestor: Optional[UnifiedTreeNode],
) -> None:
self._inc_session_coverage(session_id, leaf)
if old_ancestor is not None:
self._dec_session_coverage(session_id, old_ancestor)
def _recede_session_coverage(
self,
session_id: str,
leaf: UnifiedTreeNode,
fallback: Optional[UnifiedTreeNode],
) -> None:
self._dec_session_coverage(session_id, leaf)
if fallback is not None:
self._inc_session_coverage(session_id, fallback)
def validate_session_state(
self,
reachable_nodes: set[UnifiedTreeNode],
report_error: Callable[[str], None],
) -> None:
super().validate_session_state(reachable_nodes, report_error)
ct = self.component_type
for session_id, covered_by_leaf in self._session_leaf_covered_len.items():
for leaf, covered_len in covered_by_leaf.items():
if leaf not in self._session_leaves.get(session_id, ()):
report_error(
f"{ct} session {session_id!r} coverage leaf {leaf.id} is not indexed"
)
if covered_len <= 0:
report_error(
f"{ct} session {session_id!r} leaf {leaf.id} covered_len={covered_len}"
)
for session_id, leaves in self._session_leaves.items():
covered_by_leaf = self._session_leaf_covered_len.get(session_id, {})
for leaf in leaves:
if leaf not in covered_by_leaf:
report_error(
f"{ct} session {session_id!r} leaf {leaf.id} has no coverage record"
)
def _translate_full_to_swa(self, full_indices: torch.Tensor) -> torch.Tensor:
return self.cache.token_to_kv_pool_allocator.translate_loc_from_full_to_swa(
full_indices
@@ -335,6 +424,10 @@ class SWAComponent(TreeComponent):
new_parent.component_data[self.component_type].lock_ref = child.component_data[
self.component_type
].lock_ref
new_parent.component_data[self.component_type].session_ref = (
child.component_data[self.component_type].session_ref
)
assert new_parent.component_data[self.component_type].session_ids is None
child_swa_value = child.component_data[self.component_type].value
if child_swa_value is not None:
@@ -421,9 +514,14 @@ class SWAComponent(TreeComponent):
def _evict_device_start(self, request_cnt: int) -> None:
"""Begin the device-eviction walk from this component's LRU cursor."""
self._evict_device_request_cnt = request_cnt
self._evict_device_cursor = self.tree_core.lru_lists[
self.component_type
].get_lru_no_lock()
if self.tree_core.enable_session_radix_cache:
lru = self.tree_core.lru_lists[self.component_type]
lru.cursor_begin()
self._evict_device_cursor = lru.cursor_next()
else:
self._evict_device_cursor = self.tree_core.lru_lists[
self.component_type
].get_lru_no_lock()
def _evict_device_next_node(
self,
@@ -432,14 +530,18 @@ class SWAComponent(TreeComponent):
host_frees: dict[ComponentType, list[torch.Tensor]],
) -> Optional[NodeId]:
"""Return the next device-leaf node for the driver to evict, or None.
Internal nodes are tombstoned inline (no IO); the cursor is re-validated
(reset to LRU head) if the previous node's eviction removed it."""
Internal nodes are tombstoned inline (no IO). If the previous node's
eviction removed the cursor, the walk resumes from the partition
sentinel with session refs on, else it restarts at the LRU tail."""
ct = self.component_type
lru = self.tree_core.lru_lists[ct]
enabled = self.tree_core.enable_session_radix_cache
if self._evict_device_cursor is not None and not lru.in_list(
self._evict_device_cursor
):
self._evict_device_cursor = lru.get_lru_no_lock()
self._evict_device_cursor = (
lru.cursor_next() if enabled else lru.get_lru_no_lock()
)
while (
tracker[ct] < self._evict_device_request_cnt
and self._evict_device_cursor is not None
@@ -447,10 +549,15 @@ class SWAComponent(TreeComponent):
):
x = self._evict_device_cursor
assert x.component_data[ct].value is not None
if x in self.tree_core.evictable_device_leaves:
self._evict_device_cursor = lru.get_prev_no_lock(x)
if x in self.tree_core.evictable_device_leaves and (
not enabled or self._can_evict_leaf_atomically(x)
):
self._evict_device_cursor = (
lru.cursor_next() if enabled else lru.get_prev_no_lock(x)
)
return x.id
x_next = lru.get_prev_no_lock(x)
if not enabled:
x_next = lru.get_prev_no_lock(x)
self.tree_core._evict_component_and_detach_lru(
x,
self,
@@ -462,11 +569,13 @@ class SWAComponent(TreeComponent):
self.tree_core._cascade_evict(
x, self, tracker, device_frees=device_frees, host_frees=host_frees
)
self._evict_device_cursor = x_next
self._evict_device_cursor = lru.cursor_next() if enabled else x_next
return None
def _evict_device_end(self) -> None:
"""Clear the device-eviction walk cursor state."""
if self.tree_core.enable_session_radix_cache:
self.tree_core.lru_lists[self.component_type].cursor_end()
self._evict_device_cursor = None
def acquire_component_lock(
@@ -938,11 +1047,19 @@ class SWAComponent(TreeComponent):
Host leaves: atomic eviction via _evict_host_leaf."""
ct = self.component_type
host_lru = self.tree_core.host_lru_lists[ct]
x = host_lru.get_lru_no_host_lock()
enabled = self.tree_core.enable_session_radix_cache
if enabled:
host_lru.cursor_begin()
x = host_lru.cursor_next(host_lock=True)
else:
x = host_lru.get_lru_no_host_lock()
while tracker[ct] < num_tokens and x is not None and host_lru.in_list(x):
x_next = host_lru.get_prev_no_host_lock(x)
if not enabled:
x_next = host_lru.get_prev_no_host_lock(x)
cd = x.component_data[ct]
if x in self.tree_core.evictable_host_leaves:
if x in self.tree_core.evictable_host_leaves and (
not enabled or self._can_evict_leaf_atomically(x)
):
self.tree_core._evict_host_leaf(x, tracker, device_frees, host_frees)
else:
assert cd.host_value is not None
@@ -962,7 +1079,12 @@ class SWAComponent(TreeComponent):
host_frees=host_frees,
target=EvictLayer.HOST,
)
x = x_next
if enabled:
x = host_lru.cursor_next(host_lock=True)
else:
x = x_next
if enabled:
host_lru.cursor_end()
def free_host_values(self, host_values: list[torch.Tensor]) -> None:
if self._swa_kv_pool_host is None:
python/sglang/srt/mem_cache/unified_cache/components/tree_component.py
+182
View File
@@ -2,6 +2,7 @@ from __future__ import annotations
import dataclasses
from abc import ABC, abstractmethod
from collections import defaultdict
from enum import Enum, IntFlag
from typing import TYPE_CHECKING, Any, Callable, Optional, Sequence
@@ -50,6 +51,8 @@ class ComponentData:
metadata: dict[str, Any] = dataclasses.field(default_factory=dict)
host_value: Optional[torch.Tensor] = None
host_lock_ref: int = 0
session_ref: int = 0
session_ids: Optional[set[str]] = None
class EvictLayer(IntFlag):
@@ -112,10 +115,189 @@ class TreeComponent(ABC):
# Populated when the component passed to TreeCore constructor.
self.tree_core: Optional[UnifiedTreeCore] = None
self.is_evict_device_ongoing = False
# Per-session frontier nodes (the deepest registered node per cached
# path), not physical tree leaves: a frontier node may have children.
self._session_leaves: dict[str, set[UnifiedTreeNode]] = defaultdict(set)
# Subclasses MUST set this as a class attribute (not @property)
component_type: ComponentType
def reset_session_state(self) -> None:
self._session_leaves = defaultdict(set)
def session_ref(self, node: UnifiedTreeNode) -> int:
return node.component_data[self.component_type].session_ref
def _can_evict_leaf_atomically(self, node: UnifiedTreeNode) -> bool:
# Don't allow a non-session-ref component to cascade-evict
# a high-priority session-ref component.
if self.session_ref(node) > 0:
return True
priority = self.eviction_priority(is_leaf=False)
return not any(
comp.eviction_priority(is_leaf=False) >= priority
and comp.session_ref(node) > 0
for comp in self.tree_core.components
)
def _refresh_session_partition(self, node: UnifiedTreeNode) -> None:
ct = self.component_type
lru = self.tree_core.lru_lists[ct]
if lru.in_list(node):
lru.reset_node_mru(node)
host_lru = self.tree_core.host_lru_lists[ct]
if host_lru.in_list(node):
host_lru.reset_node_mru(node)
def _find_reusable_session_leaf(
self, node: Optional[UnifiedTreeNode]
) -> Optional[UnifiedTreeNode]:
root_node = self.tree_core.root_node
if node is None or node is root_node:
return None
validator = self.create_match_validator(match_device_only=False)
cur = node
while cur is not None and cur is not root_node:
if validator(cur):
return cur
cur = cur.parent
return None
def resolve_session_leaf(
self, req: Req, last_node: Optional[UnifiedTreeNode]
) -> Optional[UnifiedTreeNode]:
return self._find_reusable_session_leaf(last_node)
def _mark_session_leaf(self, session_id: str, leaf: UnifiedTreeNode) -> None:
self._session_leaves[session_id].add(leaf)
cd = leaf.component_data[self.component_type]
if cd.session_ids is None:
cd.session_ids = set()
cd.session_ids.add(session_id)
def _unmark_session_leaf(self, session_id: str, leaf: UnifiedTreeNode) -> None:
leaves = self._session_leaves.get(session_id)
assert leaves is not None and leaf in leaves
leaves.remove(leaf)
cd = leaf.component_data[self.component_type]
assert cd.session_ids is not None and session_id in cd.session_ids
cd.session_ids.remove(session_id)
if not cd.session_ids:
cd.session_ids = None
if not leaves:
self._session_leaves.pop(session_id, None)
def _nearest_session_ancestor(
self, new_leaf: UnifiedTreeNode, current_leaves: set[UnifiedTreeNode]
) -> Optional[UnifiedTreeNode]:
cur = new_leaf.parent
while cur is not None and cur is not self.tree_core.root_node:
if cur in current_leaves:
return cur
cur = cur.parent
return None
@abstractmethod
def _dec_session_coverage(self, session_id: str, leaf: UnifiedTreeNode) -> None:
"""Remove one session reference from the coverage anchored at `leaf`."""
...
@abstractmethod
def _advance_session_coverage(
self,
session_id: str,
leaf: UnifiedTreeNode,
old_ancestor: Optional[UnifiedTreeNode],
) -> None:
"""Move the session's coverage from `old_ancestor` forward to `leaf`."""
...
@abstractmethod
def _recede_session_coverage(
self,
session_id: str,
leaf: UnifiedTreeNode,
fallback: Optional[UnifiedTreeNode],
) -> None:
"""Move the session's coverage backward from `leaf` to `fallback`."""
...
def register_session_leaf(
self, session_id: str, leaf: Optional[UnifiedTreeNode]
) -> None:
if (
not self.tree_core.enable_session_radix_cache
or leaf is None
or leaf is self.tree_core.root_node
):
return
current_leaves = self._session_leaves[session_id]
if leaf in current_leaves:
return
old_ancestor = self._nearest_session_ancestor(leaf, current_leaves)
self._advance_session_coverage(session_id, leaf, old_ancestor)
self._mark_session_leaf(session_id, leaf)
if old_ancestor is not None:
self._unmark_session_leaf(session_id, old_ancestor)
def release_session(self, session_id: str) -> int:
leaves = tuple(self._session_leaves.get(session_id, ()))
for leaf in leaves:
self._dec_session_coverage(session_id, leaf)
self._unmark_session_leaf(session_id, leaf)
return len(leaves)
def discard_deleted_session_leaf(self, node: UnifiedTreeNode) -> None:
cd = node.component_data[self.component_type]
for session_id in tuple(cd.session_ids or ()):
leaves = self._session_leaves.get(session_id)
assert leaves is not None and node in leaves
fallback = self._find_reusable_session_leaf(node.parent)
if fallback is not None and fallback in self._session_leaves.get(
session_id, ()
):
fallback = None
self._recede_session_coverage(session_id, node, fallback)
self._unmark_session_leaf(session_id, node)
if fallback is not None:
self._mark_session_leaf(session_id, fallback)
def validate_session_state(
self,
reachable_nodes: set[UnifiedTreeNode],
report_error: Callable[[str], None],
) -> None:
reachable_nodes = set(reachable_nodes)
ct = self.component_type
for session_id, leaves in self._session_leaves.items():
if not leaves:
report_error(f"{ct} session {session_id!r} has an empty leaf index")
for leaf in leaves:
if leaf not in reachable_nodes:
report_error(
f"{ct} session {session_id!r} indexes unreachable leaf {leaf.id}"
)
if session_id not in (leaf.component_data[ct].session_ids or ()):
report_error(
f"{ct} session {session_id!r} leaf {leaf.id} is missing its marker"
)
for node in reachable_nodes:
cd = node.component_data[ct]
if cd.session_ref < 0:
report_error(f"node {node.id} {ct} session_ref={cd.session_ref}")
if cd.session_ids is not None and not cd.session_ids:
report_error(f"node {node.id} {ct} has an empty session_ids marker")
for session_id in cd.session_ids or ():
if node not in self._session_leaves.get(session_id, ()):
report_error(
f"node {node.id} {ct} session {session_id!r} marker is not indexed"
)
def node_has_component_data(
self, node: UnifiedTreeNode, target: EvictLayer = EvictLayer.DEVICE
) -> bool:
python/sglang/srt/mem_cache/unified_cache/session_ref_tracker.py
+118
View File
@@ -0,0 +1,118 @@
"""Session ref tracking for UnifiedRadixCache (``--enable-session-radix-cache``):
tag each request's KV by session_id for each tree component; ``release_radix_session``
(close) releases a session's tagged reference.
"""
from __future__ import annotations
import logging
from collections import OrderedDict
from dataclasses import dataclass
from typing import TYPE_CHECKING, Optional
if TYPE_CHECKING:
from sglang.srt.managers.schedule_batch import Req
from sglang.srt.mem_cache.unified_cache.components.tree_component import (
TreeComponent,
)
from sglang.srt.mem_cache.unified_cache.unified_tree_core import UnifiedTreeCore
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
@dataclass(kw_only=True)
class UnifiedSessionRefTracker:
"""Tags radix KV by session id; ``release_radix_session`` (close) releases a
session's tagged reference. Each component maintains its own session_ids,
session_ref and so on."""
components: tuple[TreeComponent, ...]
tree_core: UnifiedTreeCore
enable_session_radix_cache: bool
def __post_init__(self) -> None:
self.reset()
def reset(self) -> None:
self._closed_session_ids: OrderedDict[str, None] = OrderedDict()
self._session_incarnation_counter: int = 0
self._session_generations: dict[str, int] = {}
for component in self.components:
component.reset_session_state()
def session_id_for_req(self, req: Req) -> Optional[str]:
session_id = req.session_id
if session_id is None and req.session is not None:
session_id = req.session.session_id
return session_id
def register_session_ref(self, req: Req) -> None:
"""Register a non-streaming request's reusable leaves with each component."""
if not self.enable_session_radix_cache:
return
session = req.session
if session is not None and session.streaming:
return
session_id = self.session_id_for_req(req)
if session_id is None or session_id in self._closed_session_ids:
return
current_generation = self._session_generations.get(session_id)
if current_generation is None or req.session_generation != current_generation:
logger.warning("register_session_ref called for stale request; Skip it.")
return
assert req.last_node is not None
last_node = self.tree_core.node_by_id(req.last_node)
if last_node is self.tree_core.root_node:
return
for component in self.components:
leaf = component.resolve_session_leaf(req, last_node)
component.register_session_leaf(session_id, leaf)
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 open_radix_session(self, session_id: str) -> Optional[int]:
self._closed_session_ids.pop(session_id, None)
self._session_incarnation_counter += 1
self._session_generations[session_id] = self._session_incarnation_counter
return self._session_incarnation_counter
def ensure_session_generation(self, session_id: str) -> int:
generation = self._session_generations.get(session_id)
if generation is None:
generation = self.open_radix_session(session_id)
return generation
def release_radix_session(self, session_id: str) -> int:
if not self.enable_session_radix_cache or session_id is None:
return 0
if session_id in self._closed_session_ids:
return 0
self._remember_closed_session(session_id)
self._session_generations.pop(session_id, None)
indexed = 0
for component in self.components:
indexed += component.release_session(session_id)
logger.info(
"release_session %s: indexed %d component leaves",
session_id,
indexed,
)
return 0
python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py
+108 -12
View File
@@ -21,7 +21,7 @@ import sys
from array import array
from collections import defaultdict
from enum import Enum, auto
from typing import TYPE_CHECKING, Any, NamedTuple, Optional, Sequence
from typing import TYPE_CHECKING, Any, Callable, NamedTuple, Optional, Sequence
import msgspec
import torch
@@ -161,6 +161,7 @@ class UnifiedLRUList:
component_type: ComponentType,
tree_components: tuple[ComponentType, ...],
use_host_ptr: bool = False,
is_referenced: Optional[Callable[[UnifiedTreeNode], bool]] = None,
):
self.component_type = component_type
# Pointer slot: host LRU uses offset slots so device/host pointers
@@ -171,6 +172,19 @@ class UnifiedLRUList:
self.head.lru_next[self._pt] = self.tail
self.tail.lru_prev[self._pt] = self.head
self.cache: dict[int, UnifiedTreeNode] = {}
# Session partition: [head .. mid) holds session-referenced nodes and
# (mid .. tail] unreferenced ones, so evictions directly walks (tail -> head).
self._is_referenced = is_referenced
self.mid: Optional[UnifiedTreeNode] = None
self.cursor: Optional[UnifiedTreeNode] = None
if is_referenced is not None:
self.mid = UnifiedTreeNode(tree_components)
self.cursor = UnifiedTreeNode(tree_components)
for ct in tree_components:
for node in (self.mid, self.cursor):
node.component_data[ct].lock_ref = 1
node.component_data[ct].host_lock_ref = 1
self._add_node_after(self.head, self.mid)
def _add_node_after(self, prev_node: UnifiedTreeNode, new_node: UnifiedTreeNode):
pt = self._pt
@@ -180,7 +194,10 @@ class UnifiedLRUList:
prev_node.lru_next[pt] = new_node
def _add_node(self, node: UnifiedTreeNode):
self._add_node_after(self.head, node)
if self._is_referenced is None or self._is_referenced(node):
self._add_node_after(self.head, node)
else:
self._add_node_after(self.mid, node)
def _remove_node(self, node: UnifiedTreeNode):
pt = self._pt
@@ -205,19 +222,50 @@ class UnifiedLRUList:
self._remove_node(node)
self._add_node(node)
def cursor_begin(self):
if self.cursor.lru_prev[self._pt] is not None:
self._remove_node(self.cursor)
self._add_node_after(self.tail.lru_prev[self._pt], self.cursor)
def cursor_next(self, *, host_lock: bool = False):
pt = self._pt
ct = self.component_type
x = self.cursor.lru_prev[pt]
while x is not self.head:
cd = x.component_data[ct]
if (cd.host_lock_ref if host_lock else cd.lock_ref) == 0:
break
x = x.lru_prev[pt]
if x is self.head:
return None
self._remove_node(self.cursor)
self._add_node_after(x.lru_prev[pt], self.cursor)
return x
def cursor_end(self):
"""Unlink the walk-cursor sentinel after a walk."""
self._remove_node(self.cursor)
def reset_node_and_parents_mru(
self,
node: UnifiedTreeNode,
root_node: UnifiedTreeNode,
should_include,
):
prev_node = self.head
is_referenced = self._is_referenced
prev_ref = self.head
prev_unref = self.head if self.mid is None else self.mid
while node != root_node:
if should_include(node):
assert node.id in self.cache
part = is_referenced is None or is_referenced(node)
self._remove_node(node)
self._add_node_after(prev_node, node)
prev_node = node
if part:
self._add_node_after(prev_ref, node)
prev_ref = node
else:
self._add_node_after(prev_unref, node)
prev_unref = node
node = node.parent
def reset_node_and_window_ancestors_mru(
@@ -227,14 +275,21 @@ class UnifiedLRUList:
window_size: int,
should_include,
):
prev_node = self.head
is_referenced = self._is_referenced
prev_ref = self.head
prev_unref = self.head if self.mid is None else self.mid
accumulated = 0
while node != root_node and accumulated < window_size:
if should_include(node):
assert node.id in self.cache
part = is_referenced is None or is_referenced(node)
self._remove_node(node)
self._add_node_after(prev_node, node)
prev_node = node
if part:
self._add_node_after(prev_ref, node)
prev_ref = node
else:
self._add_node_after(prev_unref, node)
prev_unref = node
accumulated += len(node.key)
node = node.parent
@@ -334,6 +389,7 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
self.write_through_threshold = 256
self.is_write_back = False
self.has_swa_host_pool = False
self.enable_session_radix_cache = params.enable_session_radix_cache
self.eviction_strategy = get_eviction_strategy(params.eviction_policy.lower())
# ``device`` is derived from the construction-time allocator; the
@@ -361,6 +417,11 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
# ==== Tree API ====
def _session_lru_predicate(self, ct: ComponentType):
if not self.enable_session_radix_cache or ct is ComponentType.FULL:
return None
return lambda node: node.component_data[ct].session_ref > 0
def reset(self) -> None:
"""Rebuild the root, LRUs, sizes, evictable-leaf sets, and the empty
match result."""
@@ -382,13 +443,21 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
self.component_protected_size_ = {ct: 0 for ct in self.component_types}
self.lru_lists = {
ct: UnifiedLRUList(ct, self.component_types) for ct in self.component_types
ct: UnifiedLRUList(
ct, self.component_types, is_referenced=self._session_lru_predicate(ct)
)
for ct in self.component_types
}
self.evictable_device_leaves: set[UnifiedTreeNode] = set()
self.evictable_host_leaves: set[UnifiedTreeNode] = set()
self.host_lru_lists = {
ct: UnifiedLRUList(ct, self.component_types, use_host_ptr=True)
ct: UnifiedLRUList(
ct,
self.component_types,
use_host_ptr=True,
is_referenced=self._session_lru_predicate(ct),
)
for ct in self.component_types
}
@@ -767,7 +836,12 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
self.root_node.priority = max(self.root_node.priority, priority)
if len(key) == 0:
return InsertStepResult(
actions=[], result=InsertResult(prefix_len=0, mamba_exist=True)
actions=[],
result=InsertResult(
prefix_len=0,
mamba_exist=True,
last_device_node=self.root_node.id,
),
)
self._ongoing_insert_walk_state = _InsertWalkState(
@@ -902,7 +976,10 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
# e.g. Mamba attaches mamba_value to the leaf node
# All hooks run before their emitted actions execute; an action failure
# fail-stops the process, so partial-commit state is never observed.
state.result = InsertResult(prefix_len=state.total_prefix_length)
state.result = InsertResult(
prefix_len=state.total_prefix_length,
last_device_node=state.target_node.id,
)
for component in self.components:
component.commit_insert_component_data(
node=state.target_node,
@@ -1287,6 +1364,7 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
is_leaf = node in self.evictable_host_leaves
trigger_priority = trigger.eviction_priority(is_leaf)
base_evicted = False
for comp in self.components:
if comp.eviction_priority(is_leaf) <= trigger_priority:
@@ -1304,6 +1382,8 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
continue
if EvictLayer.HOST in target and cd.host_lock_ref != 0:
continue
if cd.session_ref > 0 and trigger.session_ref(node) == 0:
continue
if EvictLayer.DEVICE in target:
assert cd.lock_ref == 0
if EvictLayer.HOST in target:
@@ -1316,6 +1396,8 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
device_frees=device_frees,
host_frees=host_frees,
)
if comp.component_type == BASE_COMPONENT_TYPE:
base_evicted = True
# Now that all components (including SWA which depends on Full.value)
# have been freed, we can safely tombstone Full.value.
@@ -1326,9 +1408,15 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
):
node.component_data[trigger.component_type].value = None
if EvictLayer.DEVICE in target and base_evicted:
node.component_data[BASE_COMPONENT_TYPE].value = None
self._update_evictable_leaf_sets(node)
def _remove_leaf_from_parent(self, node: UnifiedTreeNode):
for component in self.components:
component.discard_deleted_session_leaf(node)
key = node.key.child_key(self.page_size)
v = node.parent.children.pop(key, None)
assert v == node
@@ -1876,6 +1964,10 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
f"[Leaf] {len(overlap)} in both sets: {[n.id for n in list(overlap)[:5]]}"
)
if self.enable_session_radix_cache:
for component in self.components:
component.validate_session_state(all_node_set, E)
# Stale nodes: leaf sets must only contain tree-reachable nodes
stale = self.evictable_device_leaves - all_node_set
if stale:
@@ -2011,6 +2103,10 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
while x is not None and x != lru.tail:
if x.lru_prev[pt] != prev:
errors.append(f"[{label}][{ct}] broken prev at node {x.id}")
if x is lru.mid or x is lru.cursor:
prev = x
x = x.lru_next[pt]
continue
if x.id not in lru.cache:
errors.append(f"[{label}][{ct}] node {x.id} in list not cache")
if x.id in visited:
python/sglang/srt/mem_cache/unified_radix_cache.py
+35
View File
@@ -54,6 +54,9 @@ from sglang.srt.mem_cache.unified_cache.components import (
SWAComponent,
TreeComponent,
)
from sglang.srt.mem_cache.unified_cache.session_ref_tracker import (
UnifiedSessionRefTracker,
)
from sglang.srt.mem_cache.unified_cache.tree_core_registry import create_tree_core
from sglang.srt.mem_cache.unified_cache.unified_tree_core import ( # noqa: F401
NodeId,
@@ -135,6 +138,7 @@ class UnifiedRadixCache(BasePrefixCache):
assert params.tree_components is not None
self.tree_components = tuple(params.tree_components)
self.enable_session_radix_cache = params.enable_session_radix_cache
component_registry = COMPONENT_REGISTRY
if params.component_registry_override:
component_registry = {
@@ -170,6 +174,13 @@ class UnifiedRadixCache(BasePrefixCache):
for component in self.components.values():
component.tree_core = self.tree_core
# Session ref tracking (--enable-session-radix-cache).
self.session_refs = UnifiedSessionRefTracker(
components=self._components_tuple,
tree_core=self.tree_core,
enable_session_radix_cache=self.enable_session_radix_cache,
)
self.sidecar_pool_specs: list[SidecarPoolSpec] = []
# Streaming session: embedded StreamingSession with self as inner.
@@ -276,6 +287,7 @@ class UnifiedRadixCache(BasePrefixCache):
def _reset_full(self) -> None:
"""Full reset: destroy entire tree and all state."""
self.tree_core.reset()
self.session_refs.reset()
# Reset Controller.
self.session.slots.clear()
@@ -668,12 +680,23 @@ class UnifiedRadixCache(BasePrefixCache):
self._dec_req_lock(req, skip_swa=req.swa_prefix_lock_released)
if is_insert and result is not None and result.last_device_node is not None:
req.last_node = result.last_device_node
# cleanup
for comp in self._components_tuple:
comp.cleanup_after_caching_req(
req, is_finished=True, insert_result=result, insert_params=insert_params
)
if self.enable_session_radix_cache and result is not None:
from sglang.srt.managers.schedule_batch import FINISH_ABORT
if req.finished_reason is not None and not isinstance(
req.finished_reason, FINISH_ABORT
):
self.session_refs.register_session_ref(req)
def cache_unfinished_req(self, req: Req, chunked: bool = False, **kwargs) -> None:
if self.session.try_cache_unfinished_req(req, chunked=chunked, **kwargs):
return
@@ -1998,6 +2021,17 @@ class UnifiedRadixCache(BasePrefixCache):
def supports_mamba(self) -> bool:
return self.is_mamba_enabled
# ---- Session radix cache API (delegates to composed UnifiedSessionRefTracker) ----
def open_radix_session(self, session_id: str) -> Optional[int]:
return self.session_refs.open_radix_session(session_id)
def ensure_session_generation(self, session_id: str) -> int:
return self.session_refs.ensure_session_generation(session_id)
def release_radix_session(self, session_id: str) -> int:
return self.session_refs.release_radix_session(session_id)
# ---- Streaming session API (delegates to composed StreamingSession) ----
def supports_streaming_session(self) -> bool:
@@ -2055,6 +2089,7 @@ class UnifiedRadixCache(BasePrefixCache):
return self.tree_core.all_mamba_values_flatten()
def available_and_evictable_str(self) -> str:
# TODO(zhangmj): need more detailed log info for session reference.
if self.supports_swa():
full_available_size = self.token_to_kv_pool_allocator.full_available_size()
else:
python/sglang/srt/server_args.py
+1 -6
View File
@@ -1409,7 +1409,7 @@ class ServerArgs:
] = 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.",
"Track per-session references on UnifiedRadixCache KV: eviction consumes unreferenced entries before referenced ones, and closing a session only dereferences its KV.",
NS("memory"),
] = False
@@ -7568,11 +7568,6 @@ class ServerArgs:
envs.SGLANG_OPT_FP8_WO_A_GEMM.set(False)
def _handle_cache_compatibility(self):
if self.enable_session_radix_cache and self.radix_eviction_policy != "priority":
raise ValueError(
"--enable-session-radix-cache requires --radix-eviction-policy priority"
)
if self.enable_hierarchical_cache and self.disable_radix_cache:
raise ValueError(
"The arguments enable-hierarchical-cache and disable-radix-cache are mutually exclusive "
python/sglang/srt/session/session_controller.py
+1
View File
@@ -431,6 +431,7 @@ class SessionController:
mm.release_features()
node.req.multimodal_inputs = None
self.tree_cache.release_radix_session(session_id)
self.tree_cache.release_session(session_id)
del self.sessions[session_id]
log_info_on_rank0(