integrate streaming session into UnifiedRadixCache (#23145)

python/sglang/srt/managers/scheduler.py

@@ -865,6 +865,7 @@ class Scheduler(
                         ComponentType.SWA if self.is_hybrid_swa else ComponentType.MAMBA
                     )
                 params.tree_components = tuple(tree_components)
+                params.enable_streaming_session = server_args.enable_streaming_session
                 self.tree_cache = UnifiedRadixCache(params)
             elif self.is_hybrid_swa:
                 from sglang.srt.mem_cache.swa_radix_cache import SWARadixCache
@@ -889,7 +890,10 @@ class Scheduler(
             else:
                 self.tree_cache = RadixCache(params)
 
-        if server_args.enable_streaming_session:
+        if (
+            server_args.enable_streaming_session
+            and not self.tree_cache.supports_streaming_session()
+        ):
             self.tree_cache = SessionAwareCache(self.tree_cache)
 
         if self.enable_hisparse:

python/sglang/srt/managers/scheduler_runtime_checker_mixin.py

@@ -9,7 +9,6 @@ from typing import TYPE_CHECKING, List, Optional, Tuple
 from sglang.srt.disaggregation.utils import DisaggregationMode
 from sglang.srt.environ import envs
 from sglang.srt.observability.metrics_collector import QueueCount
-from sglang.srt.session.session_aware_cache import SessionAwareCache
 from sglang.srt.utils.common import ceil_align, raise_error_or_warn
 from sglang.srt.utils.request_logger import disable_request_logging
 from sglang.srt.utils.watchdog import WatchdogRaw
@@ -157,24 +156,16 @@ class SchedulerRuntimeCheckerMixin:
         return idxs
 
     def _session_held_tokens(self: Scheduler) -> int:
-        if isinstance(self.tree_cache, SessionAwareCache):
-            return self.tree_cache.session_held_tokens(self._active_pool_idxs())
-        return 0
+        return self.tree_cache.session_held_tokens(self._active_pool_idxs())
 
     def _session_held_full_tokens(self: Scheduler) -> int:
-        if isinstance(self.tree_cache, SessionAwareCache):
-            return self.tree_cache.session_held_full_tokens(self._active_pool_idxs())
-        return 0
+        return self.tree_cache.session_held_full_tokens(self._active_pool_idxs())
 
     def _session_held_swa_tokens(self: Scheduler) -> int:
-        if isinstance(self.tree_cache, SessionAwareCache):
-            return self.tree_cache.session_held_swa_tokens(self._active_pool_idxs())
-        return 0
+        return self.tree_cache.session_held_swa_tokens(self._active_pool_idxs())
 
     def _session_held_req_count(self: Scheduler) -> int:
-        if isinstance(self.tree_cache, SessionAwareCache):
-            return self.tree_cache.session_held_req_count()
-        return 0
+        return self.tree_cache.session_held_req_count()
 
     def get_pool_stats(self: Scheduler) -> PoolStats:
         if self.is_hybrid_swa:

python/sglang/srt/mem_cache/base_prefix_cache.py

@@ -262,6 +262,24 @@ class BasePrefixCache(ABC, PrefixCacheTrait):
     def supports_mamba(self) -> bool:
         return False
 
+    def supports_streaming_session(self) -> bool:
+        return False
+
+    def release_session(self, session_id: str) -> None:
+        pass
+
+    def session_held_tokens(self, active_pool_idxs: Optional[set] = None) -> int:
+        return 0
+
+    def session_held_full_tokens(self, active_pool_idxs: Optional[set] = None) -> int:
+        return 0
+
+    def session_held_swa_tokens(self, active_pool_idxs: Optional[set] = None) -> int:
+        return 0
+
+    def session_held_req_count(self, active_pool_idxs: Optional[set] = None) -> int:
+        return 0
+
     def is_chunk_cache(self) -> bool:
         return False
 

python/sglang/srt/mem_cache/cache_init_params.py

@@ -42,3 +42,5 @@ class CacheInitParams:
     cache_ttl_seconds: Optional[float] = None
 
     tree_components: Optional[tuple[ComponentType, ...]] = None
+
+    enable_streaming_session: bool = False

python/sglang/srt/mem_cache/unified_radix_cache.py

@@ -4,7 +4,7 @@ import logging
 import time
 from collections import defaultdict
 from functools import partial
-from typing import TYPE_CHECKING, Optional
+from typing import TYPE_CHECKING, Any, Optional
 
 import torch
 
@@ -40,6 +40,7 @@ from sglang.srt.mem_cache.unified_cache_components import (
     get_and_increase_time_counter,
 )
 from sglang.srt.mem_cache.utils import convert_to_bigram_key
+from sglang.srt.session.session_aware_cache import SessionAwareCache
 
 if TYPE_CHECKING:
     from sglang.srt.managers.schedule_batch import Req
@@ -208,6 +209,14 @@ class UnifiedRadixCache(BasePrefixCache):
             self.key_convert_fn = convert_to_bigram_key
         else:
             self.key_convert_fn = lambda key: key
+
+        # Streaming session: embedded SessionAwareCache with self as inner.
+        # Dispatch methods below pre-check conditions so the session's
+        # internal fall-through to self.inner.xxx never fires -- no recursion.
+        self.session: Optional[SessionAwareCache] = (
+            SessionAwareCache(inner=self) if params.enable_streaming_session else None
+        )
+
         self.reset()
         logger.info(f"Init Unified RadixTree with components {self.tree_components}")
 
@@ -222,8 +231,15 @@ class UnifiedRadixCache(BasePrefixCache):
         self.lru_lists = {
             ct: UnifiedLRUList(ct, self.tree_components) for ct in self.tree_components
         }
+        if self.session is not None:
+            self.session.slots.clear()
 
     def match_prefix(self, params: MatchPrefixParams) -> MatchResult:
+        if self.session is not None:
+            result = self.session.try_match_prefix(params)
+            if result is not None:
+                return result
+
         key = params.key
         key, _ = maybe_bigram_convert(self.is_eagle, key)
         if self.disable or len(key) == 0:
@@ -272,7 +288,11 @@ class UnifiedRadixCache(BasePrefixCache):
             mamba_num_evicted=tracker.get(ComponentType.MAMBA, 0),
         )
 
-    def inc_lock_ref(self, node: UnifiedTreeNode) -> IncLockRefResult:
+    def inc_lock_ref(self, node: Any) -> IncLockRefResult:
+        if self.session is not None:
+            result = self.session.try_inc_lock_ref(node)
+            if result is not None:
+                return result
         if self.disable:
             return IncLockRefResult()
         result = IncLockRefResult()
@@ -281,8 +301,12 @@ class UnifiedRadixCache(BasePrefixCache):
         return result
 
     def dec_lock_ref(
-        self, node: UnifiedTreeNode, params: Optional[DecLockRefParams] = None
+        self, node: Any, params: Optional[DecLockRefParams] = None
     ) -> DecLockRefResult:
+        if self.session is not None:
+            result = self.session.try_dec_lock_ref(node, params)
+            if result is not None:
+                return result
         if self.disable:
             return DecLockRefResult()
         for component in self._components_tuple:
@@ -290,7 +314,12 @@ class UnifiedRadixCache(BasePrefixCache):
         # TODO: delta is not aggregated from components; no caller uses it yet.
         return DecLockRefResult()
 
-    def cache_finished_req(self, req: Req, is_insert: bool = True) -> None:
+    def cache_finished_req(self, req: Req, is_insert: bool = True, **kwargs) -> None:
+        if self.session is not None and self.session.try_cache_finished_req(
+            req, is_insert=is_insert, **kwargs
+        ):
+            return
+
         kv_committed_len = req.pop_committed_kv_cache()
 
         if self.disable:
@@ -359,7 +388,12 @@ class UnifiedRadixCache(BasePrefixCache):
                 req, is_finished=True, insert_result=result, insert_params=insert_params
             )
 
