[Session] Add streaming mode with SessionAwareCache fast path (#19171)

Co-authored-by: hnyls2002 <lsyincs@gmail.com>

python/sglang/srt/entrypoints/engine.py

@@ -492,12 +492,18 @@ class Engine(EngineBase):
         self,
         capacity_of_str_len: int,
         session_id: Optional[str] = None,
+        streaming: bool = False,
+        timeout: Optional[float] = None,
     ) -> str:
         """Open a session for multi-turn conversation with shared context.
 
         Args:
             capacity_of_str_len: Maximum string length capacity for the session.
             session_id: Optional session ID. If not provided, a UUID will be generated.
+            streaming: Use low-overhead path for realtime streaming (append-only mode).
+            timeout: If set, the session is automatically closed after being inactive
+                for this many seconds. Inactivity is measured from session open or the
+                most recent request submission.
 
         Returns:
             The session ID (either the provided one or a newly generated UUID).
@@ -505,6 +511,8 @@ class Engine(EngineBase):
         obj = OpenSessionReqInput(
             capacity_of_str_len=capacity_of_str_len,
             session_id=session_id,
+            streaming=streaming,
+            timeout=timeout,
         )
         return self.loop.run_until_complete(
             self.tokenizer_manager.open_session(obj, None)

python/sglang/srt/managers/io_struct.py

@@ -1623,6 +1623,8 @@ class ConfigureLoggingReq(BaseReq):
 class OpenSessionReqInput(BaseReq):
     capacity_of_str_len: int
     session_id: Optional[str] = None
+    streaming: Optional[bool] = None
+    timeout: Optional[float] = None
 
 
 @dataclass

python/sglang/srt/managers/schedule_batch.py

@@ -93,6 +93,7 @@ if TYPE_CHECKING:
     from typing import Any, Dict
 
     from sglang.srt.configs.model_config import ModelConfig
+    from sglang.srt.managers.session_controller import Session
     from sglang.srt.observability.scheduler_metrics_mixin import PrefillStats
     from sglang.srt.speculative.eagle_info import EagleDraftInput
     from sglang.srt.speculative.spec_info import SpecInput, SpeculativeAlgorithm
@@ -499,7 +500,7 @@ class Req(ReqDllmMixin):
         lora_id: Optional[str] = None,
         input_embeds: Optional[List[List[float]]] = None,
         token_type_ids: List[int] = None,
-        session_id: Optional[str] = None,
+        session: Optional[Session] = None,
         custom_logit_processor: Optional[str] = None,
         require_reasoning: bool = False,
         return_hidden_states: bool = False,
@@ -535,7 +536,7 @@ class Req(ReqDllmMixin):
         self.output_ids = []
         # fill_ids = origin_input_ids + output_ids. Updated if chunked.
         self.fill_ids = []
-        self.session_id = session_id
+        self.session = session
         self.input_embeds = input_embeds
 
         # For req-level memory management
@@ -874,7 +875,7 @@ class Req(ReqDllmMixin):
             match_result = tree_cache.match_prefix(
                 MatchPrefixParams(
                     key=RadixKey(token_ids=token_ids, extra_key=self.extra_key),
-                    req=self if tree_cache.supports_mamba() else None,
+                    req=self,
                     cow_mamba=tree_cache.supports_mamba(),
                 )
             )

python/sglang/srt/managers/scheduler.py

@@ -721,6 +721,10 @@ class Scheduler(
             else:
                 self.tree_cache = RadixCache(params)
 
+        from sglang.srt.mem_cache.session_aware_cache import SessionAwareCache
+
+        self.tree_cache = SessionAwareCache(self.tree_cache)
+
         if (
             server_args.disaggregation_mode == "decode"
             and server_args.disaggregation_decode_enable_offload_kvcache
@@ -751,6 +755,7 @@ class Scheduler(
         self.num_retracted_reqs: int = 0
         self.num_paused_reqs: int = 0
         self.sessions: Dict[str, Session] = {}
+        self._last_reap_sessions: float = 0.0
         self.forward_sleep_time = None
         self._engine_paused = False
 
@@ -1118,7 +1123,8 @@ class Scheduler(
                 result = self.run_batch(batch)
                 self.process_batch_result(batch, result)
             else:
-                # When the server is idle, do self-check and re-init some states
+                # When the server is idle, do self-check and re-init some states.
+                # Skip if there are any streaming sessions (latency sensitive).
                 self.self_check_during_idle()
 
             # Update last_batch
@@ -1349,7 +1355,10 @@ class Scheduler(
         return work_reqs, control_reqs
 
     def process_input_requests(self, recv_reqs: List):
-
+        now = time.monotonic()
+        if now - self._last_reap_sessions > 1.0:  # reap sessions every second
+            self._last_reap_sessions = now
+            self.reap_timed_out_sessions()
         for recv_req in recv_reqs:
             # If it is a health check generation request and there are running requests, ignore it.
             if is_health_check_generate_req(recv_req) and (
@@ -1458,7 +1467,7 @@ class Scheduler(
             if not req.finished() or not (mm_inputs := req.multimodal_inputs):
                 continue
             # For session requests, keep mm_inputs for the next request
-            if req.session_id:
+            if req.session:
                 continue
             # For non-session requests, clear features and mm_inputs
             for item in mm_inputs.mm_items:
@@ -2929,17 +2938,37 @@ class Scheduler(
             return OpenSessionReqOutput(session_id, False)
         else:
             self.sessions[session_id] = Session(
-                recv_req.capacity_of_str_len, session_id
+                recv_req.capacity_of_str_len,
+                session_id,
+                streaming=bool(recv_req.streaming),
+                timeout=recv_req.timeout,
             )
             return OpenSessionReqOutput(session_id, True)
 
     def close_session(self, recv_req: CloseSessionReqInput):
-        # handle error
         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:
-            del self.sessions[session_id]
+            self._close_session(session_id)
+
+    def _close_session(self, session_id: str):
+        session = self.sessions[session_id]
+        if session.streaming and session.req_nodes:
+            assert len(session.req_nodes) == 1
+            req = next(iter(session.req_nodes.values())).req
+            if not req.finished():
+                req.session = None
+        self.tree_cache.release_session(session_id)
+        del self.sessions[session_id]
+
+    def reap_timed_out_sessions(self):
+        timed_out = [
+            sid for sid, session in self.sessions.items() if session.is_timed_out()
+        ]
+        for sid in timed_out:
+            logger.info(f"Session {sid} timed out, closing.")
+            self._close_session(sid)
 
     def maybe_sleep_on_idle(self):
         if self.idle_sleeper is not None:

python/sglang/srt/managers/scheduler_output_processor_mixin.py

@@ -175,7 +175,6 @@ class SchedulerOutputProcessorMixin:
                         release_kv_cache(req, self.tree_cache)
                         req.time_stats.set_completion_time()
                     elif not batch.decoding_reqs or req not in batch.decoding_reqs:
-                        # This updates radix so others can match
                         self.tree_cache.cache_unfinished_req(req)
 
                     self.maybe_collect_customized_info(i, req, logits_output)

python/sglang/srt/managers/scheduler_runtime_checker_mixin.py

@@ -92,10 +92,11 @@ class SchedulerRuntimeCheckerMixin:
             swa_available_size,
             swa_evictable_size,
         ) = self._get_swa_token_info()
-        memory_leak = full_num_used != 0 or swa_num_used != 0
+        session_held = self.tree_cache.session_held_tokens()
+        memory_leak = (full_num_used - session_held) != 0 or swa_num_used != 0
         token_msg = (
             f"{self.full_tokens_per_layer=}, {full_available_size=}, {full_evictable_size=}, {self.tree_cache.full_protected_size()=}\n"
-            f"{self.swa_tokens_per_layer=}, {swa_available_size=}, {swa_evictable_size=}, {self.tree_cache.swa_protected_size()=}\n"
+            f"{self.swa_tokens_per_layer=}, {swa_available_size=}, {swa_evictable_size=}, {self.tree_cache.swa_protected_size()=}, {session_held=}\n"
         )
         return memory_leak, token_msg
 
@@ -110,8 +111,9 @@ class SchedulerRuntimeCheckerMixin:
             mamba_available_size,
             mamba_evictable_size,
         ) = self._get_mamba_token_info()
+        session_held = self.tree_cache.session_held_tokens()
         memory_leak = (
-            full_num_used != self.tree_cache.full_protected_size()
+            full_num_used != self.tree_cache.full_protected_size() + session_held
             or mamba_num_used != self.tree_cache.mamba_protected_size()
         )
         if memory_leak:
@@ -150,14 +152,11 @@ class SchedulerRuntimeCheckerMixin:
     def _check_radix_cache_memory(self: Scheduler):
         _, _, available_size, evictable_size = self._get_token_info()
         protected_size = self.tree_cache.protected_size()
+        session_held = self.tree_cache.session_held_tokens()
         memory_leak = (available_size + evictable_size) != (
-            # self.max_total_num_tokens
+            self.max_total_num_tokens - protected_size - session_held
-            # if not self.enable_hierarchical_cache
-            # else self.max_total_num_tokens - protected_size
-            self.max_total_num_tokens
-            - protected_size
         )
-        token_msg = f"{self.max_total_num_tokens=}, {available_size=}, {evictable_size=}, {protected_size=}\n"
+        token_msg = f"{self.max_total_num_tokens=}, {available_size=}, {evictable_size=}, {protected_size=}, {session_held=}\n"
         return memory_leak, token_msg
 
     def _get_batch_uncached_size(self: Scheduler, batch: ScheduleBatch) -> int:
@@ -205,7 +204,14 @@ class SchedulerRuntimeCheckerMixin:
             log_msg = f"[Mem Check (BUSY)] {available_size=}, {evictable_size=}, {protected_size=}, {uncached_size=}"
             logger.info(log_msg)
 
-        total_tokens = available_size + evictable_size + protected_size + uncached_size
+        session_held = self.tree_cache.session_held_tokens()
+        total_tokens = (
+            available_size
+            + evictable_size
+            + protected_size
+            + uncached_size
+            + session_held
+        )
         assert (
             total_tokens == self.max_total_num_tokens
         ), f"Mem Leak Detected! {total_tokens=} vs {self.max_total_num_tokens=}"
@@ -218,10 +224,12 @@ class SchedulerRuntimeCheckerMixin:
         else:
             req_total_size = self.req_to_token_pool.size
 
-        if len(self.req_to_token_pool.free_slots) != req_total_size:
+        session_req_count = self.tree_cache.session_held_req_count()
+        if len(self.req_to_token_pool.free_slots) + session_req_count != req_total_size:
             msg = (
                 "req_to_token_pool memory leak detected!"
                 f"available_size={len(self.req_to_token_pool.free_slots)}, "
+                f"session_held={session_req_count}, "
                 f"total_size={self.req_to_token_pool.size}\n"
             )
             raise_error_or_warn(

python/sglang/srt/managers/session_controller.py

@@ -11,6 +11,7 @@
 # ==============================================================================
 
 import logging
+import time
 import uuid
 from typing import Dict, Optional
 
@@ -65,25 +66,59 @@ class SessionReqNode:
 
 
 class Session:
-    def __init__(self, capacity_of_str_len: int, session_id: Optional[str] = None):
+    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] = {}
 
+    def is_timed_out(self) -> bool:
+        if self.timeout is None:
+            return False
+        return time.monotonic() - self.last_active_time > self.timeout
+
     def create_req(self, req: TokenizedGenerateReqInput, tokenizer, vocab_size: int):
         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
-        if session_params.replace:
+        abort_message = ""
+        if self.streaming:
+            # Streaming sessions: only simple appends allowed; reject otherwise.
+            if 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
+                _, last_req_node = self.req_nodes.popitem()
+                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()
@@ -93,18 +128,22 @@ class Session:
             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():
-                        logging.warning(
-                            "The request in a session is appending to a request that hasn't finished."
-                        )
                         abort = True
+                        abort_message = "Session request is appending to a request that hasn't finished."
+                        logging.warning(abort_message)
 
         if last_req is not None:
             # trim bos token if it is an append
-            if tokenizer is not None and req.input_ids[0] == tokenizer.bos_token_id:
+            if (
+                tokenizer is not None
+                and req.input_ids
+                and req.input_ids[0] == tokenizer.bos_token_id
+            ):
                 req.input_ids = req.input_ids[1:]
 
             input_ids = (
@@ -136,6 +175,7 @@ class Session:
         else:
             input_ids = req.input_ids
             input_ids_unpadded = req.input_ids
+
         new_req = Req(
             rid=req.rid,
             origin_input_text=None,
@@ -143,7 +183,7 @@ class Session:
             origin_input_ids_unpadded=input_ids_unpadded,
             sampling_params=req.sampling_params,
             lora_id=req.lora_id,
-            session_id=self.session_id,
+            session=self,
             custom_logit_processor=req.custom_logit_processor,
             stream=req.stream,
             return_logprob=req.return_logprob,
@@ -156,7 +196,11 @@ class Session:
         new_req.tokenizer = tokenizer
 
         if abort:
-            new_req.set_finish_with_abort("Invalid request session id")
+            new_req.set_finish_with_abort(abort_message)
+        elif self.streaming:
+            if last_req is not None:
+                last_req.session = None
+            self.req_nodes[req.rid] = SessionReqNode(new_req)
         else:
             new_req_node = SessionReqNode(new_req, last_req_node)
             self.req_nodes[req.rid] = new_req_node

python/sglang/srt/mem_cache/common.py

@@ -478,6 +478,13 @@ def release_kv_cache(req: Req, tree_cache: BasePrefixCache, is_insert: bool = Tr
 
     tree_cache.cache_finished_req(req, is_insert=is_insert)
 
+    # FIXME: SessionAwareCache.cache_finished_req sets req_pool_idx = None to
+    # transfer KV ownership to the SessionSlot, so we skip the remaining
+    # cleanup (overalloc free + pool slot free). This means over-allocated
+    # tokens from speculative decoding are NOT freed between turns.
+    if req.req_pool_idx is None:
+        return
+
     start_p, end_p = req.pop_overallocated_kv_cache()
 
     global_server_args = get_global_server_args()

python/sglang/srt/mem_cache/memory_pool.py

@@ -158,7 +158,7 @@ class ReqToTokenPool:
         reusing = [i for i, r in enumerate(reqs) if r.req_pool_idx is not None]
         if not any(r.is_dllm() for r in reqs):
             assert (
-                len(reusing) <= 1
+                sum(1 for i in reusing if reqs[i].is_chunked > 0) <= 1
             ), "only one chunked request may reuse req_pool_idx in a batch"
         assert all(
             reqs[i].is_chunked > 0 or reqs[i].kv_committed_len > 0 for i in reusing

python/sglang/srt/mem_cache/session_aware_cache.py

@@ -0,0 +1,311 @@
+from __future__ import annotations
+
+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,
+    EvictParams,
+    EvictResult,
+    MatchPrefixParams,
+    MatchResult,
+)
+
+if TYPE_CHECKING:
+    from sglang.srt.managers.schedule_batch import Req
+
+
+class _VirtualNode:
+    """Sentinel node for streaming session requests.
+
+    Passed to inc_lock_ref / dec_lock_ref so the wrapper 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
+
+    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
+
+        req.req_pool_idx = None
+        req.mamba_pool_idx = 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
+
+        self.req_pool_idx = None
+        self.mamba_pool_idx = None
+
+
+def _is_streaming(req: Optional[Req]) -> bool:
+    return req is not None and req.session is not None and req.session.streaming
+
+
+class SessionAwareCache(BasePrefixCache):
+    """Decorator around any BasePrefixCache that manages streaming session KV.
+
+    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.
+    """
+
+    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
+
+    # -- BasePrefixCache abstract methods --
+
+    def reset(self):
+        self.slots.clear()
+        self.inner.reset()
+
+    def match_prefix(self, params: MatchPrefixParams) -> MatchResult:
+        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)
+
+        slot.restore_to_req(req)
+
+        max_prefix_len = len(params.key.token_ids)
+        prefix_len = min(req.kv_committed_len, max_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,
+        )
+
+    def cache_finished_req(self, req: Req, is_insert: bool = True, **kwargs):
+        if not _is_streaming(req):
+            return self.inner.cache_finished_req(req, is_insert=is_insert, **kwargs)
+
+        session_id = req.session.session_id
+        slot = self.slots.get(session_id)
+        is_first = slot is None
+        if is_first:
+            slot = SessionSlot()
+            self.slots[session_id] = slot
+
+        slot.save_from_req(req, is_first=is_first)
+
+    def cache_unfinished_req(self, req: Req, **kwargs):
+        if _is_streaming(req) and req.session.session_id in self.slots:
+            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):
+        if isinstance(node, _VirtualNode):
+            return None
+        return self.inner.inc_lock_ref(node)
+
+    def dec_lock_ref(self, node: Any, swa_uuid_for_lock: Optional[str] = None):
+        if isinstance(node, _VirtualNode):
+            return
+        if swa_uuid_for_lock is not None:
+            return self.inner.dec_lock_ref(node, swa_uuid_for_lock)
+        return self.inner.dec_lock_ref(node)
+
+    # -- Session lifecycle --
+
+    def release_session(self, session_id: str):
+        """Release all KV resources held by a streaming session."""
+        slot = self.slots.pop(session_id, None)
+        if slot is None:
+            return
+
+        if slot.last_node is not None:
+            if slot.swa_uuid_for_lock is not None:
+                self.inner.dec_lock_ref(slot.last_node, slot.swa_uuid_for_lock)
+            else:
+                self.inner.dec_lock_ref(slot.last_node)
+
+        if slot.req_pool_idx is not None:
+            start = slot.cache_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)
+
+    def session_held_tokens(self) -> int:
+        """Total KV tokens held by session slots, not tracked by the tree."""
+        total = 0
+        for slot in self.slots.values():
+            if slot.req_pool_idx is not None:
+                total += slot.kv_allocated_len - slot.cache_protected_len
+        return total
+
+    def session_held_req_count(self) -> int:
+        """Number of req pool slots held by session slots."""
+        return sum(1 for s in self.slots.values() if s.req_pool_idx is not None)
+
+    # -- 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, last_host_node, host_hit_length):
+        return self.inner.init_load_back(last_host_node, host_hit_length)
+
+    def ready_to_load_host_cache(self):
+        return self.inner.ready_to_load_host_cache()
+
+    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 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()
+
+    # Forward attribute access for cache-specific methods (e.g. sanity_check,
+    # sliding_window_size, all_values_flatten, etc.)
+    def __getattr__(self, name):
+        return getattr(self.inner, name)

