diff --git a/python/sglang/srt/entrypoints/engine.py b/python/sglang/srt/entrypoints/engine.py index 8246c849e..42ce76841 100644 --- a/python/sglang/srt/entrypoints/engine.py +++ b/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) diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py index d5a8338b9..9326c691c 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/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 diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 5576cd577..03f86c615 100644 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/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(), ) ) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 8f1dc0db4..ab51a1a52 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/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: diff --git a/python/sglang/srt/managers/scheduler_output_processor_mixin.py b/python/sglang/srt/managers/scheduler_output_processor_mixin.py index 01eb9a16d..cafa87b09 100644 --- a/python/sglang/srt/managers/scheduler_output_processor_mixin.py +++ b/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) diff --git a/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py b/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py index 484a949f5..0ac5563c9 100644 --- a/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py +++ b/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 - # if not self.enable_hierarchical_cache - # else self.max_total_num_tokens - protected_size - self.max_total_num_tokens - - protected_size + self.max_total_num_tokens - protected_size - session_held ) - 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( diff --git a/python/sglang/srt/managers/session_controller.py b/python/sglang/srt/managers/session_controller.py index 64ea7232e..bbc619191 100644 --- a/python/sglang/srt/managers/session_controller.py +++ b/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 diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index 76de4912b..5a759ed11 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/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() diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index c51e004fb..d654b7f31 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/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 diff --git a/python/sglang/srt/mem_cache/session_aware_cache.py b/python/sglang/srt/mem_cache/session_aware_cache.py new file mode 100644 index 000000000..512e1347b --- /dev/null +++ b/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) diff --git a/test/manual/test_session_control.py b/test/registered/sessions/test_session_control.py similarity index 82% rename from test/manual/test_session_control.py rename to test/registered/sessions/test_session_control.py index 99b112802..d74a39f67 100644 --- a/test/manual/test_session_control.py +++ b/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 ): diff --git a/test/registered/sessions/test_session_latency.py b/test/registered/sessions/test_session_latency.py new file mode 100644 index 000000000..7298641db --- /dev/null +++ b/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()