-    def cache_unfinished_req(self, req: Req, chunked=False) -> None:
+    def cache_unfinished_req(self, req: Req, chunked=False, **kwargs) -> None:
+        if self.session is not None and self.session.try_cache_unfinished_req(
+            req, chunked=chunked, **kwargs
+        ):
+            return
+
         token_ids = req.fill_ids
 
         if self.disable:
@@ -779,6 +813,35 @@ class UnifiedRadixCache(BasePrefixCache):
     def supports_mamba(self) -> bool:
         return ComponentType.MAMBA in self.components
 
+    # ---- Streaming session API (delegates to composed SessionImpl) ----
+
+    def supports_streaming_session(self) -> bool:
+        return self.session is not None
+
+    def release_session(self, session_id: str) -> None:
+        if self.session is not None:
+            self.session.release_session(session_id)
+
+    def session_held_tokens(self, active_pool_idxs: Optional[set] = None) -> int:
+        if self.session is None:
+            return 0
+        return self.session.session_held_tokens(active_pool_idxs)
+
+    def session_held_full_tokens(self, active_pool_idxs: Optional[set] = None) -> int:
+        if self.session is None:
+            return 0
+        return self.session.session_held_full_tokens(active_pool_idxs)
+
+    def session_held_swa_tokens(self, active_pool_idxs: Optional[set] = None) -> int:
+        if self.session is None:
+            return 0
+        return self.session.session_held_swa_tokens(active_pool_idxs)
+
+    def session_held_req_count(self, active_pool_idxs: Optional[set] = None) -> int:
+        if self.session is None:
+            return 0
+        return self.session.session_held_req_count(active_pool_idxs)
+
     def evictable_size(self) -> int:
         return self.component_evictable_size_.get(BASE_COMPONENT_TYPE, 0)
 