test/registered/sessions/test_session_control.py

@@ -3,11 +3,13 @@ Usage:
 python3 -m unittest test_session_control.TestSessionControl.test_session_control
 python3 -m unittest test_session_control.TestSessionControl.test_session_control_with_branching
 python3 -m unittest test_session_control.TestSessionControl.test_session_control_backtrack_with_abort
+python3 -m unittest test_session_control.TestSessionControl.test_streaming_session
 python3 -m unittest test_session_control.TestSessionControlVision.test_session_control
 """
 
 import asyncio
 import json
+import time
 import unittest
 
 import aiohttp
@@ -15,6 +17,7 @@ import requests
 
 from sglang.srt.utils import kill_process_tree
 from sglang.srt.utils.hf_transformers_utils import get_tokenizer
+from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.test_utils import (
     DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
     DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
@@ -23,6 +26,8 @@ from sglang.test.test_utils import (
     popen_launch_server,
 )
 
+register_cuda_ci(est_time=60, suite="stage-b-test-large-1-gpu")
+
 
 def remove_prefix(text: str, prefix: str) -> str:
     return text[len(prefix) :] if text.startswith(prefix) else text
@@ -429,6 +434,162 @@ class TestSessionControl(unittest.TestCase):
         asyncio.run(self.run_session_control_backtrack_with_abort(replace=True))
         asyncio.run(self.run_session_control_backtrack_with_abort(replace=False))
 
+    def test_streaming_session(self, gen_len=12):
+        chunks = [
+            "Let me tell you something about France.",
+            "The capital of France is",
+            "The population of the city is",
+        ]
+        tokenizer = get_tokenizer(self.model)
+        chunks_ids = [tokenizer.encode(x) for x in chunks]
+        for i in range(1, len(chunks_ids)):
+            if chunks_ids[i][0] == tokenizer.bos_token_id:
+                chunks_ids[i] = chunks_ids[i][1:]
+
+        # === Part 1: streaming session ===
+        requests.post(self.base_url + "/flush_cache")
+        session_id = requests.post(
+            self.base_url + "/open_session",
+            json={"capacity_of_str_len": 1000, "streaming": True},
+        ).json()
+        rid = None
+        outputs_from_session = []
+
+        prev_kv_len = 0
+        for turn_idx, chunk_ids in enumerate(chunks_ids):
+            response = requests.post(
+                self.base_url + "/generate",
+                json={
+                    "input_ids": chunk_ids,
+                    "session_params": {"id": session_id, "rid": rid},
+                    "sampling_params": {
+                        "temperature": 0,
+                        "max_new_tokens": gen_len,
+                        "no_stop_trim": True,
+                        "skip_special_tokens": False,
+                    },
+                },
+            ).json()
+            rid = response["meta_info"]["id"]
+            outputs_from_session.append(response["text"])
+            cached = response["meta_info"]["cached_tokens"]
+            prompt_tokens = response["meta_info"]["prompt_tokens"]
+            completion_tokens = response["meta_info"]["completion_tokens"]
+
+            if turn_idx == 0:
+                # Turn 1 should have no cache hit (cache was flushed).
+                self.assertEqual(
+                    cached, 0, "Turn 1 should have 0 cached tokens (clean start)"
+                )
+            else:
+                # Turns 2+ inherit KV from the previous turn (via inherit_kv_states,
+                # not radix tree matching). cached_tokens reflects the inherited prefix.
+                self.assertEqual(
+                    cached,
+                    prev_kv_len,
+                    f"Turn {turn_idx + 1}: should inherit {prev_kv_len} KV tokens from previous turn",
+                )
+            prev_kv_len = prompt_tokens + completion_tokens
+
+        # Close the session before checking cache/memory state.
+        ret = requests.post(
+            self.base_url + "/close_session",
+            json={"session_id": session_id},
+        )
+        self.assertEqual(ret.status_code, 200)
+
+        # === Cache verification (after close, before flush) ===
+
+        # Assertion 2: turn 1's prompt was inserted to the cache.
+        verify_resp = requests.post(
+            self.base_url + "/generate",
+            json={
+                "input_ids": chunks_ids[0],
+                "sampling_params": {"temperature": 0, "max_new_tokens": 1},
+            },
+        ).json()
+        self.assertGreater(
+            verify_resp["meta_info"]["cached_tokens"],
+            0,
+            "Turn 1's prompt should be cached in the radix tree",
+        )
+
+        # Assertion 3 (insertion): turn 2's prompt tokens should NOT be in cache.
+        # The tree should only contain turn 1's extent (prompt + output from
+        # cache_unfinished_req during decode). Turn 2's prompt starts fresh tokens
+        # that were never inserted.
+        verify_resp2 = requests.post(
+            self.base_url + "/generate",
+            json={
+                "input_ids": chunks_ids[1],
+                "sampling_params": {"temperature": 0, "max_new_tokens": 1},
+            },
+        ).json()
+        self.assertEqual(
+            verify_resp2["meta_info"]["cached_tokens"],
+            0,
+            "Turn 2's prompt should not be in cache (no insertion for turns 2+)",
+        )
+
+        # === Memory verification ===
+
+        # Assertion 4 & 5: KV is released properly and no memory leak.
+        # SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_IDLE is True by default;
+        # the scheduler will crash if it detects a leak during idle.
+        time.sleep(2)
+        health_resp = requests.get(self.base_url + "/health")
+        self.assertEqual(
+            health_resp.status_code,
+            200,
+            "Server should be healthy after session close (no memory leak)",
+        )
+
+        # After flush, all cache should be reclaimed.
+        requests.post(self.base_url + "/flush_cache")
+        verify_resp3 = requests.post(
+            self.base_url + "/generate",
+            json={
+                "input_ids": chunks_ids[0],
+                "sampling_params": {"temperature": 0, "max_new_tokens": 1},
+            },
+        ).json()
+        self.assertEqual(
+            verify_resp3["meta_info"]["cached_tokens"],
+            0,
+            "After session close + flush, cache should be fully reclaimed",
+        )
+
+        # === Part 2: non-session baseline for output comparison ===
+        requests.post(self.base_url + "/flush_cache")
+
+        outputs_normal = []
+        input_ids = chunks_ids[0][:]
+        for i in range(len(chunks_ids)):
+            response = requests.post(
+                self.base_url + "/generate",
+                json={
+                    "input_ids": input_ids,
+                    "sampling_params": {
+                        "temperature": 0,
+                        "max_new_tokens": gen_len,
+                        "no_stop_trim": True,
+                        "skip_special_tokens": False,
+                    },
+                },
+            ).json()
+            outputs_normal.append(response["text"])
+            if i + 1 < len(chunks_ids):
+                out_ids = tokenizer.encode(response["text"])
+                if out_ids and out_ids[0] == tokenizer.bos_token_id:
+                    out_ids = out_ids[1:]
+                input_ids = input_ids + out_ids + chunks_ids[i + 1]
+
+        print("outputs from streaming session:")
+        print(outputs_from_session)
+        print("outputs from normal queries:")
+        print(outputs_normal)
+        self.assertEqual(outputs_from_session, outputs_normal)
+
     def run_session_control_with_branching(
         self, root_prompt, chunks_per_step, gen_len=16
     ):

test/registered/sessions/test_session_latency.py

@@ -0,0 +1,399 @@
+"""
+Benchmark: Streaming Session Inter-Turn Latency
+
+Measures per-turn latency across three modes as context grows:
+  - no_session:        re-send full context each turn (radix tree prefix match)
+  - regular_session:   session append (radix tree insert + match)
+  - streaming_session: session append (O(1) KV direct transfer)
+
+Each mode runs NUM_CONCURRENT parallel sessions, each doing NUM_TURNS sequential
+requests (16 input / 8 output per turn).
+
+Usage:
+    python -m pytest bench_session_latency.py -s
+    python -m unittest bench_session_latency.BenchSessionLatency.test_streaming_session
+    python -m unittest bench_session_latency.BenchSessionLatency
+"""
+
+import time
+import unittest
+from concurrent.futures import ThreadPoolExecutor
+from dataclasses import dataclass, field
+from typing import Dict, List, Optional
+
+import requests
+from tabulate import tabulate
+
+from sglang.srt.utils import kill_process_tree
+from sglang.srt.utils.hf_transformers_utils import get_tokenizer
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.test_utils import (
+    DEFAULT_MODEL_NAME_FOR_TEST,
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    popen_launch_server,
+)
+
+register_cuda_ci(est_time=100, suite="stage-b-test-large-1-gpu")
+
+NUM_TURNS = 300
+INPUT_LEN = 16
+GEN_LEN = 8
+NUM_CONCURRENT = 4
+TAIL_TURNS = 10
+SAMPLE_TURNS = 8
+
+FILLER_TEXT = (
+    "The quick brown fox jumps over the lazy dog. "
+    "Pack my box with five dozen liquor jugs. "
+    "How vexingly quick daft zebras jump. "
+    "Sphinx of black quartz, judge my vow. "
+) * 200
+
+SAMPLING_PARAMS = {
+    "temperature": 0,
+    "max_new_tokens": GEN_LEN,
+    "no_stop_trim": True,
+    "skip_special_tokens": False,
+    "ignore_eos": True,
+}
+
+
+@dataclass
+class TurnResult:
+    turn: int
+    context_len: int
+    cached_tokens: int
+    prompt_tokens: int
+    completion_tokens: int
+    client_latency_ms: float
+    e2e_latency_ms: float
+
+
+@dataclass
+class ModeResult:
+    mode: str
+    turns: List[TurnResult] = field(default_factory=list)
+    outputs: List[str] = field(default_factory=list)
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def _generate_input_chunks(
+    tokenizer, num_turns: int, input_len: int, offset: int = 0
+) -> List[List[int]]:
+    all_ids = tokenizer.encode(FILLER_TEXT)
+    if all_ids and all_ids[0] == tokenizer.bos_token_id:
+        all_ids = all_ids[1:]
+
+    start = offset * num_turns * input_len
+    needed = start + num_turns * input_len
+    while len(all_ids) < needed:
+        all_ids = all_ids + all_ids
+    chunks = [
+        all_ids[start + i * input_len : start + (i + 1) * input_len]
+        for i in range(num_turns)
+    ]
+
+    if tokenizer.bos_token_id is not None:
+        chunks[0] = [tokenizer.bos_token_id] + chunks[0]
+
+    return chunks
+
+
+def _send_generate(base_url: str, payload: dict) -> dict:
+    resp = requests.post(base_url + "/generate", json=payload)
+    if resp.status_code != 200:
+        raise RuntimeError(f"Generate failed ({resp.status_code}): {resp.text}")
+    return resp.json()
+
+
+def _record_turn(
+    turn_idx: int, context_len: int, meta: dict, client_latency_ms: float
+) -> TurnResult:
+    return TurnResult(
+        turn=turn_idx + 1,
+        context_len=context_len,
+        cached_tokens=meta["cached_tokens"],
+        prompt_tokens=meta["prompt_tokens"],
+        completion_tokens=meta["completion_tokens"],
+        client_latency_ms=client_latency_ms,
+        e2e_latency_ms=meta.get("e2e_latency", 0) * 1000,
+    )
+
+
+# ---------------------------------------------------------------------------
+# Single-session runners (called by worker threads)
+# ---------------------------------------------------------------------------
+
+
+def _run_one_no_session(
+    base_url: str, tokenizer, chunks: List[List[int]]
+) -> ModeResult:
+    result = ModeResult(mode="no_session")
+    accumulated_ids: List[int] = []
+
+    for turn_idx, chunk_ids in enumerate(chunks):
+        accumulated_ids.extend(chunk_ids)
+
+        t0 = time.perf_counter()
+        response = _send_generate(
+            base_url,
+            {"input_ids": accumulated_ids, "sampling_params": SAMPLING_PARAMS},
+        )
+        client_lat = (time.perf_counter() - t0) * 1000
+
+        meta = response["meta_info"]
+        result.turns.append(
+            _record_turn(turn_idx, len(accumulated_ids), meta, client_lat)
+        )
+        result.outputs.append(response["text"])
+
+        output_ids = tokenizer.encode(response["text"])
+        if output_ids and output_ids[0] == tokenizer.bos_token_id:
+            output_ids = output_ids[1:]
+        accumulated_ids.extend(output_ids)
+
+    return result
+
+
+def _run_one_session(
+    base_url: str, chunks: List[List[int]], streaming: bool = False
+) -> ModeResult:
+    mode = "streaming_session" if streaming else "regular_session"
+    result = ModeResult(mode=mode)
+
+    capacity = sum(len(c) for c in chunks) + len(chunks) * GEN_LEN + 1024
+    open_payload: dict = {"capacity_of_str_len": capacity}
+    if streaming:
+        open_payload["streaming"] = True
+    session_id = requests.post(base_url + "/open_session", json=open_payload).json()
+
+    rid = None
+    context_len = 0
+
+    for turn_idx, chunk_ids in enumerate(chunks):
+        context_len += len(chunk_ids)
+
+        t0 = time.perf_counter()
+        response = _send_generate(
+            base_url,
+            {
+                "input_ids": chunk_ids,
+                "session_params": {"id": session_id, "rid": rid},
+                "sampling_params": SAMPLING_PARAMS,
+            },
+        )
+        client_lat = (time.perf_counter() - t0) * 1000
+
+        meta = response["meta_info"]
+        rid = meta["id"]
+        context_len += meta["completion_tokens"]
+
+        result.turns.append(_record_turn(turn_idx, context_len, meta, client_lat))
+        result.outputs.append(response["text"])
+
+    requests.post(base_url + "/close_session", json={"session_id": session_id})
+    return result
+
+
+# ---------------------------------------------------------------------------
+# Stats & reporting
+# ---------------------------------------------------------------------------
+
+
+def _collect_latencies(
+    results: List[ModeResult], last_n: Optional[int] = None
+) -> List[float]:
+    lats = []
+    for r in results:
+        turns = r.turns[1:]  # skip turn 1
+        if last_n is not None:
+            turns = r.turns[-last_n:]
+        lats.extend(t.client_latency_ms for t in turns)
+    return lats
+
+
+def _avg(values: List[float]) -> float:
+    return sum(values) / len(values) if values else 0.0
+
+
+def _print_mode_table(result: ModeResult, label: str = ""):
+    tag = f"{result.mode} ({label})" if label else result.mode
+    print(f"\n  [{tag}]  {len(result.turns)} turns")
+
+    n = len(result.turns)
+    if n <= SAMPLE_TURNS * 2:
+        indices = list(range(n))
+    else:
+        indices = list(range(SAMPLE_TURNS)) + [-1] + list(range(n - SAMPLE_TURNS, n))
+
+    rows = []
+    for idx in indices:
+        if idx == -1:
+            rows.append(["..."] * 5)
+            continue
+        t = result.turns[idx]
+        rows.append(
+            [
+                t.turn,
+                t.context_len,
+                t.cached_tokens,
+                f"{t.client_latency_ms:.1f}ms",
+                f"{t.e2e_latency_ms:.1f}ms",
+            ]
+        )
+    print(
+        tabulate(
+            rows,
+            headers=["Turn", "Context", "Cached", "Client Lat", "E2E Lat"],
+            colalign=("right",) * 5,
+        )
+    )
+
+
+def _print_summary(all_results: Dict[str, List[ModeResult]]):
+    stats = [
+        (
+            mode,
+            _avg(_collect_latencies(rs)),
+            _avg(_collect_latencies(rs, last_n=TAIL_TURNS)),
+        )
+        for mode, rs in all_results.items()
+    ]
+    base_all, base_tail = (stats[0][1] or 1.0), (stats[0][2] or 1.0)
+    tail_label = f"last {TAIL_TURNS}"
+
+    print(f"\n  SUMMARY  ({NUM_CONCURRENT} sessions x {NUM_TURNS} turns)")
+    rows = [
+        [
+            mode,
+            f"{a:.1f}ms",
+            f"{t:.1f}ms",
+            f"{base_all / a:.2f}x" if a else "inf",
+            f"{base_tail / t:.2f}x" if t else "inf",
+        ]
+        for mode, a, t in stats
+    ]
+    print(
+        tabulate(
+            rows,
+            headers=[
+                "Mode",
+                "Avg (all)",
+                f"Avg ({tail_label})",
+                "Speedup (all)",
+                f"Speedup ({tail_label})",
+            ],
+            colalign=("left", "right", "right", "right", "right"),
+        )
+    )
+
+
+# ---------------------------------------------------------------------------
+# Test class
+# ---------------------------------------------------------------------------
+
+
+class BenchSessionLatency(CustomTestCase):
+
+    @classmethod
+    def setUpClass(cls):
+        cls.model = DEFAULT_MODEL_NAME_FOR_TEST
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        cls.process = popen_launch_server(
+            cls.model,
+            cls.base_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=["--attention-backend", "flashinfer"],
+        )
+        cls.tokenizer = get_tokenizer(cls.model)
+
+        requests.post(cls.base_url + "/flush_cache")
+        _send_generate(
+            cls.base_url,
+            {
+                "input_ids": cls.tokenizer.encode("Hello world"),
+                "sampling_params": {"temperature": 0, "max_new_tokens": 1},
+            },
+        )
+
+        cls.all_results: Dict[str, List[ModeResult]] = {}
+
+    @classmethod
+    def tearDownClass(cls):
+        if len(cls.all_results) > 1:
+            _print_summary(cls.all_results)
+        kill_process_tree(cls.process.pid)
+
+    def _run_concurrent_no_session(self) -> List[ModeResult]:
+        requests.post(self.base_url + "/flush_cache")
+
+        def run_one(session_idx):
+            chunks = _generate_input_chunks(
+                self.tokenizer, NUM_TURNS, INPUT_LEN, offset=session_idx
+            )
+            return _run_one_no_session(self.base_url, self.tokenizer, chunks)
+
+        with ThreadPoolExecutor(max_workers=NUM_CONCURRENT) as pool:
+            return list(pool.map(run_one, range(NUM_CONCURRENT)))
+
+    def _run_concurrent_session(self, streaming: bool = False) -> List[ModeResult]:
+        requests.post(self.base_url + "/flush_cache")
+
+        def run_one(session_idx):
+            chunks = _generate_input_chunks(
+                self.tokenizer, NUM_TURNS, INPUT_LEN, offset=session_idx
+            )
+            return _run_one_session(self.base_url, chunks, streaming=streaming)
+
+        with ThreadPoolExecutor(max_workers=NUM_CONCURRENT) as pool:
+            return list(pool.map(run_one, range(NUM_CONCURRENT)))
+
+    # ------------------------------------------------------------------
+    # Test methods
+    # ------------------------------------------------------------------
+
+    def test_no_session(self):
+        results = self._run_concurrent_no_session()
+        self.__class__.all_results["no_session"] = results
+        _print_mode_table(results[0], label="session 0")
+
+    def test_regular_session(self):
+        results = self._run_concurrent_session(streaming=False)
+        self.__class__.all_results["regular_session"] = results
+        _print_mode_table(results[0], label="session 0")
+
+    def test_streaming_session(self):
+        results = self._run_concurrent_session(streaming=True)
+        self.__class__.all_results["streaming_session"] = results
+        _print_mode_table(results[0], label="session 0")
+
+        reg_list = self.__class__.all_results.get("regular_session")
+        if reg_list:
+            reg_out = reg_list[0].outputs
+            stm_out = results[0].outputs
+            mismatches = sum(1 for a, b in zip(reg_out, stm_out) if a != b)
+            self.assertEqual(
+                mismatches,
+                0,
+                f"regular vs streaming (session 0): {mismatches}/{len(reg_out)} turns differ",
+            )
+
+            reg_tail = _avg(_collect_latencies(reg_list, last_n=TAIL_TURNS))
+            stm_tail = _avg(_collect_latencies(results, last_n=TAIL_TURNS))
+            speedup = reg_tail / stm_tail if stm_tail > 0 else float("inf")
+            self.assertGreaterEqual(
+                speedup,
+                2.0,
+                f"streaming should be >=2x faster on last {TAIL_TURNS} turns "
+                f"(regular={reg_tail:.1f}ms, streaming={stm_tail:.1f}ms, speedup={speedup:.2f}x)",
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()