[core] Always-on StreamingSession in UnifiedRadixCache (#23202)

python/sglang/srt/managers/scheduler.py

@@ -205,8 +205,8 @@ from sglang.srt.parser.reasoning_parser import ReasoningParser
 from sglang.srt.plugins import load_plugins
 from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo
 from sglang.srt.server_args import PortArgs, ServerArgs, get_global_server_args
-from sglang.srt.session.session_aware_cache import SessionAwareCache
 from sglang.srt.session.session_controller import SessionController
+from sglang.srt.session.streaming_session import StreamingSession
 from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
 from sglang.srt.utils import (
     DynamicGradMode,
@@ -865,7 +865,6 @@ 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
@@ -894,7 +893,7 @@ class Scheduler(
             server_args.enable_streaming_session
             and not self.tree_cache.supports_streaming_session()
         ):
-            self.tree_cache = SessionAwareCache(self.tree_cache)
+            self.tree_cache = StreamingSession(self.tree_cache)
 
         if self.enable_hisparse:
             # Coordinator was created inside ModelRunner.initialize() before CUDA graph capture

python/sglang/srt/mem_cache/cache_init_params.py

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

python/sglang/srt/mem_cache/common.py

@@ -478,7 +478,7 @@ def release_kv_cache(req: Req, tree_cache: BasePrefixCache, is_insert: bool = Tr
 
     tree_cache.cache_finished_req(req, is_insert=is_insert)
 
-    # SessionAwareCache.cache_finished_req handles speculative tail trim
+    # StreamingSession.cache_finished_req handles speculative tail trim
     # and bookkeeping flag sync internally, then sets req_pool_idx = None.
     if req.req_pool_idx is None:
         return

python/sglang/srt/mem_cache/unified_radix_cache.py

@@ -40,7 +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
+from sglang.srt.session.streaming_session import StreamingSession
 
 if TYPE_CHECKING:
     from sglang.srt.managers.schedule_batch import Req
@@ -210,12 +210,12 @@ class UnifiedRadixCache(BasePrefixCache):
         else:
             self.key_convert_fn = lambda key: key
 
-        # Streaming session: embedded SessionAwareCache with self as inner.
+        # Streaming session: embedded StreamingSession with self as inner.
+        # Always on -- zero overhead when no streaming session is open (the
+        # try_* entries short-circuit on non-streaming reqs / real TreeNodes).
         # 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.session = StreamingSession(inner=self)
 
         self.reset()
         logger.info(f"Init Unified RadixTree with components {self.tree_components}")
@@ -231,14 +231,12 @@ 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()
+        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
+        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)
@@ -289,10 +287,9 @@ class UnifiedRadixCache(BasePrefixCache):
         )
 
     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
+        result = self.session.try_inc_lock_ref(node)
+        if result is not None:
+            return result
         if self.disable:
             return IncLockRefResult()
         result = IncLockRefResult()
@@ -303,10 +300,9 @@ class UnifiedRadixCache(BasePrefixCache):
     def dec_lock_ref(
         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
+        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:
@@ -315,9 +311,7 @@ class UnifiedRadixCache(BasePrefixCache):
         return DecLockRefResult()
 
     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
-        ):
+        if self.session.try_cache_finished_req(req, is_insert=is_insert, **kwargs):
             return
 
         kv_committed_len = req.pop_committed_kv_cache()
@@ -389,9 +383,7 @@ class UnifiedRadixCache(BasePrefixCache):
             )
 
     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
-        ):
+        if self.session.try_cache_unfinished_req(req, chunked=chunked, **kwargs):
             return
 
         token_ids = req.fill_ids
@@ -813,33 +805,24 @@ class UnifiedRadixCache(BasePrefixCache):
     def supports_mamba(self) -> bool:
         return ComponentType.MAMBA in self.components
 
-    # ---- Streaming session API (delegates to composed SessionImpl) ----
+    # ---- Streaming session API (delegates to composed StreamingSession) ----
 
     def supports_streaming_session(self) -> bool:
-        return self.session is not None
+        return True
 
     def release_session(self, session_id: str) -> None:
-        if self.session is not None:
-            self.session.release_session(session_id)
+        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:
@@ -964,7 +947,7 @@ class UnifiedRadixCache(BasePrefixCache):
         # 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():
+        if self.session.any_holding_kv():
             return
         try:
             # 1. Collect all nodes from tree

python/sglang/srt/server_args.py

@@ -4485,7 +4485,7 @@ class ServerArgs:
             "--enable-streaming-session",
             action="store_true",
             default=ServerArgs.enable_streaming_session,
-            help="Enable streaming session mode and SessionAwareCache wrapper.",
+            help="Enable streaming session mode and StreamingSession wrapper.",
         )
         parser.add_argument(
             "--random-seed",

python/sglang/srt/session/streaming_session.py

@@ -113,11 +113,11 @@ def _is_streaming(req: Optional[Req]) -> bool:
     return req is not None and req.session is not None and req.session.streaming
 
 
-class SessionAwareCache(BasePrefixCache):
+class StreamingSession(BasePrefixCache):
     """Adds streaming-session KV save/restore on top of any BasePrefixCache.
 
-    Works both as an external wrapper (``SessionAwareCache(RadixCache(...))``)
-    and in embedded composition (``SessionAwareCache(inner=self)``). For the
+    Works both as an external wrapper (``StreamingSession(RadixCache(...))``)
+    and in embedded composition (``StreamingSession(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.