From efc267ca29154636c06a58d7f70fa1636f021b78 Mon Sep 17 00:00:00 2001
From: Liangsheng Yin <hnyls2002@gmail.com>
Date: Wed, 15 Apr 2026 14:15:46 -0700
Subject: [PATCH] streaming session: trim spec v2 overshoot in
 cache_finished_req (#22897)

---
 .../srt/mem_cache/session_aware_cache.py      | 56 +++++++++++++------
 1 file changed, 38 insertions(+), 18 deletions(-)

diff --git a/python/sglang/srt/mem_cache/session_aware_cache.py b/python/sglang/srt/mem_cache/session_aware_cache.py
index bf40a3cde..be56c2b0c 100644
--- a/python/sglang/srt/mem_cache/session_aware_cache.py
+++ b/python/sglang/srt/mem_cache/session_aware_cache.py
@@ -263,6 +263,11 @@ class SessionAwareCache(BasePrefixCache):
             slot = SessionSlot()
             self.slots[session_id] = slot
 
+        finished_len = (
+            req.finished_len if req.finished_len is not None else len(req.output_ids)
+        )
+        self._trim_overshoot(req, finished_len)
+
         slot.save_from_req(req, is_first=is_first)
 
         # Update req_nodes to this successfully finished request.
@@ -271,25 +276,12 @@ 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.
+        """match_prefix path: free orphaned KV in [prefix_len, kv_allocated_len)
+        before 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.
         """
-        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)
+        self._free_kv_aligned(slot.req_pool_idx, prefix_len, slot.kv_allocated_len)
         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)
@@ -297,6 +289,34 @@ class SessionAwareCache(BasePrefixCache):
         req.kv_committed_len = min(req.kv_committed_len, prefix_len)
         req.swa_evicted_seqlen = min(req.swa_evicted_seqlen, prefix_len)
 
+    def _trim_overshoot(self, req: Req, finished_len: int):
+        """Trim slot KV to finished_len boundary. Spec v2 may overshoot
+        max_new_tokens (verify round commits M+1 at a time); next turn's
+        input is output_ids[:finished_len], so positions past that must
+        be released to avoid token/KV mismatch.
+        """
+        target = len(req.origin_input_ids) + finished_len
+        self._free_kv_aligned(req.req_pool_idx, target, req.kv_allocated_len)
+        req.kv_allocated_len = min(req.kv_allocated_len, target)
+        req.kv_committed_len = min(req.kv_committed_len, target)
+        req.output_ids = req.output_ids[:finished_len]
+
+    def _free_kv_aligned(self, pool_idx: int, target: int, end: int):
+        """Free req_to_token[pool_idx, ceil_align(target):end). Page-aligned
+        because PagedTokenToKVPoolAllocator.free returns whole pages
+        (free_index // page_size), so partial-page free would corrupt pages
+        still holding committed tokens. The range [target, ceil_align(target))
+        stays attached until release_session frees the whole page.
+        """
+        if end <= target:
+            return
+        start = target
+        if self.page_size > 1:
+            start = ceil_align(start, self.page_size)
+        if start < end:
+            tail = self.req_to_token_pool.req_to_token[pool_idx, start:end]
+            self.token_to_kv_pool_allocator.free(tail)
+
     @staticmethod
     def _mark_kv_freed(req: Req):
         """Set bookkeeping flags so busy check skips this finished req."""