@@ -898,6 +961,11 @@ class UnifiedRadixCache(BasePrefixCache):
         """Thorough sanity check: verify LRU membership, lock state, linked-list
         integrity, and evictable sizes for every component.
         Expensive — use only in tests or idle checks."""
+        # Skip when streaming sessions hold tree locks: the check asserts
+        # all nodes are unlocked during idle, which streaming sessions break
+        # by design (they hold a first-turn lock across turns).
+        if self.session is not None and self.session.any_holding_kv():
+            return
         try:
             # 1. Collect all nodes from tree
             all_nodes = self._collect_all_nodes()

python/sglang/srt/session/session_aware_cache.py

@@ -29,7 +29,7 @@ logger = logging.getLogger(__name__)
 class _VirtualNode:
     """Sentinel node for streaming session requests.
 
-    Passed to inc_lock_ref / dec_lock_ref so the wrapper can distinguish
+    Passed to inc_lock_ref / dec_lock_ref so the cache can distinguish
     streaming-session locks (no-op) from real radix-tree locks (forwarded).
     """
 
@@ -114,11 +114,13 @@ def _is_streaming(req: Optional[Req]) -> bool:
 
 
 class SessionAwareCache(BasePrefixCache):
-    """Decorator around any BasePrefixCache that manages streaming session KV.
+    """Adds streaming-session KV save/restore on top of any BasePrefixCache.
 
