[Session] Add streaming mode with SessionAwareCache fast path (#19171)
Co-authored-by: hnyls2002 <lsyincs@gmail.com>
This commit is contained in:
Aurick Qiao
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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(),
|
||||
)
|
||||
)
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
+161
@@ -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
|
||||
):
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user