from __future__ import annotations import logging from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, Dict, Optional import torch from sglang.srt.mem_cache.base_prefix_cache import ( BasePrefixCache, DecLockRefParams, DecLockRefResult, EvictParams, EvictResult, IncLockRefResult, InitLoadBackParams, MatchPrefixParams, MatchResult, ) from sglang.srt.utils.common import ceil_align if TYPE_CHECKING: from sglang.srt.managers.schedule_batch import Req logger = logging.getLogger(__name__) class _VirtualNode: """Sentinel node for streaming session requests. Passed to inc_lock_ref / dec_lock_ref so the cache can distinguish streaming-session locks (no-op) from real radix-tree locks (forwarded). """ pass @dataclass class SessionSlot: """Holds KV state between streaming session turns.""" virtual_node: _VirtualNode = field(default_factory=_VirtualNode) # KV pool state (None means no KV is currently held by this slot) req_pool_idx: Optional[int] = None kv_committed_len: int = 0 kv_allocated_len: int = 0 # First req's radix tree node (for dec_lock_ref on session close) last_node: Any = None cache_protected_len: int = 0 swa_uuid_for_lock: Optional[str] = None # SWA state swa_evicted_seqlen: int = 0 # Mamba states mamba_pool_idx: Any = None mamba_ping_pong_track_buffer: Any = None mamba_next_track_idx: Any = None mamba_last_track_seqlen: Any = None mamba_branching_seqlen: Any = None @property def is_holding_kv(self) -> bool: """Whether this slot currently holds KV pool resources.""" return self.req_pool_idx is not None def save_from_req(self, req: Req, is_first: bool): """Save KV state from a finishing request into this slot.""" self.req_pool_idx = req.req_pool_idx self.kv_committed_len = req.kv_committed_len self.kv_allocated_len = req.kv_allocated_len self.swa_evicted_seqlen = req.swa_evicted_seqlen if is_first: self.last_node = req.last_node self.cache_protected_len = req.cache_protected_len self.swa_uuid_for_lock = req.swa_uuid_for_lock self.mamba_pool_idx = req.mamba_pool_idx self.mamba_ping_pong_track_buffer = req.mamba_ping_pong_track_buffer self.mamba_next_track_idx = req.mamba_next_track_idx self.mamba_last_track_seqlen = req.mamba_last_track_seqlen self.mamba_branching_seqlen = req.mamba_branching_seqlen # Ownership has transferred to the slot. Null *all* of the req's # references so any later alloc()/free path that inspects the req # (e.g. the alloc-skip check on `req.mamba_ping_pong_track_buffer # is None`, or the retract cleanup) sees no dangling pointers # into slot-owned tensors. Without this the alloc path can decide # the req still has a ping-pong buffer and skip alloc, causing # the slot's tensor to be reused by a new req and leaked when # the slot is later freed. req.req_pool_idx = None req.mamba_pool_idx = None req.mamba_ping_pong_track_buffer = None req.mamba_next_track_idx = None req.mamba_last_track_seqlen = None req.mamba_branching_seqlen = None def restore_to_req(self, req: Req): """Restore KV state from this slot into an incoming request.""" req.req_pool_idx = self.req_pool_idx req.kv_committed_len = self.kv_committed_len req.kv_allocated_len = self.kv_allocated_len req.swa_evicted_seqlen = self.swa_evicted_seqlen req.swa_uuid_for_lock = self.swa_uuid_for_lock req.mamba_pool_idx = self.mamba_pool_idx req.mamba_ping_pong_track_buffer = self.mamba_ping_pong_track_buffer req.mamba_next_track_idx = self.mamba_next_track_idx req.mamba_last_track_seqlen = self.mamba_last_track_seqlen req.mamba_branching_seqlen = self.mamba_branching_seqlen # NOTE: req_pool_idx and mamba_pool_idx are intentionally NOT cleared # from the slot. During chunked prefill, a request may be rejected by # the scheduler (e.g. budget exhausted) and retried in the next cycle. # Each retry calls match_prefix -> restore_to_req again, so the slot # must remain intact for idempotent restoration. def _is_streaming(req: Optional[Req]) -> bool: return req is not None and req.session is not None and req.session.streaming class StreamingSession(BasePrefixCache): """Adds streaming-session KV save/restore on top of any BasePrefixCache. 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. """ def __init__(self, inner: BasePrefixCache): self.inner = inner self.slots: Dict[str, SessionSlot] = {} # -- Forward PrefixCacheTrait properties to inner cache -- @property def req_to_token_pool(self): return self.inner.req_to_token_pool @req_to_token_pool.setter def req_to_token_pool(self, value): self.inner.req_to_token_pool = value @property def token_to_kv_pool_allocator(self): return self.inner.token_to_kv_pool_allocator @token_to_kv_pool_allocator.setter def token_to_kv_pool_allocator(self, value): self.inner.token_to_kv_pool_allocator = value @property def page_size(self): return self.inner.page_size @page_size.setter def page_size(self, value): self.inner.page_size = value @property def disable(self): return self.inner.disable @disable.setter def disable(self, value): self.inner.disable = value @property def metrics_collector(self): return self.inner.metrics_collector @metrics_collector.setter 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() # -- 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 slot.restore_to_req(req) # token_ids = get_fill_ids()[:input_len-1] (1-token logit reserve # already applied). min handles retract retry where committed_len # can exceed len(token_ids) by 1. prefix_len = min(req.kv_committed_len, len(params.key)) # Streaming sessions are append-only (session_controller rollback # ensures req_nodes always points to the last successful req). assert prefix_len >= slot.cache_protected_len, ( f"streaming session prefix shrank: {prefix_len=} < " f"{slot.cache_protected_len=}" ) # Free orphaned tail: alloc_for_extend will overwrite # req_to_token[prefix_len:] with new indices. The range # [prefix_len, kv_allocated_len) has stale indices from the # previous turn's decode (e.g. alloc-commit gap on retract, # or speculative draft tokens). self._free_tail(slot, req, prefix_len) device_indices = self.req_to_token_pool.req_to_token[ req.req_pool_idx, :prefix_len ].to(dtype=torch.int64) return MatchResult( device_indices=device_indices, last_device_node=slot.virtual_node, last_host_node=slot.virtual_node, best_match_node=slot.virtual_node, cache_protected_len=slot.cache_protected_len, ) 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 False from sglang.srt.managers.schedule_batch import FINISH_ABORT session_id = req.session.session_id slot = self.slots.get(session_id) is_first = slot is None # Mid-processing abort only. Pre-aborted reqs have session=None # (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. if isinstance(req.finished_reason, FINISH_ABORT): if slot is None: # First-request mid-processing abort: create ephemeral # slot from req state so release_session handles cleanup. # Include last_node/cache_protected_len from the req so # release_session calls dec_lock_ref on the tree lock. # Also carry the mamba refs over so _free_slot_mamba can # return the (possibly extra_buffer ping-pong) slots to # the mamba pool; otherwise the abort orphans them. slot = SessionSlot( req_pool_idx=req.req_pool_idx, kv_allocated_len=req.kv_allocated_len, last_node=req.last_node, cache_protected_len=req.cache_protected_len, swa_uuid_for_lock=req.swa_uuid_for_lock, mamba_pool_idx=req.mamba_pool_idx, mamba_ping_pong_track_buffer=req.mamba_ping_pong_track_buffer, ) self.slots[session_id] = slot # Slot now owns the mamba state — drop the req's refs so # the abort fall-through doesn't double-free. req.mamba_pool_idx = None req.mamba_ping_pong_track_buffer = None slot.kv_allocated_len = max(slot.kv_allocated_len, req.kv_allocated_len) self.release_session(session_id) req.req_pool_idx = None req.session.abort_req() self._mark_kv_freed(req) return True if is_first: slot = SessionSlot() self.slots[session_id] = slot finished_len = ( req.finished_len if req.finished_len is not None else len(req.output_ids) ) target = len(req.origin_input_ids) + finished_len self._trim_overshoot(req, finished_len) slot.save_from_req(req, is_first=is_first) # Inherit the authoritative finished length on the slot, not the lagging # req clock (under overlap + honest committed the clock lags the in-flight # verify by ~1, which would short-change inheritance). Clamp to allocated # to keep committed <= allocated for prepare_for_decode. slot.kv_committed_len = min(target, slot.kv_allocated_len) # Update req_nodes to this successfully finished request. req.session.finish_req(req) self._mark_kv_freed(req) return True 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, : req.extend_range.end ] 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 # -- BasePrefixCache abstract methods: thin adapters over try_handle_* -- 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 self.inner.cache_finished_req(req, is_insert=is_insert, **kwargs) def cache_unfinished_req(self, req: Req, **kwargs): 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: 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: 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) -> None: slot = self.slots.pop(session_id, None) if slot is None: return protected_len = slot.cache_protected_len lock_node = slot.last_node tokens_freed = ( max(0, slot.kv_allocated_len - protected_len) if slot.is_holding_kv else 0 ) logger.info( "Session KV released: %s (%d tokens freed)", session_id, tokens_freed ) if lock_node is not None: if slot.swa_uuid_for_lock is not None: self.inner.dec_lock_ref( lock_node, DecLockRefParams(swa_uuid_for_lock=slot.swa_uuid_for_lock), ) else: self.inner.dec_lock_ref(lock_node) if slot.is_holding_kv: start = protected_len end = slot.kv_allocated_len if start < end: kv_indices = self.req_to_token_pool.req_to_token[ slot.req_pool_idx, start:end ] self.token_to_kv_pool_allocator.free(kv_indices) self.req_to_token_pool.free_slots.append(slot.req_pool_idx) self._free_slot_mamba(slot) def release_radix_session(self, session_id: str) -> None: self.inner.release_radix_session(session_id) 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 -- 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. """ total = 0 for slot in self.slots.values(): in_batch = ( active_pool_idxs is not None and slot.req_pool_idx in active_pool_idxs ) if slot.is_holding_kv and not in_batch: allocated = ceil_align(slot.kv_allocated_len, self.page_size) total += allocated - slot.cache_protected_len return total def session_held_full_tokens(self, active_pool_idxs: Optional[set] = None) -> int: """An alias to align the naming style of SWA""" return self.session_held_tokens(active_pool_idxs) def session_held_swa_tokens(self, active_pool_idxs: Optional[set] = None) -> int: """Total SWA tokens held by session slots, not tracked by the tree.""" total = 0 for slot in self.slots.values(): in_batch = ( active_pool_idxs is not None and slot.req_pool_idx in active_pool_idxs ) if slot.is_holding_kv and not in_batch: allocated = ceil_align(slot.kv_allocated_len, self.page_size) total += allocated - max( slot.cache_protected_len, slot.swa_evicted_seqlen ) return total def session_held_req_count(self, active_pool_idxs: Optional[set] = None) -> int: """Number of req pool slots held by session slots.""" def _owned(s): in_batch = ( active_pool_idxs is not None and s.req_pool_idx in active_pool_idxs ) return s.is_holding_kv and not in_batch return sum(_owned(s) for s in self.slots.values()) def session_held_mamba_slots(self, active_pool_idxs: Optional[set] = None) -> int: """Total mamba_pool entries held by session slots (mamba_pool_idx + mamba_ping_pong_track_buffer). Excludes slots whose owning req is currently in the batch -- those slots are counted via the normal alloc/free paths (same convention as the sibling ``session_held_*`` accessors). """ total = 0 for slot in self.slots.values(): in_batch = ( active_pool_idxs is not None and slot.req_pool_idx in active_pool_idxs ) if in_batch: continue if slot.mamba_pool_idx is not None: total += slot.mamba_pool_idx.numel() if slot.mamba_ping_pong_track_buffer is not None: total += slot.mamba_ping_pong_track_buffer.numel() return total def _free_slot_mamba(self, slot: SessionSlot) -> None: """Return a session slot's mamba pool state to the allocator.""" mamba_allocator = getattr(self.req_to_token_pool, "mamba_allocator", None) if mamba_allocator is None: return if slot.mamba_pool_idx is not None: mamba_allocator.free(slot.mamba_pool_idx.unsqueeze(0)) slot.mamba_pool_idx = None if slot.mamba_ping_pong_track_buffer is not None: mamba_allocator.free(slot.mamba_ping_pong_track_buffer) slot.mamba_ping_pong_track_buffer = None # -- 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): return self.inner.evictable_size() def full_evictable_size(self): return self.inner.full_evictable_size() def swa_evictable_size(self): return self.inner.swa_evictable_size() def protected_size(self): return self.inner.protected_size() def full_protected_size(self): return self.inner.full_protected_size() def swa_protected_size(self): return self.inner.swa_protected_size() def total_size(self): return self.inner.total_size() def pretty_print(self): return self.inner.pretty_print() def init_load_back(self, params: InitLoadBackParams): return self.inner.init_load_back(params) def ready_to_load_host_cache(self): return self.inner.ready_to_load_host_cache() def flush_write_through_acks(self) -> None: return self.inner.flush_write_through_acks() def check_hicache_events(self): return self.inner.check_hicache_events() def take_events(self): return self.inner.take_events() def supports_swa(self): return self.inner.supports_swa() 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() def is_tree_cache(self): return self.inner.is_tree_cache() def available_and_evictable_str(self): return self.inner.available_and_evictable_str() def init_metrics_collector(self): return self.inner.init_metrics_collector() def sanity_check(self): # Skip inner sanity check when sessions hold tree locks, because # the check asserts all nodes are unlocked during idle. if self.any_holding_kv(): return self.inner.sanity_check() # Forward attribute access for cache-specific methods (e.g. # sliding_window_size, all_values_flatten, etc.) def __getattr__(self, name): return getattr(self.inner, name)