From a490632416fdc7b2edef5b0e551b9bffdc189dbf Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Mon, 20 Apr 2026 22:59:50 -0700 Subject: [PATCH] Opt-in strip of thinking tokens from radix cache (#23315) Co-authored-by: ianliuy Co-authored-by: Wen-xuan-Xu --- python/sglang/srt/managers/schedule_batch.py | 11 +++- python/sglang/srt/mem_cache/common.py | 4 +- python/sglang/srt/server_args.py | 8 +++ .../test_unified_radix_cache_unittest.py | 53 ++++++++++++++++++- 4 files changed, 72 insertions(+), 4 deletions(-) diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 771584e44..de7405434 100644 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/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. diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index 59c1dd974..555bd97a4 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/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=}" diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 51e408dc7..94158e7b5 100644 --- a/python/sglang/srt/server_args.py +++ b/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", diff --git a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py index 7014290bd..eab8f6bfc 100644 --- a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py +++ b/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)