From ce31934ca80e87b7b0769e142e27daafc517e740 Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Wed, 15 Apr 2026 01:44:27 -0700 Subject: [PATCH] Streaming session: fix retract tail leak via _free_tail (#22862) --- python/sglang/srt/mem_cache/common.py | 44 +------------- .../srt/mem_cache/session_aware_cache.py | 59 +++++++++++++++---- 2 files changed, 49 insertions(+), 54 deletions(-) diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index 1a43f29cb..f59d9c7bc 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/python/sglang/srt/mem_cache/common.py @@ -9,7 +9,6 @@ import triton.language as tl from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache, EvictParams from sglang.srt.mem_cache.memory_pool import HybridReqToTokenPool, ReqToTokenPool -from sglang.srt.mem_cache.session_aware_cache import SessionAwareCache, _is_streaming from sglang.srt.mem_cache.swa_memory_pool import SWATokenToKVPoolAllocator from sglang.srt.server_args import get_global_server_args from sglang.srt.utils import support_triton @@ -477,50 +476,11 @@ def release_kv_cache(req: Req, tree_cache: BasePrefixCache, is_insert: bool = Tr req.mamba_pool_idx = None return - # Streaming sessions transfer req_pool ownership into SessionSlot objects. - # Trim any speculative tail before that transfer, otherwise later turns - # restore only the committed prefix and can strand unreachable KV pages. - # - # Aborted streaming-session requests (e.g. input too long) skip the - # streaming path entirely. match_prefix did not restore the slot's KV - # state, so the request has a fresh pool slot that should be freed by - # cache_finished_req below (which also sets req_pool_idx = None). - from sglang.srt.managers.schedule_batch import FINISH_ABORT - - is_streaming_session = isinstance(tree_cache, SessionAwareCache) and _is_streaming( - req - ) - is_aborted_streaming = is_streaming_session and isinstance( - getattr(req, "finished_reason", None), FINISH_ABORT - ) - if is_streaming_session and not is_aborted_streaming: - start_p, end_p = req.pop_overallocated_kv_cache() - page_size = get_global_server_args().page_size - if page_size > 1: - start_p = ceil_align(start_p, page_size) - if start_p < end_p: - indices_to_free = tree_cache.req_to_token_pool.req_to_token[ - req.req_pool_idx - ][start_p:end_p] - tree_cache.token_to_kv_pool_allocator.free(indices_to_free) - req.kv_allocated_len = req.kv_committed_len - tree_cache.cache_finished_req(req, is_insert=is_insert) - # SessionAwareCache.cache_finished_req sets req_pool_idx = None to transfer - # KV ownership to the SessionSlot, so the remaining cleanup is skipped. - # Streaming-session specific overalloc trimming must therefore happen - # before cache_finished_req above. + # SessionAwareCache.cache_finished_req handles speculative tail trim + # and bookkeeping flag sync internally, then sets req_pool_idx = None. if req.req_pool_idx is None: - if is_streaming_session: - # The request no longer owns any KV once SessionAwareCache either - # transfers it into the session slot or frees it on abort. Mark - # both bookkeeping flags so busy-time memory checks do not keep - # counting this finished request as uncached KV. - if not req.kv_committed_freed: - req.pop_committed_kv_cache() - if not req.kv_overallocated_freed: - req.pop_overallocated_kv_cache() return start_p, end_p = req.pop_overallocated_kv_cache() diff --git a/python/sglang/srt/mem_cache/session_aware_cache.py b/python/sglang/srt/mem_cache/session_aware_cache.py index 36bde595b..bf40a3cde 100644 --- a/python/sglang/srt/mem_cache/session_aware_cache.py +++ b/python/sglang/srt/mem_cache/session_aware_cache.py @@ -193,10 +193,25 @@ class SessionAwareCache(BasePrefixCache): slot.restore_to_req(req) - # logprob_start_len is already forced to -1 for streaming sessions - # (in Req.init_next_round_input), so the prefix key is not truncated - # and we can directly reuse the committed KV length. - prefix_len = min(req.kv_committed_len, max(len(params.key.token_ids) - 1, 0)) + # token_ids = fill_ids[:input_len-1] (1-token logit reserve already + # applied). min handles retract retry where committed_len can + # exceed len(token_ids) by 1. + prefix_len = min(req.kv_committed_len, len(params.key.token_ids)) + + # Streaming sessions are append-only (session_controller rollback + # ensures req_nodes always points to the last successful req). + assert prefix_len >= slot.cache_protected_len, ( + f"streaming session prefix shrank: {prefix_len=} < " + f"{slot.cache_protected_len=}" + ) + + # Free orphaned tail: alloc_for_extend will overwrite + # req_to_token[prefix_len:] with new indices. The range + # [prefix_len, kv_allocated_len) has stale indices from the + # previous turn's decode (e.g. alloc-commit gap on retract, + # or speculative draft tokens). + self._free_tail(slot, req, prefix_len) + device_indices = self.req_to_token_pool.req_to_token[ req.req_pool_idx, :prefix_len ].to(dtype=torch.int64) @@ -255,6 +270,33 @@ class SessionAwareCache(BasePrefixCache): self._mark_kv_freed(req) + def _free_tail(self, slot: SessionSlot, req: Req, prefix_len: int): + """Free KV in [prefix_len, kv_allocated_len) before the next + alloc_for_extend overwrites it. The gap appears when spec + decoding pushes allocated above committed, or when retract + retry's logit-reserve pulls prefix_len below committed. + Free start is ceil-aligned to page_size: PagedTokenToKVPoolAllocator + frees by whole pages, so partial-page free would corrupt pages + still holding committed tokens; the gap stays attached until + release_session. + """ + if prefix_len >= slot.kv_allocated_len: + return + free_start = prefix_len + if self.page_size > 1: + free_start = ceil_align(free_start, self.page_size) + if free_start < slot.kv_allocated_len: + tail_indices = self.req_to_token_pool.req_to_token[ + slot.req_pool_idx, free_start : slot.kv_allocated_len + ] + self.token_to_kv_pool_allocator.free(tail_indices) + slot.kv_allocated_len = prefix_len + slot.kv_committed_len = min(slot.kv_committed_len, prefix_len) + slot.swa_evicted_seqlen = min(slot.swa_evicted_seqlen, prefix_len) + req.kv_allocated_len = prefix_len + req.kv_committed_len = min(req.kv_committed_len, prefix_len) + req.swa_evicted_seqlen = min(req.swa_evicted_seqlen, prefix_len) + @staticmethod def _mark_kv_freed(req: Req): """Set bookkeeping flags so busy check skips this finished req.""" @@ -298,14 +340,7 @@ class SessionAwareCache(BasePrefixCache): # -- Session lifecycle -- def release_session(self, session_id: str): - """Release all KV resources held by a streaming session. - - `slot.last_node` + `slot.cache_protected_len` are trusted directly: radix - tree splits mutate TreeNode objects in place (see RadixCache._split_node), - so a saved TreeNode reference remains valid and the locked prefix length - is unchanged. No rematch needed -- and `match_prefix` here would cause - further splits that disturb accounting. - """ + """Release all KV resources held by a streaming session.""" slot = self.slots.pop(session_id, None) if slot is None: return