Streaming session: fix retract tail leak via _free_tail (#22862)

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()

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