chore: import upstream snapshot with attribution

python/sglang/srt/session/session_controller.py

@@ -0,0 +1,484 @@
+# Copyright 2023-2024 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+from __future__ import annotations
+
+import logging
+import time
+import uuid
+from array import array
+from typing import TYPE_CHECKING, Dict, Optional
+
+from sglang.srt.managers.io_struct import (
+    CloseSessionReqInput,
+    OpenSessionReqInput,
+    OpenSessionReqOutput,
+    TokenizedGenerateReqInput,
+)
+from sglang.srt.managers.schedule_batch import FINISH_ABORT, Req
+from sglang.srt.utils.common import log_info_on_rank0
+
+if TYPE_CHECKING:
+    from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache
+
+logger = logging.getLogger(__name__)
+
+
+class SessionReqNode:
+    def __init__(
+        self,
+        req: Req,
+        parent: Optional[SessionReqNode] = None,
+        children=None,
+    ):
+        self.req = req
+        self.parent = parent
+        if parent is not None:
+            parent.children.append(self)
+        self.children = [] if not children else children
+
+    def clear_children(self, req_dict):
+        for req_node in self.children:
+            req_node.clear(req_dict)
+        self.children = []
+
+    def clear(self, req_dict):
+        for req_node in self.children:
+            req_node.clear(req_dict)
+
+        if self.req.finished_reason is None:
+            self.req.to_finish = FINISH_ABORT()
+        del req_dict[self.req.rid]
+
+    def abort(self):
+        if self.req.finished_reason is None:
+            self.req.to_finish = FINISH_ABORT()
+
+    def __str__(self):
+        return self._str_helper(self.req.rid)
+
+    def _str_helper(self, prefix=""):
+        if len(self.children) == 0:
+            return prefix + "\n"
+        else:
+            origin_prefix = prefix
+            prefix += " -- " + self.children[0].req.rid
+            ret = self.children[0]._str_helper(prefix)
+            for child in self.children[1:]:
+                prefix = " " * len(origin_prefix) + " \\- " + child.req.rid
+                ret += child._str_helper(prefix)
+            return ret
+
+
+class Session:
+    def __init__(
+        self,
+        capacity_of_str_len: int,
+        session_id: Optional[str] = None,
+        streaming: bool = False,
+        timeout: Optional[float] = None,
+    ):
+        self.session_id = session_id if session_id is not None else uuid.uuid4().hex
+        self.capacity_of_str_len = capacity_of_str_len
+        self.streaming = streaming
+        self.timeout = timeout
+        self.last_active_time: float = time.monotonic()
+        self.req_nodes: Dict[str, SessionReqNode] = {}
+        self.close_on_finish: bool = False
+        self._inflight: bool = False
+        # Token-array lengths of last_req as of its finish_req. The share path
+        # appends speculatively beyond these; only finish_req confirms them, so
+        # _share_token_arrays trims back first (heals aborted turns).
+        self.committed_origin_len: Optional[int] = None
+        self.committed_unpadded_len: Optional[int] = None
+        self.committed_fill_len: Optional[int] = None
+
+    def is_timed_out(self) -> bool:
+        if self.timeout is None:
+            return False
+        return time.monotonic() - self.last_active_time > self.timeout
+
+    @staticmethod
+    def _strip_bos_token(req: TokenizedGenerateReqInput, tokenizer) -> None:
+        """Trim a leading BOS on an appended turn; shift mm offsets to match."""
+        if not (
+            tokenizer is not None
+            and req.input_ids
+            and req.input_ids[0] == tokenizer.bos_token_id
+        ):
+            return
+        req.input_ids = req.input_ids[1:]
+        if req.mm_inputs:
+            for item in req.mm_inputs.mm_items:
+                if item.offsets:
+                    if any(s == 0 for s, _ in item.offsets):
+                        logging.warning(
+                            "mm_item offset starts at 0 (BOS position), "
+                            "clamping to 0 after BOS strip"
+                        )
+                    item.offsets = [
+                        (max(0, s - 1), max(0, e - 1)) for s, e in item.offsets
+                    ]
+
+    def _share_token_arrays(self, last_req: Req, new_input_ids):
+        """Plain streaming append: reuse last_req's token arrays in place.
+
+        Trims each array back to its committed length first — an earlier turn
+        may have appended its tokens and then aborted before finish_req, and
+        req_nodes still points at last_req, so anything beyond the committed
+        lengths is unconfirmed. Then extends with last turn's output and the
+        new input. Returns (input_ids, input_ids_unpadded, carry_fill);
+        carry_fill (== the new origin) spares the first fill_ids rebuild.
+        """
+        out_tail = last_req.output_ids[: last_req.sampling_params.max_new_tokens]
+
+        input_ids = last_req.origin_input_ids
+        del input_ids[self.committed_origin_len :]
+        if last_req.origin_input_ids_unpadded is input_ids:
+            input_ids_unpadded = input_ids
+        else:
+            input_ids_unpadded = last_req.origin_input_ids_unpadded
+            del input_ids_unpadded[self.committed_unpadded_len :]
+
+        carry_fill = last_req.full_untruncated_fill_ids
+        if (
+            not isinstance(carry_fill, array)
+            or carry_fill is input_ids
+            or carry_fill is input_ids_unpadded
+        ):
+            # Unexpected type or aliased with an origin array (extending it
+            # below would double-append): let _refresh_fill_ids rebuild.
+            carry_fill = None
+        else:
+            del carry_fill[self.committed_fill_len :]
+            baked = len(carry_fill) - len(input_ids)
+            if 0 <= baked <= len(out_tail):
+                carry_fill.extend(out_tail[baked:])
+                carry_fill.extend(new_input_ids)
+            else:
+                carry_fill = None
+
+        input_ids.extend(out_tail)
+        input_ids.extend(new_input_ids)
+        if input_ids_unpadded is not input_ids:
+            input_ids_unpadded.extend(out_tail)
+            input_ids_unpadded.extend(new_input_ids)
+        return input_ids, input_ids_unpadded, carry_fill
+
+    @staticmethod
+    def _concat_token_arrays(
+        last_req: Req, req: TokenizedGenerateReqInput, session_params
+    ):
+        """Copy-based assembly for replace/offset/drop_previous_output turns."""
+        out_tail = last_req.output_ids[: last_req.sampling_params.max_new_tokens]
+
+        input_ids = last_req.origin_input_ids + out_tail
+        if session_params.drop_previous_output:
+            input_ids = last_req.origin_input_ids[:]
+        if session_params.offset and session_params.offset != 0:
+            input_ids = input_ids[: session_params.offset] + req.input_ids
+        else:
+            input_ids += req.input_ids
+
+        input_ids_unpadded = last_req.origin_input_ids_unpadded + out_tail
+        if session_params.drop_previous_output:
+            input_ids_unpadded = last_req.origin_input_ids_unpadded[:]
+        if session_params.offset and session_params.offset != 0:
+            input_ids_unpadded = (
+                input_ids_unpadded[: session_params.offset] + req.input_ids
+            )
+        else:
+            input_ids_unpadded += req.input_ids
+        return input_ids, input_ids_unpadded
+
+    def create_req(
+        self,
+        req: TokenizedGenerateReqInput,
+        tokenizer,
+        vocab_size: int,
+        eos_token_ids=None,
+    ):
+        assert req.session_params is not None
+        self.last_active_time = time.monotonic()
+        session_params = req.session_params
+
+        last_req_node = None
+        last_req = None
+        abort = False
+        abort_message = ""
+        if self.streaming:
+            # Streaming sessions: only simple appends allowed; reject otherwise.
+            if self._inflight:
+                abort = True
+                abort_message = "Streaming session already has an active request."
+            elif session_params.replace:
+                abort = True
+                abort_message = "Streaming sessions do not support replace."
+            elif session_params.drop_previous_output:
+                abort = True
+                abort_message = (
+                    "Streaming sessions do not support drop_previous_output."
+                )
+            elif session_params.offset and session_params.offset != 0:
+                abort = True
+                abort_message = "Streaming sessions do not support offset."
+            elif self.req_nodes:
+                assert len(self.req_nodes) == 1
+                # Peek (don't pop) the single req_node. req_nodes is updated
+                # only in finish_req after the request completes successfully.
+                [last_req_node] = self.req_nodes.values()
+                last_req = last_req_node.req
+        elif session_params.replace:
+            if session_params.rid is None:
+                for _, req_node in self.req_nodes.items():
+                    req_node.clear(self.req_nodes)
+            else:
+                if session_params.rid not in self.req_nodes:
+                    abort = True
+                    abort_message = "Invalid request session id"
+                else:
+                    last_req_node = self.req_nodes[session_params.rid]
+                    last_req_node.abort()
+                    last_req = last_req_node.req
+                    last_req_node.clear_children(self.req_nodes)
+        else:
+            if session_params.rid is not None:
+                if session_params.rid not in self.req_nodes:
+                    abort = True
+                    abort_message = "Invalid request session id"
+                else:
+                    last_req_node = self.req_nodes[session_params.rid]
+                    last_req = last_req_node.req
+                    if not last_req.finished():
+                        abort = True
+                        abort_message = "Session request is appending to a request that hasn't finished."
+                        logging.warning(abort_message)
+
+        carry_fill = None
+        if last_req is not None:
+            self._strip_bos_token(req, tokenizer)
+            # In-place sharing is only safe for the plain streaming append:
+            # streaming sessions allow a single inflight request, last_req has
+            # finished, and the committed_* lengths recorded by finish_req let
+            # _share_token_arrays trim away tokens appended by an aborted turn.
+            # offset / drop_previous_output rewrite history and must copy.
+            can_share_token_arrays = (
+                self.streaming
+                and self.committed_origin_len is not None
+                and not session_params.drop_previous_output
+                and not (session_params.offset and session_params.offset != 0)
+            )
+            if can_share_token_arrays:
+                input_ids, input_ids_unpadded, carry_fill = self._share_token_arrays(
+                    last_req, req.input_ids
+                )
+            else:
+                input_ids, input_ids_unpadded = self._concat_token_arrays(
+                    last_req, req, session_params
+                )
+        else:
+            input_ids = req.input_ids
+            input_ids_unpadded = req.input_ids
+
+        new_req = Req(
+            rid=req.rid,
+            origin_input_text=None,
+            origin_input_ids=input_ids,
+            origin_input_ids_unpadded=input_ids_unpadded,
+            sampling_params=req.sampling_params,
+            lora_id=req.lora_id,
+            session=self,
+            custom_logit_processor=req.custom_logit_processor,
+            stream=req.stream,
+            return_logprob=req.return_logprob,
+            top_logprobs_num=req.top_logprobs_num,
+            token_ids_logprob=req.token_ids_logprob,
+            vocab_size=vocab_size,
+            eos_token_ids=eos_token_ids,
+            require_reasoning=req.require_reasoning,
+            return_hidden_states=req.return_hidden_states,
+            return_routed_experts=req.return_routed_experts,
+            routed_experts_start_len=req.routed_experts_start_len,
+            priority=req.priority,
+            routing_key=req.routing_key,
+            extra_key=req.extra_key,
+            http_worker_ipc=req.http_worker_ipc,
+            time_stats=req.time_stats,
+        )
+        if last_req is not None:
+            new_req.multimodal_inputs = last_req.multimodal_inputs
+        new_req.tokenizer = tokenizer
+        if carry_fill is not None:
+            new_req.full_untruncated_fill_ids = carry_fill
+
+        if abort:
+            new_req.set_finish_with_abort(abort_message)
+        elif self.streaming:
+            # req_nodes is NOT updated here — finish_req() handles it.
+            self._inflight = True
+        else:
+            new_req_node = SessionReqNode(new_req, last_req_node)
+            self.req_nodes[req.rid] = new_req_node
+
+        return new_req
+
+    def finish_req(self, req):
+        """Update req_nodes after a streaming request finishes successfully."""
+        self._inflight = False
+        if self.req_nodes:
+            [prev_node] = self.req_nodes.values()
+            prev_node.req.session = None
+            self.req_nodes.clear()
+        self.req_nodes[req.rid] = SessionReqNode(req)
+        # Confirm this req's token arrays as the session's rollback point.
+        self.committed_origin_len = len(req.origin_input_ids)
+        self.committed_unpadded_len = len(req.origin_input_ids_unpadded)
+        self.committed_fill_len = len(req.full_untruncated_fill_ids)
+
+    def abort_req(self):
+        """Clear inflight flag on abort (req_nodes stays unchanged)."""
+        self._inflight = False
+
+
+class SessionController:
+    def __init__(self, tree_cache: BasePrefixCache):
+        self.sessions: Dict[str, Session] = {}
+        self._last_reap_time: float = 0.0
+        self.tree_cache = tree_cache
+
+    def __contains__(self, session_id: str) -> bool:
+        return session_id in self.sessions
+
+    def get(self, session_id: str) -> Optional[Session]:
+        return self.sessions.get(session_id)
+
+    def open(self, recv_req: OpenSessionReqInput) -> OpenSessionReqOutput:
+        session_id = recv_req.session_id
+        if session_id in self.sessions:
+            logger.warning(f"session id {session_id} already exist, cannot open.")
+            return OpenSessionReqOutput(session_id=session_id, success=False)
+        elif session_id is None:
+            logger.warning("session id is None, cannot open.")
+            return OpenSessionReqOutput(session_id=session_id, success=False)
+        else:
+            self.sessions[session_id] = Session(
+                recv_req.capacity_of_str_len,
+                session_id,
+                streaming=bool(recv_req.streaming),
+                timeout=recv_req.timeout,
+            )
+            log_info_on_rank0(
+                logger, f"Session opened: {session_id} (active={len(self.sessions)})"
+            )
+            return OpenSessionReqOutput(session_id=session_id, success=True)
+
+    def close(self, recv_req: CloseSessionReqInput):
+        session_id = recv_req.session_id
+        if session_id not in self.sessions:
+            logger.warning(f"session id {session_id} does not exist, cannot delete.")
+        else:
+            self._close(session_id)
+
+    def _close(self, session_id: str):
+        session = self.sessions[session_id]
+        req = None
+        has_unfinished_request = False
+        if session.streaming and session._inflight:
+            has_unfinished_request = True
+        elif session.streaming and session.req_nodes:
+            assert len(session.req_nodes) == 1
+            [last_node] = session.req_nodes.values()
+            req = last_node.req
+            if not req.finished():
+                has_unfinished_request = True
+
+        if has_unfinished_request:
+            # An in-flight request is still decoding on this session's KV
+            # memory. Freeing now would corrupt the scheduler. Mark the
+            # session for deferred cleanup: the request keeps its session
+            # reference so cache_finished_req takes the streaming path,
+            # and we schedule release_session for after it completes.
+            session.close_on_finish = True
+            logger.info(
+                "Deferring session close for %s (unfinished request)",
+                session_id,
+            )
+            return
+
+        # No owning request -- safe to release immediately.
+        if session.streaming and session.req_nodes:
+            req = next(iter(session.req_nodes.values())).req
+            req.session = None
+
+        # Release multimodal features held by session requests.
+        # Session reqs skip the normal mm cleanup path (scheduler and
+        # output_processor) so features stay alive until the session closes.
+        seen_mm = set()
+        for node in session.req_nodes.values():
+            mm = node.req.multimodal_inputs
+            if mm is not None and id(mm) not in seen_mm:
+                seen_mm.add(id(mm))
+                mm.release_features()
+            node.req.multimodal_inputs = None
+
+        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)})"
+        )
+
+    def maybe_reap(self, now: float, interval: float = 1.0):
+        # reap sessions every second
+        if now - self._last_reap_time > interval:
+            self._last_reap_time = now
+
+            # Finish deferred closes for sessions whose requests completed.
+            pending = [
+                sid
+                for sid, session in self.sessions.items()
+                if session.close_on_finish and self._all_requests_finished(session)
+            ]
+            for sid in pending:
+                log_info_on_rank0(
+                    logger, f"Deferred close ready for session {sid}, releasing."
+                )
+                # Reset close_on_finish so _close proceeds with the release.
+                self.sessions[sid].close_on_finish = False
+                self._close(sid)
+
+            timed_out = [
+                sid for sid, session in self.sessions.items() if session.is_timed_out()
+            ]
+            for sid in timed_out:
+                log_info_on_rank0(logger, f"Session {sid} timed out, closing.")
+                self._close(sid)
+
+    @staticmethod
+    def _all_requests_finished(session: Session) -> bool:
+        if not session.req_nodes:
+            return True
+        return all(node.req.finished() for node in session.req_nodes.values())
+
+    @staticmethod
+    def adjust_mm_offsets(recv_req: TokenizedGenerateReqInput, req: Req, image_inputs):
+        # For session requests, adjust mm_inputs offsets by the prefix length.
+        # Session.create_req prepends previous context to origin_input_ids,
+        # so offsets from the new prompt need to be shifted.
+        if len(recv_req.input_ids) >= len(req.origin_input_ids):
+            return
+        prefix_len = len(req.origin_input_ids) - len(recv_req.input_ids)
+        for mm_item in image_inputs.mm_items:
+            if mm_item.offsets:
+                mm_item.offsets = [
+                    (start + prefix_len, end + prefix_len)
+                    for start, end in mm_item.offsets
+                ]

python/sglang/srt/session/streaming_session.py

@@ -0,0 +1,647 @@
+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)