-    Non-streaming requests are pure pass-through. Streaming requests have their
-    KV lifecycle managed by SessionSlot objects, avoiding any invasive changes
-    to the scheduling pipeline.
+    Works both as an external wrapper (``SessionAwareCache(RadixCache(...))``)
+    and in embedded composition (``SessionAwareCache(inner=self)``). For the
+    embedded case, the composing cache must pre-check dispatch conditions
+    (``_is_streaming`` / ``find_active_slot`` / ``has_slot``) so the internal
+    fall-through to ``self.inner.xxx`` never fires -- otherwise it recurses.
     """
 
     def __init__(self, inner: BasePrefixCache):
@@ -167,30 +169,67 @@ class SessionAwareCache(BasePrefixCache):
     def metrics_collector(self, value):
         self.inner.metrics_collector = value
 
+    # -- Condition helpers (used by embedded-mode callers for pre-dispatch) --
+
+    def has_slot(self, session_id: str) -> bool:
+        return session_id in self.slots
+
+    def any_holding_kv(self) -> bool:
+        return any(s.is_holding_kv for s in self.slots.values())
+
+    # -- Try-handle entries for composition (see class docstring) --
+
+    def try_inc_lock_ref(self, node: Any) -> Optional[IncLockRefResult]:
+        """No-op lock if ``node`` is a session-internal sentinel; returns
+        None to tell the caller to run its raw tree lock path."""
+        if isinstance(node, _VirtualNode):
+            return IncLockRefResult()
+        return None
+
+    def try_dec_lock_ref(
+        self, node: Any, params: Optional[DecLockRefParams] = None
+    ) -> Optional[DecLockRefResult]:
+        if isinstance(node, _VirtualNode):
+            return DecLockRefResult()
+        return None
+
+    def find_active_slot(self, req: Req) -> Optional[SessionSlot]:
+        """Returns an active slot for this req, or None.
+
+        Side effect: if req is pre-aborted (to_finish set, e.g. input too
+        long), detach it from the session so cache_finished_req treats it
+        as a normal req. The slot stays intact for the next request.
+        """
+        if not _is_streaming(req):
+            return None
+        slot = self.slots.get(req.session.session_id)
+        if slot is None or slot.req_pool_idx is None:
+            return None
+        if req.to_finish is not None:
+            req.session.abort_req()
+            req.session = None
+            return None
+        return slot
+
     # -- BasePrefixCache abstract methods --
 
     def reset(self):
         self.slots.clear()
         self.inner.reset()
 
-    def match_prefix(self, params: MatchPrefixParams) -> MatchResult:
+    # -- Streaming entries: contract with embedded composers (e.g.
+    # UnifiedRadixCache) is a uniform "try_handle_*" pattern. Each method
+    # executes the streaming body if applicable and signals whether the
+    # caller still needs to run its raw path.
+
+    def try_match_prefix(self, params: MatchPrefixParams) -> Optional[MatchResult]:
+        """Returns a MatchResult iff the request hits an active session slot;
+        otherwise None (caller falls back to its raw match)."""
+        slot = self.find_active_slot(params.req)
+        if slot is None:
+            return None
+
         req = params.req
-        if not _is_streaming(req):
-            return self.inner.match_prefix(params)
-
-        session_id = req.session.session_id
-        slot = self.slots.get(session_id)
-        if slot is None or slot.req_pool_idx is None:
-            return self.inner.match_prefix(params)
-
-        # Pre-aborted req (scheduler-level abort, e.g. input too long):
-        # detach from session so cache_finished_req treats it as a normal
-        # req. The slot stays intact for the next request.
-        if req.to_finish is not None:
-            req.session.abort_req()
-            req.session = None
-            return self.inner.match_prefix(params)
-
         slot.restore_to_req(req)
 
         # token_ids = fill_ids[:input_len-1] (1-token logit reserve already
@@ -223,9 +262,13 @@ class SessionAwareCache(BasePrefixCache):
             cache_protected_len=slot.cache_protected_len,
         )
 
-    def cache_finished_req(self, req: Req, is_insert: bool = True, **kwargs):
+    def try_cache_finished_req(
+        self, req: Req, is_insert: bool = True, **kwargs
+    ) -> bool:
+        """Handles a streaming-session finish (save slot / mid-abort nuke).
+        Returns True if handled; False means caller runs its raw path."""
         if not _is_streaming(req):
-            return self.inner.cache_finished_req(req, is_insert=is_insert, **kwargs)
+            return False
 
         from sglang.srt.managers.schedule_batch import FINISH_ABORT
 
@@ -234,7 +277,7 @@ class SessionAwareCache(BasePrefixCache):
         is_first = slot is None
 
         # Mid-processing abort only. Pre-aborted reqs have session=None
-        # (set in match_prefix) and never reach here.
+        # (set in find_active_slot) and never reach here.
         # Nuke all KV via release_session, delete slot. Token IDs stay
         # in req_nodes (finish_req was never called -> last successful
         # req). Next request re-prefills from scratch.
@@ -257,7 +300,7 @@ class SessionAwareCache(BasePrefixCache):
             req.req_pool_idx = None
             req.session.abort_req()
             self._mark_kv_freed(req)
-            return
+            return True
 
         if is_first:
             slot = SessionSlot()
@@ -274,94 +317,66 @@ class SessionAwareCache(BasePrefixCache):
         req.session.finish_req(req)
 
         self._mark_kv_freed(req)
+        return True
 
-    def _free_tail(self, slot: SessionSlot, req: Req, prefix_len: int):
-        """match_prefix path: free orphaned KV in [prefix_len, kv_allocated_len)
-        before alloc_for_extend overwrites it. The gap appears when spec
-        decoding pushes allocated above committed, or when retract retry's
-        logit-reserve pulls prefix_len below committed.
-        """
-        self._free_kv_aligned(slot.req_pool_idx, prefix_len, slot.kv_allocated_len)
-        slot.kv_allocated_len = prefix_len
-        slot.kv_committed_len = min(slot.kv_committed_len, prefix_len)
-        slot.swa_evicted_seqlen = min(slot.swa_evicted_seqlen, prefix_len)
-        req.kv_allocated_len = prefix_len
-        req.kv_committed_len = min(req.kv_committed_len, prefix_len)
-        req.swa_evicted_seqlen = min(req.swa_evicted_seqlen, prefix_len)
+    def try_cache_unfinished_req(
+        self, req: Req, chunked: bool = False, **kwargs
+    ) -> bool:
+        """Handles a streaming-session mid-flight cache op:
+          - chunked prefill: snapshot current KV as prefix, skip radix
+          - subsequent turn: skip radix (slot already holds KV)
+        Returns False for first-turn non-chunked (caller must run raw radix
+        insert to set up the initial tree lock)."""
+        if not _is_streaming(req):
+            return False
+        if chunked:
+            kv_indices = self.req_to_token_pool.req_to_token[
+                req.req_pool_idx, : len(req.fill_ids)
+            ]
+            req.prefix_indices = kv_indices.to(dtype=torch.int64, copy=True)
+            return True
+        if req.session.session_id in self.slots:
+            return True
+        return False
 
-    def _trim_overshoot(self, req: Req, finished_len: int):
-        """Trim slot KV to finished_len boundary. Spec v2 may overshoot
-        max_new_tokens (verify round commits M+1 at a time); next turn's
-        input is output_ids[:finished_len], so positions past that must
-        be released to avoid token/KV mismatch.
-        """
-        target = len(req.origin_input_ids) + finished_len
-        self._free_kv_aligned(req.req_pool_idx, target, req.kv_allocated_len)
-        req.kv_allocated_len = min(req.kv_allocated_len, target)
-        req.kv_committed_len = min(req.kv_committed_len, target)
-        req.swa_evicted_seqlen = min(req.swa_evicted_seqlen, target)
-        req.output_ids = req.output_ids[:finished_len]
+    # -- BasePrefixCache abstract methods: thin adapters over try_handle_* --
 
-    def _free_kv_aligned(self, pool_idx: int, target: int, end: int):
-        """Free req_to_token[pool_idx, ceil_align(target):end). Page-aligned
-        because PagedTokenToKVPoolAllocator.free returns whole pages
-        (free_index // page_size), so partial-page free would corrupt pages
-        still holding committed tokens. The range [target, ceil_align(target))
-        stays attached until release_session frees the whole page.
-        """
-        if end <= target:
+    def match_prefix(self, params: MatchPrefixParams) -> MatchResult:
+        result = self.try_match_prefix(params)
+        if result is not None:
+            return result
+        return self.inner.match_prefix(params)
+
+    def cache_finished_req(self, req: Req, is_insert: bool = True, **kwargs):
+        if self.try_cache_finished_req(req, is_insert=is_insert, **kwargs):
             return
-        start = target
-        if self.page_size > 1:
-            start = ceil_align(start, self.page_size)
-        if start < end:
-            tail = self.req_to_token_pool.req_to_token[pool_idx, start:end]
-            self.token_to_kv_pool_allocator.free(tail)
-
-    @staticmethod
-    def _mark_kv_freed(req: Req):
-        """Set bookkeeping flags so busy check skips this finished req."""
-        if not req.kv_committed_freed:
-            req.pop_committed_kv_cache()
-        if not req.kv_overallocated_freed:
-            req.pop_overallocated_kv_cache()
+        self.inner.cache_finished_req(req, is_insert=is_insert, **kwargs)
 
     def cache_unfinished_req(self, req: Req, **kwargs):
-        if _is_streaming(req):
-            # in chunked_prefill for streaming, we skip the stash path which triggers radix.
-            # only the last chunk in first turn trigger a full prompt radix insert.
-            if kwargs.get("chunked", False):
-                kv_indices = self.req_to_token_pool.req_to_token[
-                    req.req_pool_idx, : len(req.fill_ids)
-                ]
-                req.prefix_indices = kv_indices.to(dtype=torch.int64, copy=True)
-                return
-            if req.session.session_id in self.slots:
-                # Subsequent turns: slot exists, skip inner entirely.
-                return
-            # First turn (no slot): fall through to inner for lock management,
-            # tree insertion, and cache_protected_len updates between chunks.
+        if self.try_cache_unfinished_req(req, **kwargs):
+            return
         self.inner.cache_unfinished_req(req, **kwargs)
 
     def evict(self, params: EvictParams) -> EvictResult:
         return self.inner.evict(params)
 
     def inc_lock_ref(self, node: Any) -> IncLockRefResult:
-        if isinstance(node, _VirtualNode):
-            return IncLockRefResult()
+        result = self.try_inc_lock_ref(node)
+        if result is not None:
+            return result
         return self.inner.inc_lock_ref(node)
 
     def dec_lock_ref(
         self, node: Any, params: Optional[DecLockRefParams] = None
     ) -> DecLockRefResult:
-        if isinstance(node, _VirtualNode):
-            return DecLockRefResult()
+        result = self.try_dec_lock_ref(node, params)
+        if result is not None:
+            return result
         return self.inner.dec_lock_ref(node, params)
 
     # -- Session lifecycle --
 
-    def release_session(self, session_id: str):
-        """Release all KV resources held by a streaming session."""
+    def release_session(self, session_id: str) -> None:
         slot = self.slots.pop(session_id, None)
         if slot is None:
             return
@@ -396,7 +411,7 @@ class SessionAwareCache(BasePrefixCache):
     def session_held_tokens(self, active_pool_idxs: Optional[set] = None) -> int:
         """Total KV tokens held by session slots, not tracked by the tree.
 
-        Excludes slots whose KV is currently owned by an owning request —
+        Excludes slots whose KV is currently owned by an owning request --
         those tokens are counted via uncached_size in the busy mem check.
         A slot's pool_idx being in active_pool_idxs indicates a req owns it.
         """
@@ -439,6 +454,59 @@ class SessionAwareCache(BasePrefixCache):
 
         return sum(_owned(s) for s in self.slots.values())
 
+    # -- Internal helpers (streaming body bits) --
+
+    def _free_tail(self, slot: SessionSlot, req: Req, prefix_len: int) -> None:
+        """match_prefix path: free orphaned KV in [prefix_len, kv_allocated_len)
+        before alloc_for_extend overwrites it. The gap appears when spec
+        decoding pushes allocated above committed, or when retract retry's
+        logit-reserve pulls prefix_len below committed.
+        """
+        self._free_kv_aligned(slot.req_pool_idx, prefix_len, slot.kv_allocated_len)
+        slot.kv_allocated_len = prefix_len
+        slot.kv_committed_len = min(slot.kv_committed_len, prefix_len)
+        slot.swa_evicted_seqlen = min(slot.swa_evicted_seqlen, prefix_len)
+        req.kv_allocated_len = prefix_len
+        req.kv_committed_len = min(req.kv_committed_len, prefix_len)
+        req.swa_evicted_seqlen = min(req.swa_evicted_seqlen, prefix_len)
+
+    def _trim_overshoot(self, req: Req, finished_len: int) -> None:
+        """Trim slot KV to finished_len boundary. Spec v2 may overshoot
+        max_new_tokens (verify round commits M+1 at a time); next turn's
+        input is output_ids[:finished_len], so positions past that must
+        be released to avoid token/KV mismatch.
+        """
+        target = len(req.origin_input_ids) + finished_len
+        self._free_kv_aligned(req.req_pool_idx, target, req.kv_allocated_len)
+        req.kv_allocated_len = min(req.kv_allocated_len, target)
+        req.kv_committed_len = min(req.kv_committed_len, target)
+        req.swa_evicted_seqlen = min(req.swa_evicted_seqlen, target)
+        req.output_ids = req.output_ids[:finished_len]
+
+    def _free_kv_aligned(self, pool_idx: int, target: int, end: int) -> None:
+        """Free req_to_token[pool_idx, ceil_align(target):end). Page-aligned
+        because PagedTokenToKVPoolAllocator.free returns whole pages
+        (free_index // page_size), so partial-page free would corrupt pages
+        still holding committed tokens. The range [target, ceil_align(target))
+        stays attached until release_session frees the whole page.
+        """
+        if end <= target:
+            return
+        start = target
+        if self.page_size > 1:
+            start = ceil_align(start, self.page_size)
+        if start < end:
+            tail = self.req_to_token_pool.req_to_token[pool_idx, start:end]
+            self.token_to_kv_pool_allocator.free(tail)
+
+    @staticmethod
+    def _mark_kv_freed(req: Req) -> None:
+        """Set bookkeeping flags so busy check skips this finished req."""
+        if not req.kv_committed_freed:
+            req.pop_committed_kv_cache()
+        if not req.kv_overallocated_freed:
+            req.pop_overallocated_kv_cache()
+
     # -- Pass-through methods --
 
     def evictable_size(self):
@@ -486,6 +554,9 @@ class SessionAwareCache(BasePrefixCache):
     def supports_mamba(self):
         return self.inner.supports_mamba()
 
+    def supports_streaming_session(self) -> bool:
+        return True
+
     def is_chunk_cache(self):
         return self.inner.is_chunk_cache()
 
@@ -501,7 +572,7 @@ class SessionAwareCache(BasePrefixCache):
     def sanity_check(self):
         # Skip inner sanity check when sessions hold tree locks, because
         # the check asserts all nodes are unlocked during idle.
-        if any(s.is_holding_kv for s in self.slots.values()):
+        if self.any_holding_kv():
             return
         self.inner.sanity_check()
 

python/sglang/srt/session/session_controller.py

@@ -24,7 +24,6 @@ from sglang.srt.managers.io_struct import (
     TokenizedGenerateReqInput,
 )
 from sglang.srt.managers.schedule_batch import FINISH_ABORT, Req
-from sglang.srt.session.session_aware_cache import SessionAwareCache
 from sglang.srt.utils.common import log_info_on_rank0
 
 if TYPE_CHECKING:
@@ -349,8 +348,7 @@ class SessionController:
                 mm.release_features()
             node.req.multimodal_inputs = None
 
-        if isinstance(self.tree_cache, SessionAwareCache):
-            self.tree_cache.release_session(session_id)
+        self.tree_cache.release_session(session_id)
         del self.sessions[session_id]
         log_info_on_rank0(
             logger, f"Session closed: {session_id} (active={len(self.sessions)})"