From 2b53e660de89c142adab497c97edf4e98e1de621 Mon Sep 17 00:00:00 2001
From: Liangsheng Yin <hnyls2002@gmail.com>
Date: Thu, 19 Mar 2026 17:09:40 -0700
Subject: [PATCH] Simplify streaming session logprob handling (#20955)

---
 python/sglang/srt/managers/schedule_batch.py  | 17 +++++++++++++
 .../srt/mem_cache/session_aware_cache.py      | 24 ++++---------------
 2 files changed, 21 insertions(+), 20 deletions(-)

diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py
index 355c302eb..d5ca9f122 100644
--- a/python/sglang/srt/managers/schedule_batch.py
+++ b/python/sglang/srt/managers/schedule_batch.py
@@ -873,6 +873,23 @@ class Req(ReqDllmMixin):
             self.fill_ids = self.origin_input_ids + self.output_ids
 
         input_len = len(self.fill_ids)
+
+        # Streaming sessions reuse committed KV from the session slot, so
+        # custom logprob_start_len is not supported — override to -1.
+        if (
+            self.session is not None
+            and self.session.streaming
+            and self.return_logprob
+            and self.logprob_start_len >= 0
+        ):
+            logger.warning(
+                "logprob_start_len=%d is not supported for streaming sessions "
+                "and will be ignored (rid=%s). Only new-token logprobs are returned.",
+                self.logprob_start_len,
+                self.rid,
+            )
+            self.logprob_start_len = -1
+
         # NOTE: the matched length is at most 1 less than the input length to enable logprob computation
         max_prefix_len = input_len - 1
         if self.return_logprob and self.logprob_start_len >= 0:
diff --git a/python/sglang/srt/mem_cache/session_aware_cache.py b/python/sglang/srt/mem_cache/session_aware_cache.py
index 3f30189ac..0298a3a5a 100644
--- a/python/sglang/srt/mem_cache/session_aware_cache.py
+++ b/python/sglang/srt/mem_cache/session_aware_cache.py
@@ -181,26 +181,10 @@ class SessionAwareCache(BasePrefixCache):
 
         slot.restore_to_req(req)
 
-        # For streaming sessions, ignore the logprob_start_len truncation on
-        # the key and reuse the full committed KV from the slot. The key may
-        # have been truncated (e.g. to length 0 when logprob_start_len=0), but
-        # the session slot already holds those KV tokens — skipping them would
-        # orphan allocated memory. Use fill_ids length (the actual input for
-        # this turn) as the upper bound instead of the truncated key.
-        # Also clamp logprob_start_len so the scheduler doesn't expect logprobs
-        # for tokens that are already in the session's committed KV.
-        input_len = (
-            len(req.fill_ids) if req.fill_ids is not None else len(params.key.token_ids)
-        )
-        max_prefix_len = max(input_len - 1, 0)
-        prefix_len = min(req.kv_committed_len, max_prefix_len)
-
-        if (
-            req.return_logprob
-            and req.logprob_start_len >= 0
-            and req.logprob_start_len < prefix_len
-        ):
-            req.logprob_start_len = prefix_len
+        # 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))
         device_indices = self.req_to_token_pool.req_to_token[
             req.req_pool_idx, :prefix_len
         ].to(dtype=torch.int64)