diff --git a/docs_new/docs.json b/docs_new/docs.json index f2d66d391..b86df1f66 100644 --- a/docs_new/docs.json +++ b/docs_new/docs.json @@ -920,6 +920,7 @@ "pages": [ "docs/advanced_features/overview", "docs/advanced_features/server_arguments", + "docs/advanced_features/session_radix_cache", "docs/advanced_features/hyperparameter_tuning", "docs/advanced_features/attention_backend", "docs/advanced_features/hisparse_guide", diff --git a/docs_new/docs/advanced_features/overview.mdx b/docs_new/docs/advanced_features/overview.mdx index 804f01bd4..4681fefdc 100644 --- a/docs_new/docs/advanced_features/overview.mdx +++ b/docs_new/docs/advanced_features/overview.mdx @@ -4,6 +4,7 @@ description: Advanced configuration, optimization, and deployment features for S --- - [Server Arguments](./server_arguments) +- [Session-Aware Radix Cache](./session_radix_cache) - [Hyperparameter Tuning](./hyperparameter_tuning) - [Attention Backend](./attention_backend) - [Speculative Decoding](./speculative_decoding) diff --git a/docs_new/docs/advanced_features/session_radix_cache.mdx b/docs_new/docs/advanced_features/session_radix_cache.mdx new file mode 100644 index 000000000..4db1ed410 --- /dev/null +++ b/docs_new/docs/advanced_features/session_radix_cache.mdx @@ -0,0 +1,57 @@ +--- +title: "Session-Aware Radix Cache" +metatags: + description: "Keep active sessions ahead of unrelated KV in UnifiedRadixCache eviction order." +--- + +Session-aware radix caching improves cache hits for long-lived, multi-turn workloads under memory pressure. It registers reusable KV to a session and evicts unreferenced KV before KV still referenced by an active session. + +Session references are soft protection, not memory pins. Referenced KV can still be evicted when reclaiming unreferenced KV is insufficient. + +## Enable the cache + +This feature is implemented only by `UnifiedRadixCache`. + +```bash Command +SGLANG_ENABLE_UNIFIED_RADIX_TREE=1 python3 -m sglang.launch_server \ + --model-path MODEL_PATH \ + --enable-session-radix-cache +``` + +## Pass and close the session + +Your application must pass the same top-level `session_id` on every request in a session. The ID labels cache references only; it does not append or reconstruct conversation context, so each request must still contain the intended prompt. + +```bash Command +curl http://localhost:30000/generate \ + -H "Content-Type: application/json" \ + -d '{ + "text": "FULL_PROMPT_FOR_THIS_TURN", + "sampling_params": {"max_new_tokens": 128}, + "session_id": "agent-42" + }' +``` + +When a request finishes, SGLang automatically registers its reusable cache leaves under the `session_id`. This reference-only workflow does not require an `/open_session` call. + +Call `/close_session` when the application session ends, including error and cancellation paths: + +```bash Command +curl -X POST http://localhost:30000/close_session \ + -H "Content-Type: application/json" \ + -d '{"session_id": "agent-42"}' +``` + +Closing removes the session's references but does not immediately free its KV. The KV remains reusable and returns to the normal eviction order. + +## Eviction behavior + +The cache tracks references independently for each UnifiedRadixCache component. Device and host eviction use the same session preference. + +| Component | Referenced data | Eviction order | +| --- | --- | --- | +| Full attention | The reusable prefix path from the registered leaf to the root | Unreferenced nodes first, then referenced nodes with fewer session references, then the configured policy such as LRU | +| Sliding-window attention (SWA) | The reusable tail covering the sliding window plus page-alignment allowance | Two LRU passes: unreferenced nodes first, then referenced nodes if more space is required | +| Mamba | The reusable state on the registered leaf | Two LRU passes: unreferenced nodes first, then referenced nodes if more space is required | + +UnifiedRadixCache still applies component cascade rules. Evicting an internal Full node also evicts its SWA and Mamba data; evicting SWA also evicts Mamba data; evicting Mamba affects only Mamba. Evicting a leaf removes all component data on that leaf. diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index 9a28d739d..5642c64ef 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -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) diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index f61b1e4fd..a370c136f 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -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 diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 6589c6a78..5a0b8264f 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -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) diff --git a/python/sglang/srt/mem_cache/radix_cache.py b/python/sglang/srt/mem_cache/radix_cache.py index 4ac746313..2916f8c64 100644 --- a/python/sglang/srt/mem_cache/radix_cache.py +++ b/python/sglang/srt/mem_cache/radix_cache.py @@ -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) diff --git a/python/sglang/srt/mem_cache/registry.py b/python/sglang/srt/mem_cache/registry.py index e84a01f12..63eb13065 100644 --- a/python/sglang/srt/mem_cache/registry.py +++ b/python/sglang/srt/mem_cache/registry.py @@ -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 diff --git a/python/sglang/srt/mem_cache/session_radix_cache.py b/python/sglang/srt/mem_cache/session_radix_cache.py deleted file mode 100644 index 951c30e6f..000000000 --- a/python/sglang/srt/mem_cache/session_radix_cache.py +++ /dev/null @@ -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 diff --git a/python/sglang/srt/mem_cache/unified_cache/components/full_component.py b/python/sglang/srt/mem_cache/unified_cache/components/full_component.py index 7558254b7..3ec838c38 100644 --- a/python/sglang/srt/mem_cache/unified_cache/components/full_component.py +++ b/python/sglang/srt/mem_cache/unified_cache/components/full_component.py @@ -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( diff --git a/python/sglang/srt/mem_cache/unified_cache/components/mamba_component.py b/python/sglang/srt/mem_cache/unified_cache/components/mamba_component.py index c42b33615..a4efe3e4a 100644 --- a/python/sglang/srt/mem_cache/unified_cache/components/mamba_component.py +++ b/python/sglang/srt/mem_cache/unified_cache/components/mamba_component.py @@ -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: diff --git a/python/sglang/srt/mem_cache/unified_cache/components/swa_component.py b/python/sglang/srt/mem_cache/unified_cache/components/swa_component.py index d586714a9..d35ef0731 100644 --- a/python/sglang/srt/mem_cache/unified_cache/components/swa_component.py +++ b/python/sglang/srt/mem_cache/unified_cache/components/swa_component.py @@ -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: diff --git a/python/sglang/srt/mem_cache/unified_cache/components/tree_component.py b/python/sglang/srt/mem_cache/unified_cache/components/tree_component.py index fb8db4ec2..bc7b83a6a 100644 --- a/python/sglang/srt/mem_cache/unified_cache/components/tree_component.py +++ b/python/sglang/srt/mem_cache/unified_cache/components/tree_component.py @@ -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: diff --git a/python/sglang/srt/mem_cache/unified_cache/session_ref_tracker.py b/python/sglang/srt/mem_cache/unified_cache/session_ref_tracker.py new file mode 100644 index 000000000..07cc37248 --- /dev/null +++ b/python/sglang/srt/mem_cache/unified_cache/session_ref_tracker.py @@ -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 diff --git a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py index ff89de1ce..e0110067c 100644 --- a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py +++ b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py @@ -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: diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index 80cfd171a..de2405f89 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -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: diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index c297e8e33..36cb7c977 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -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 " diff --git a/python/sglang/srt/session/session_controller.py b/python/sglang/srt/session/session_controller.py index 0d6866a0e..789b9c3d6 100644 --- a/python/sglang/srt/session/session_controller.py +++ b/python/sglang/srt/session/session_controller.py @@ -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( diff --git a/test/manual/core/test_session_radix_cache.py b/test/manual/core/test_session_radix_cache.py deleted file mode 100644 index 2ad8c9c80..000000000 --- a/test/manual/core/test_session_radix_cache.py +++ /dev/null @@ -1,174 +0,0 @@ -"""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, - enable_session_radix_cache=True, - ) - ) - - 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_disabled_cache_does_not_tag_session_kv(self): - self.cache.enable_session_radix_cache = False - self._insert([1, 2, 3, 4]) - self._tag([1, 2, 3, 4], "S") - self.assertIsNone(getattr(self._leaf([1, 2, 3, 4]), "session_ids", None)) - - 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_radix_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_radix_session("A"), 0) # B still holds - self.assertEqual(self._cached([1, 2, 3, 4]), 4) - self.assertEqual(self.cache.release_radix_session("B"), 1) # last holder frees - self.assertEqual(self._cached([1, 2, 3, 4]), 0) - - def test_legacy_release_does_not_release_radix_session(self): - self._insert([1, 2, 3, 4]) - self._tag([1, 2, 3, 4], "S") - self.cache.release_session("S") - self.assertEqual(self._cached([1, 2, 3, 4]), 4) - self.assertEqual(self.cache.release_radix_session("S"), 1) - - 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_radix_session("S"), 0 - ) # late close is a no-op - - def test_close_tombstone_blocks_late_finish(self): - self._insert([1, 2, 3, 4]) - self._tag([1, 2, 3, 4], "S") - self.assertEqual(self.cache.release_radix_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)) - - 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_radix_session("B"), 0) - self.assertEqual(getattr(self._leaf([1, 2, 3, 4]), "session_ids", None), {"A"}) - self.assertEqual(self.cache.release_radix_session("A"), 1) - - self._insert([5, 6, 7, 8]) - self._tag([5, 6, 7, 8], "B") - self.assertIsNone(getattr(self._leaf([5, 6, 7, 8]), "session_ids", None)) - - -if __name__ == "__main__": - unittest.main() diff --git a/test/registered/unit/mem_cache/test_session_unified_radix_cache.py b/test/registered/unit/mem_cache/test_session_unified_radix_cache.py new file mode 100644 index 000000000..55dff04d4 --- /dev/null +++ b/test/registered/unit/mem_cache/test_session_unified_radix_cache.py @@ -0,0 +1,222 @@ +"""Tests for session references on UnifiedRadixCache.""" + +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + +import ast +import unittest +from array import array +from pathlib import Path +from types import SimpleNamespace + +import torch + +from sglang.srt.mem_cache.allocator import TokenToKVPoolAllocator +from sglang.srt.mem_cache.base_prefix_cache import ( + EvictParams, + InsertParams, + MatchPrefixParams, +) +from sglang.srt.mem_cache.cache_init_params import CacheInitParams +from sglang.srt.mem_cache.memory_pool import MHATokenToKVPool, ReqToTokenPool +from sglang.srt.mem_cache.radix_cache import RadixCache, RadixKey +from sglang.srt.mem_cache.unified_cache.components import ComponentType +from sglang.srt.mem_cache.unified_radix_cache import UnifiedRadixCache +from sglang.test.test_utils import CustomTestCase + +REPO_ROOT = Path(__file__).resolve().parents[4] +MEM_CACHE_ROOT = REPO_ROOT / "python/sglang/srt/mem_cache" + + +def class_bases(path: Path, class_name: str) -> set[str]: + tree = ast.parse(path.read_text()) + class_node = next( + node + for node in tree.body + if isinstance(node, ast.ClassDef) and node.name == class_name + ) + return { + base.id if isinstance(base, ast.Name) else ast.unparse(base) + for base in class_node.bases + } + + +class TestSessionCacheOwnership(CustomTestCase): + def test_only_unified_radix_cache_owns_session_ref_tracker(self): + ordinary_mixin = MEM_CACHE_ROOT / "session_radix_cache.py" + radix_cache = MEM_CACHE_ROOT / "radix_cache.py" + hiradix_cache = MEM_CACHE_ROOT / "hiradix_cache.py" + evict_policy = MEM_CACHE_ROOT / "evict_policy.py" + unified_cache = MEM_CACHE_ROOT / "unified_radix_cache.py" + session_ref_tracker = ( + MEM_CACHE_ROOT / "unified_cache" / "session_ref_tracker.py" + ) + + self.assertFalse(ordinary_mixin.exists()) + ordinary_source = "\n".join( + path.read_text() for path in (radix_cache, hiradix_cache, evict_policy) + ) + for removed_symbol in ( + "SessionRadixCacheMixin", + "SessionAwareEvictionStrategy", + "session_ref", + "_session_on_", + "_session_forget_node", + "_account_new_evictable_node", + "_supports_session_radix_cache", + "enable_session_radix_cache", + ): + self.assertNotIn(removed_symbol, ordinary_source) + self.assertNotIn( + "SessionRadixCacheMixin", class_bases(radix_cache, "RadixCache") + ) + # Session behavior is composed, not mixed in (general-code-style rule). + self.assertEqual( + class_bases(unified_cache, "UnifiedRadixCache"), {"BasePrefixCache"} + ) + self.assertIn("UnifiedSessionRefTracker", session_ref_tracker.read_text()) + self.assertNotIn("SessionUnifiedRadixCacheMixin", unified_cache.read_text()) + + for component in ( + "full_component.py", + "swa_component.py", + "mamba_component.py", + ): + self.assertIn( + "session_ref", + ( + MEM_CACHE_ROOT / "unified_cache" / "components" / component + ).read_text(), + ) + + registry = MEM_CACHE_ROOT / "registry.py" + self.assertIn( + "--enable-session-radix-cache requires UnifiedRadixCache", + registry.read_text(), + ) + + +def make_params(enable_session: bool) -> CacheInitParams: + dtype = torch.float16 + kv_pool = MHATokenToKVPool( + size=64, + page_size=1, + dtype=dtype, + head_num=2, + head_dim=8, + layer_num=1, + device="cpu", + enable_memory_saver=False, + ) + allocator = TokenToKVPoolAllocator( + size=64, + dtype=dtype, + device="cpu", + kvcache=kv_pool, + need_sort=False, + ) + req_pool = ReqToTokenPool( + size=8, + max_context_len=128, + device="cpu", + enable_memory_saver=False, + ) + return CacheInitParams( + disable=False, + req_to_token_pool=req_pool, + token_to_kv_pool_allocator=allocator, + page_size=1, + eviction_policy="lru", + enable_session_radix_cache=enable_session, + tree_components=(ComponentType.FULL,), + ) + + +def insert(cache, token_ids): + """Insert and return the tail node; the cache boundary hands back a NodeId.""" + indices = cache.token_to_kv_pool_allocator.alloc(len(token_ids)) + node_id = cache.insert( + InsertParams( + key=RadixKey(array("q", token_ids)), + value=indices.to(torch.int64), + ) + ).last_device_node + return cache.tree_core.node_by_id(node_id) + + +def match_len(cache, token_ids) -> int: + return len( + cache.match_prefix( + MatchPrefixParams(key=RadixKey(array("q", token_ids))) + ).device_indices + ) + + +def register(cache, token_ids, session_id, generation=None): + if generation is None: + generation = cache.ensure_session_generation(session_id) + cache.session_refs.register_session_ref( + SimpleNamespace( + session_id=session_id, + session_generation=generation, + session=None, + last_node=cache.match_prefix( + MatchPrefixParams(key=RadixKey(array("q", token_ids))) + ).last_device_node, + origin_input_ids=array("q", token_ids), + output_ids=array("q"), + kv_committed_len=len(token_ids), + extra_key=None, + ) + ) + + +class TestRadixCacheSessionRemoval(CustomTestCase): + def test_plain_radix_cache_does_not_enable_session_references(self): + cache = RadixCache(make_params(enable_session=True)) + + self.assertFalse(hasattr(cache, "enable_session_radix_cache")) + self.assertFalse(hasattr(cache, "register_session_ref")) + self.assertFalse(hasattr(cache, "open_radix_session")) + + +class TestSessionUnifiedRadixCache(CustomTestCase): + def setUp(self): + self.cache = UnifiedRadixCache(make_params(enable_session=True)) + self.full = self.cache.components[ComponentType.FULL] + + def test_register_and_release_update_full_component_reference(self): + leaf = insert(self.cache, [1, 2, 3, 4]) + generation = self.cache.open_radix_session("s1") + + register(self.cache, [1, 2, 3, 4], "s1", generation) + self.assertEqual(self.full.session_ref(leaf), 1) + + self.cache.release_radix_session("s1") + self.assertEqual(self.full.session_ref(leaf), 0) + + def test_reopen_rejects_stale_generation(self): + leaf = insert(self.cache, [1, 2, 3, 4]) + old_generation = self.cache.open_radix_session("s1") + self.cache.release_radix_session("s1") + self.cache.open_radix_session("s1") + + register(self.cache, [1, 2, 3, 4], "s1", old_generation) + + self.assertEqual(self.full.session_ref(leaf), 0) + + def test_eviction_prefers_unreferenced_full_kv(self): + referenced = insert(self.cache, [1, 2, 3, 4]) + insert(self.cache, [7, 8, 9]) + register(self.cache, [1, 2, 3, 4], "s1") + + self.cache.evict(EvictParams(num_tokens=3)) + + self.assertEqual(match_len(self.cache, [7, 8, 9]), 0) + self.assertEqual(match_len(self.cache, [1, 2, 3, 4]), 4) + self.assertEqual(self.full.session_ref(referenced), 1) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/mem_cache/test_tree_core_registry.py b/test/registered/unit/mem_cache/test_tree_core_registry.py index 0769dd8ce..15efe29c7 100644 --- a/test/registered/unit/mem_cache/test_tree_core_registry.py +++ b/test/registered/unit/mem_cache/test_tree_core_registry.py @@ -65,6 +65,15 @@ class _StubFullComponent(TreeComponent): def _evict_device_end(self) -> None: pass + def _dec_session_coverage(self, session_id, leaf) -> None: + pass + + def _advance_session_coverage(self, session_id, leaf, old_ancestor) -> None: + pass + + def _recede_session_coverage(self, session_id, leaf, fallback) -> None: + pass + class _StubMambaComponent(_StubFullComponent): component_type = ComponentType.MAMBA diff --git a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py index eff882fdb..8421dd38c 100644 --- a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py +++ b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py @@ -188,6 +188,15 @@ class _FakeFullComponent(TreeComponent): def _evict_device_end(self) -> None: pass + def _dec_session_coverage(self, session_id, leaf) -> None: + pass + + def _advance_session_coverage(self, session_id, leaf, old_ancestor) -> None: + pass + + def _recede_session_coverage(self, session_id, leaf, fallback) -> None: + pass + class TestUnifiedRadixComponentRegistryOverride(CustomTestCase): def test_component_registry_override_is_instance_local(self): diff --git a/test/registered/unit/server_args/test_server_args.py b/test/registered/unit/server_args/test_server_args.py index 648ffe52c..2ea9ee132 100644 --- a/test/registered/unit/server_args/test_server_args.py +++ b/test/registered/unit/server_args/test_server_args.py @@ -1274,17 +1274,6 @@ class TestPrefillOnlyDisableKvCache(unittest.TestCase): self._validate_prefill_only_args(kv_cache_dtype=kv_cache_dtype) -class TestSessionRadixCacheServerArgs(unittest.TestCase): - def test_requires_priority_radix_eviction_policy(self): - server_args = ServerArgs( - model_path="dummy", - enable_session_radix_cache=True, - radix_eviction_policy="lru", - ) - with self.assertRaisesRegex(ValueError, "--radix-eviction-policy priority"): - server_args._handle_cache_compatibility() - - class TestCudaGraphConfigDataclassAccess(CustomTestCase): @patch( "sglang.srt.model_executor.runner_backend."