Opt-in strip of thinking tokens from radix cache (#23315)

Co-authored-by: ianliuy <ianl@alumni.usc.edu>
Co-authored-by: Wen-xuan-Xu <lilmeep727@gmail.com>

python/sglang/srt/managers/schedule_batch.py

@@ -905,13 +905,20 @@ class Req(ReqDllmMixin):
             return self.output_ids[: self.finished_len]
         return self.output_ids
 
+    def _cache_commit_len(self) -> int:
+        # Report only the prompt prefix so thinking + answer fall into the
+        # overallocated range and are reclaimed by release_kv_cache. #22373.
+        if get_global_server_args().strip_thinking_cache and self.reasoning_tokens > 0:
+            return min(self.kv_committed_len, len(self.origin_input_ids))
+        return self.kv_committed_len
+
     def pop_committed_kv_cache(self) -> int:
         """Return the length of committed KV cache and mark them as freed."""
         assert (
             not self.kv_committed_freed
         ), f"Committed KV cache already freed ({self.kv_committed_len=})"
         self.kv_committed_freed = True
-        return self.kv_committed_len
+        return self._cache_commit_len()
 
     def pop_overallocated_kv_cache(self) -> Tuple[int, int]:
         """Return the range of over-allocated KV cache and mark them as freed."""
@@ -923,7 +930,7 @@ class Req(ReqDllmMixin):
             not self.kv_overallocated_freed
         ), f"Overallocated KV cache already freed, {self.kv_committed_len=}, {self.kv_allocated_len=}"
         self.kv_overallocated_freed = True
-        return self.kv_committed_len, self.kv_allocated_len
+        return self._cache_commit_len(), self.kv_allocated_len
 
     def update_spec_acceptance_histogram(self, accepted_draft_tokens: int):
         """Update the speculative decoding acceptance histogram.

python/sglang/srt/mem_cache/common.py

@@ -489,7 +489,9 @@ def release_kv_cache(req: Req, tree_cache: BasePrefixCache, is_insert: bool = Tr
     page_size = global_server_args.page_size
     spec_algo = global_server_args.speculative_algorithm
 
-    if spec_algo is None:
+    # strip_thinking_cache intentionally reports output tokens as overallocated
+    # so they fall into the free path below (#22373).
+    if spec_algo is None and not global_server_args.strip_thinking_cache:
         assert (
             start_p == end_p
         ), f"Unexpected overallocated KV cache, {req.kv_committed_len=}, {req.kv_allocated_len=}"

python/sglang/srt/server_args.py

@@ -443,6 +443,7 @@ class ServerArgs:
     file_storage_path: str = "sglang_storage"
     enable_cache_report: bool = False
     reasoning_parser: Optional[str] = None
+    strip_thinking_cache: bool = False
     tool_call_parser: Optional[str] = None
     tool_server: Optional[str] = None
     sampling_defaults: str = "model"
@@ -4904,6 +4905,13 @@ class ServerArgs:
             default=ServerArgs.reasoning_parser,
             help=f"Specify the parser for reasoning models, supported parsers are: {list(ReasoningParser.DetectorMap.keys())}.",
         )
+        parser.add_argument(
+            "--strip-thinking-cache",
+            action="store_true",
+            help="Skip caching reasoning-model output (thinking + answer) in the "
+            "radix tree on finish; keep only the prompt prefix. Opt-in: changes "
+            "cache contents.",
+        )
         tool_call_parser_choices = list(FunctionCallParser.ToolCallParserEnum.keys())
         parser.add_argument(
             "--tool-call-parser",

test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py

@@ -30,7 +30,11 @@ from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool, SWATokenToKVPoolAllo
 from sglang.srt.mem_cache.unified_cache_components.tree_component import ComponentType
 from sglang.srt.mem_cache.unified_radix_cache import UnifiedRadixCache
 from sglang.srt.sampling.sampling_params import SamplingParams
-from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
+from sglang.srt.server_args import (
+    ServerArgs,
+    get_global_server_args,
+    set_global_server_args_for_scheduler,
+)
 from sglang.srt.utils import get_device
 from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.test_utils import CustomTestCase
@@ -485,6 +489,53 @@ class UnifiedRadixCacheSuite:
         self.assertEqual(len(m.device_indices), aligned_len)
         tree.sanity_check()
 
+    def test_cache_finished_req_strips_thinking(self):
+        tree, allocator, req_to_token_pool = build_fixture(self.cfg)
+        ps = self.cfg.page_size
+
+        req = self._make_req(req_to_token_pool)
+        prompt_ids = self._make_seq(1, 3)
+        output_ids = self._make_seq(2000, 7)
+        req.origin_input_ids = prompt_ids
+        req.output_ids = output_ids
+        req.fill_ids = prompt_ids + output_ids
+        kv_len = len(req.fill_ids)
+        kv_indices = self._alloc(allocator, kv_len)
+        req_to_token_pool.write((req.req_pool_idx, slice(0, kv_len)), kv_indices)
+        req.kv_committed_len = kv_len
+        req.kv_allocated_len = kv_len
+        req.last_node = tree.root_node
+        req.cache_protected_len = 0
+        req.swa_uuid_for_lock = None
+        req.extra_key = None
+        if self.cfg.has_mamba:
+            req.mamba_last_track_seqlen = kv_len
+        req.reasoning_tokens = 1
+
+        get_global_server_args().strip_thinking_cache = True
+        try:
+            avail_before = allocator.available_size()
+            tree.cache_finished_req(req, is_insert=True)
+            start_p, end_p = req.pop_overallocated_kv_cache()
+        finally:
+            get_global_server_args().strip_thinking_cache = False
+        if ps > 1:
+            start_p = ((start_p + ps - 1) // ps) * ps
+        if start_p < end_p:
+            allocator.free(
+                req_to_token_pool.req_to_token[req.req_pool_idx][start_p:end_p]
+            )
+
+        prompt_aligned = (len(prompt_ids) // ps) * ps
+        # Thinking+answer must not be reachable past the prompt.
+        m = tree.match_prefix(MatchPrefixParams(key=RadixKey(prompt_ids + output_ids)))
+        self.assertEqual(len(m.device_indices), prompt_aligned)
+        # Only prompt-aligned pages remain owned by the tree.
+        self.assertEqual(
+            allocator.available_size(), avail_before + kv_len - prompt_aligned
+        )
+        tree.sanity_check()
+
     def test_cache_finished_req_no_insert(self):
         tree, allocator, req_to_token_pool = build_fixture(self.cfg)
         req = self._make_req(req_to_token_pool)