diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 00972b0a3..06cd1a660 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/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: diff --git a/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py b/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py index 838670356..a88473be6 100644 --- a/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py +++ b/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: diff --git a/python/sglang/srt/mem_cache/base_prefix_cache.py b/python/sglang/srt/mem_cache/base_prefix_cache.py index d19be907f..6c66da5fe 100644 --- a/python/sglang/srt/mem_cache/base_prefix_cache.py +++ b/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 diff --git a/python/sglang/srt/mem_cache/cache_init_params.py b/python/sglang/srt/mem_cache/cache_init_params.py index 6f6fafae0..5a7c40ded 100644 --- a/python/sglang/srt/mem_cache/cache_init_params.py +++ b/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 diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index 73b645a3b..201fca688 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/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() diff --git a/python/sglang/srt/session/session_aware_cache.py b/python/sglang/srt/session/session_aware_cache.py index 9d6aabc16..c34c24651 100644 --- a/python/sglang/srt/session/session_aware_cache.py +++ b/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() diff --git a/python/sglang/srt/session/session_controller.py b/python/sglang/srt/session/session_controller.py index 00776ee31..befab0c1e 100644 --- a/python/sglang/srt/session/session_controller.py +++ b/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)})"