diff --git a/python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py b/python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py index b02109f2a..f0e86361d 100644 --- a/python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py +++ b/python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py @@ -255,7 +255,7 @@ class DecodeKVCacheOffloadManager: if req.req_pool_idx is None or req.req_pool_idx == -1: return - kv_committed_len = req.pop_committed_kv_cache() + kv_committed_len = req.effective_kv_committed_len() # Free the prefill-aligned slots. Previously this was done # eagerly in offload_kv_cache (mid-decode), which raced with @@ -276,7 +276,7 @@ class DecodeKVCacheOffloadManager: # Free over-allocated KV cache slots (e.g. from speculative decoding v2). # Without spec v2, start_p == end_p so this is a no-op. - start_p, end_p = req.pop_overallocated_kv_cache() + start_p, end_p = kv_committed_len, req.kv.kv_allocated_len if self.page_size > 1: start_p = ceil_align(start_p, self.page_size) if start_p < end_p: diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 46fa64016..071d9d50c 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -749,8 +749,6 @@ class Req(ReqDllmMixin): # For req-level memory management self.kv_committed_len = 0 self.kv: ReqKvInfo = ReqKvInfo(kv_allocated_len=0, swa_evicted_seqlen=0) - self.kv_committed_freed = False - self.kv_overallocated_freed = False # for cross-encoder model self.token_type_ids = token_type_ids @@ -1076,33 +1074,13 @@ class Req(ReqDllmMixin): or self.mamba_host_hit_length > 0 ) - def _cache_commit_len(self) -> int: + def effective_kv_committed_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_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._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.""" - - # NOTE: This function is called when there is over-allocation of KV cache. - # Over-allocation: we allocate more KV cache than the committed length. - # e.g., speculative decoding may allocate more KV cache than actually used. - assert ( - not self.kv_overallocated_freed - ), f"Overallocated KV cache already freed, {self.kv_committed_len=}, {self.kv.kv_allocated_len=}" - self.kv_overallocated_freed = True - return self._cache_commit_len(), self.kv.kv_allocated_len - def update_spec_correct_drafts_histogram(self, num_correct_drafts: int): """Update the speculative decoding acceptance histogram. @@ -1525,8 +1503,6 @@ class Req(ReqDllmMixin): self.already_computed = 0 self.kv.kv_allocated_len = 0 self.kv_committed_len = 0 - self.kv_committed_freed = False - self.kv_overallocated_freed = False self.kv.swa_evicted_seqlen = 0 self.extend_batch_idx = 0 self.decode_batch_idx = 0 diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 20f370fe3..abad9d4aa 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -2564,9 +2564,7 @@ class Scheduler( req.pending_bootstrap = False if self.enable_hicache_storage: self.tree_cache.release_aborted_request(req.rid) - if ( - req.req_pool_idx is not None or self.tree_cache.supports_mamba() - ) and not req.kv_committed_freed: + if req.req_pool_idx is not None or self.tree_cache.supports_mamba(): release_kv_cache(req, self.tree_cache, is_insert=False) self.chunked_req = None diff --git a/python/sglang/srt/managers/scheduler_components/invariant_checker.py b/python/sglang/srt/managers/scheduler_components/invariant_checker.py index b7e0c5123..0b533a84d 100644 --- a/python/sglang/srt/managers/scheduler_components/invariant_checker.py +++ b/python/sglang/srt/managers/scheduler_components/invariant_checker.py @@ -246,8 +246,7 @@ class SchedulerInvariantChecker: swa_uncached = 0 for batch in batches: for req in batch.reqs: - assert req.kv_committed_freed == req.kv_overallocated_freed - if req.kv_committed_freed or req.req_pool_idx is None: + if req.req_pool_idx is None: continue allocated_len = req.kv.kv_allocated_len diff --git a/python/sglang/srt/mem_cache/chunk_cache.py b/python/sglang/srt/mem_cache/chunk_cache.py index e155c5553..78970d058 100644 --- a/python/sglang/srt/mem_cache/chunk_cache.py +++ b/python/sglang/srt/mem_cache/chunk_cache.py @@ -76,11 +76,12 @@ class ChunkCache(BasePrefixCache): # ChunkCache does not support prefix caching, so insert is a no-op return InsertResult(prefix_len=0) - def cache_finished_req(self, req: Req, is_insert: bool = True): - kv_committed_len = req.pop_committed_kv_cache() + def cache_finished_req( + self, req: Req, is_insert: bool = True, *, kv_len_to_handle: int + ): # For decode server: if req.output_ids is empty, we want to free all req.origin_input_ids kv_indices = self.req_to_token_pool.req_to_token[ - req.req_pool_idx, :kv_committed_len + req.req_pool_idx, :kv_len_to_handle ] self.token_to_kv_pool_allocator.free(kv_indices) @@ -150,8 +151,10 @@ class PureSWAChunkCache(SWAChunkCache): prefix is released here when the request finishes. """ - def cache_finished_req(self, req: Req, is_insert: bool = True): - kv_committed_len = req.pop_committed_kv_cache() + def cache_finished_req( + self, req: Req, is_insert: bool = True, *, kv_len_to_handle: int + ): + kv_committed_len = kv_len_to_handle kv_indices = self.req_to_token_pool.req_to_token[ req.req_pool_idx, :kv_committed_len ] diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index a0a7f5708..9e4f0924b 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/python/sglang/srt/mem_cache/common.py @@ -640,17 +640,19 @@ def release_kv_cache(req: Req, tree_cache: BasePrefixCache, is_insert: bool = Tr req.mamba_pool_idx = None return + effective_kv_committed_len = req.effective_kv_committed_len() tree_cache.cache_finished_req( req, is_insert=is_insert and not getattr(req, "skip_radix_cache_insert", False), + kv_len_to_handle=effective_kv_committed_len, ) # StreamingSession.cache_finished_req handles speculative tail trim - # and bookkeeping flag sync internally, then sets req_pool_idx = None. + # internally, then sets req_pool_idx = None. if req.req_pool_idx is None: return - start_p, end_p = req.pop_overallocated_kv_cache() + start_p, end_p = effective_kv_committed_len, req.kv.kv_allocated_len global_server_args = get_server_args() page_size = global_server_args.page_size diff --git a/python/sglang/srt/mem_cache/mamba_radix_cache.py b/python/sglang/srt/mem_cache/mamba_radix_cache.py index 4b2961181..73623430d 100644 --- a/python/sglang/srt/mem_cache/mamba_radix_cache.py +++ b/python/sglang/srt/mem_cache/mamba_radix_cache.py @@ -523,20 +523,21 @@ class MambaRadixCache(KVCacheEventMixin, BasePrefixCache): ) return InsertResult(prefix_len=prefix_len, mamba_exist=mamba_exist) - def cache_finished_req(self, req: Req, is_insert: bool = True) -> None: + def cache_finished_req( + self, req: Req, is_insert: bool = True, *, kv_len_to_handle: int + ) -> None: """Cache request when it finishes.""" - kv_committed_len = req.pop_committed_kv_cache() if self.disable: kv_indices = self.req_to_token_pool.req_to_token[ - req.req_pool_idx, :kv_committed_len + req.req_pool_idx, :kv_len_to_handle ] self.token_to_kv_pool_allocator.free(kv_indices) self.req_to_token_pool.free_mamba_cache(req) return - token_ids = (req.origin_input_ids + req.output_ids)[:kv_committed_len] + token_ids = (req.origin_input_ids + req.output_ids)[:kv_len_to_handle] kv_indices = self.req_to_token_pool.req_to_token[ - req.req_pool_idx, :kv_committed_len + req.req_pool_idx, :kv_len_to_handle ] if is_insert: @@ -572,7 +573,7 @@ class MambaRadixCache(KVCacheEventMixin, BasePrefixCache): assert ( cache_len == page_aligned_len - ), f"It is required {cache_len=}, {page_aligned_len=}, {kv_committed_len=}, {len(req.origin_input_ids)=}, {len(req.output_ids)=} ping @yizhang2077 if you see this" + ), f"It is required {cache_len=}, {page_aligned_len=}, {kv_len_to_handle=}, {len(req.origin_input_ids)=}, {len(req.output_ids)=} ping @yizhang2077 if you see this" # Radix Cache takes one ref in memory pool # insert the token_ids and kv_indices into the radix tree diff --git a/python/sglang/srt/mem_cache/pure_swa_radix_cache.py b/python/sglang/srt/mem_cache/pure_swa_radix_cache.py index 639749c3e..b3d1366df 100644 --- a/python/sglang/srt/mem_cache/pure_swa_radix_cache.py +++ b/python/sglang/srt/mem_cache/pure_swa_radix_cache.py @@ -62,7 +62,9 @@ class PureSWARadixCache(RadixCache): num_tokens = max(params.num_tokens, params.swa_num_tokens) return super().evict(EvictParams(num_tokens=num_tokens)) - def cache_finished_req(self, req: Req, is_insert: bool = True): + def cache_finished_req( + self, req: Req, is_insert: bool = True, *, kv_len_to_handle: int + ): """Cache request when it finishes. Only inserts the prefill portion [0, evict_floor) into the radix tree. @@ -73,7 +75,7 @@ class PureSWARadixCache(RadixCache): if self.disable_finished_insert: is_insert = False - kv_committed_len = req.pop_committed_kv_cache() + kv_committed_len = kv_len_to_handle if self.disable: kv_indices = self.req_to_token_pool.req_to_token[ req.req_pool_idx, :kv_committed_len diff --git a/python/sglang/srt/mem_cache/radix_cache.py b/python/sglang/srt/mem_cache/radix_cache.py index 6ebf09939..b6d416168 100644 --- a/python/sglang/srt/mem_cache/radix_cache.py +++ b/python/sglang/srt/mem_cache/radix_cache.py @@ -434,21 +434,22 @@ class RadixCache(SessionRadixCacheMixin, KVCacheEventMixin, BasePrefixCache): ) return InsertResult(prefix_len=prefix_len, last_device_node=last_node) - def cache_finished_req(self, req: Req, is_insert: bool = True): + def cache_finished_req( + self, req: Req, is_insert: bool = True, *, kv_len_to_handle: int + ): """Cache request when it finishes.""" # In deterministic mode, disable finished request insertion to radix cache if self.disable_finished_insert: is_insert = False - kv_committed_len = req.pop_committed_kv_cache() if self.disable: kv_indices = self.req_to_token_pool.req_to_token[ - req.req_pool_idx, :kv_committed_len + req.req_pool_idx, :kv_len_to_handle ] self.token_to_kv_pool_allocator.free(kv_indices) return - token_ids = (req.origin_input_ids + req.output_ids)[:kv_committed_len] + token_ids = (req.origin_input_ids + req.output_ids)[:kv_len_to_handle] kv_indices = self.req_to_token_pool.req_to_token[ req.req_pool_idx, : len(token_ids) ] diff --git a/python/sglang/srt/mem_cache/radix_cache_cpp.py b/python/sglang/srt/mem_cache/radix_cache_cpp.py index 3729f397e..fb082c1f8 100644 --- a/python/sglang/srt/mem_cache/radix_cache_cpp.py +++ b/python/sglang/srt/mem_cache/radix_cache_cpp.py @@ -169,19 +169,20 @@ class RadixCacheCpp(BasePrefixCache): def total_size(self): return self.tree.total_size() - def cache_finished_req(self, req: Req, is_insert: bool = True): + def cache_finished_req( + self, req: Req, is_insert: bool = True, *, kv_len_to_handle: int + ): """Cache request when it finishes.""" assert req.req_pool_idx is not None - kv_committed_len = req.pop_committed_kv_cache() - token_ids = (req.origin_input_ids + req.output_ids)[:kv_committed_len] + token_ids = (req.origin_input_ids + req.output_ids)[:kv_len_to_handle] kv_indices = self.req_to_token_pool.req_to_token[ - req.req_pool_idx, :kv_committed_len + req.req_pool_idx, :kv_len_to_handle ].to(dtype=torch.int64, copy=True) # NOTE: our C++ implementation don't need `token_ids` and `kv_indices` to be page-aligned # it will automatically align them, but length of them should be equal old_prefix_len = len(req.prefix_indices) // self.page_size * self.page_size - page_aligned_overall_len = kv_committed_len // self.page_size * self.page_size + page_aligned_overall_len = kv_len_to_handle // self.page_size * self.page_size if is_insert: new_prefix_len = self._insert( @@ -200,7 +201,7 @@ class RadixCacheCpp(BasePrefixCache): ) # need to free the unaligned part, since it cannot be inserted into the radix tree - if page_aligned_overall_len < kv_committed_len: + if page_aligned_overall_len < kv_len_to_handle: # NOTE: sglang PagedAllocator support unaligned free (which will automatically align it) self.token_to_kv_pool_allocator.free(kv_indices[page_aligned_overall_len:]) diff --git a/python/sglang/srt/mem_cache/storage/flexkv/flexkv_radix_cache.py b/python/sglang/srt/mem_cache/storage/flexkv/flexkv_radix_cache.py index 6d9f2f618..3c3f8e562 100644 --- a/python/sglang/srt/mem_cache/storage/flexkv/flexkv_radix_cache.py +++ b/python/sglang/srt/mem_cache/storage/flexkv/flexkv_radix_cache.py @@ -379,10 +379,12 @@ class FlexKVRadixCache(RadixCache): # ------------------------------------------------------------------ def cache_finished_req( # type: ignore[override] - self, req: Req, is_insert: bool = True + self, req: Req, is_insert: bool = True, *, kv_len_to_handle: int ) -> None: """Base cache_finished_req then fire an async FlexKV store.""" - super().cache_finished_req(req, is_insert=is_insert) + super().cache_finished_req( + req, is_insert=is_insert, kv_len_to_handle=kv_len_to_handle + ) if not is_insert: self._load_markers.pop(req.rid, None) return diff --git a/python/sglang/srt/mem_cache/storage/lmcache/lmc_radix_cache.py b/python/sglang/srt/mem_cache/storage/lmcache/lmc_radix_cache.py index 96e2f18b0..897a5e61b 100644 --- a/python/sglang/srt/mem_cache/storage/lmcache/lmc_radix_cache.py +++ b/python/sglang/srt/mem_cache/storage/lmcache/lmc_radix_cache.py @@ -428,10 +428,14 @@ class LMCRadixCache(RadixCache): ) ) - def cache_finished_req(self, req: Req, is_insert: bool = True) -> None: + def cache_finished_req( + self, req: Req, is_insert: bool = True, *, kv_len_to_handle: int + ) -> None: """On request completion, insert device KV into radix and store to LMCache.""" - super().cache_finished_req(req, is_insert=is_insert) + super().cache_finished_req( + req, is_insert=is_insert, kv_len_to_handle=kv_len_to_handle + ) if not is_insert: if self._mode is LMCacheMode.MP: self._mp_load_back_markers.pop(req.rid, None) diff --git a/python/sglang/srt/mem_cache/swa_radix_cache.py b/python/sglang/srt/mem_cache/swa_radix_cache.py index ce036adc3..ef1a632c6 100644 --- a/python/sglang/srt/mem_cache/swa_radix_cache.py +++ b/python/sglang/srt/mem_cache/swa_radix_cache.py @@ -456,19 +456,20 @@ class SWARadixCache(KVCacheEventMixin, BasePrefixCache): ) return InsertResult(prefix_len=prefix_len) - def cache_finished_req(self, req: Req, is_insert: bool = True) -> None: + def cache_finished_req( + self, req: Req, is_insert: bool = True, *, kv_len_to_handle: int + ) -> None: """Cache request when it finishes.""" - kv_committed_len = req.pop_committed_kv_cache() if self.disable: kv_indices = self.req_to_token_pool.req_to_token[ - req.req_pool_idx, :kv_committed_len + req.req_pool_idx, :kv_len_to_handle ] self.token_to_kv_pool_allocator.free(kv_indices) return - token_ids = (req.origin_input_ids + req.output_ids)[:kv_committed_len] + token_ids = (req.origin_input_ids + req.output_ids)[:kv_len_to_handle] kv_indices = self.req_to_token_pool.req_to_token[ - req.req_pool_idx, :kv_committed_len + req.req_pool_idx, :kv_len_to_handle ] radix_key = RadixKey( diff --git a/python/sglang/srt/mem_cache/unified_cache_components/README.md b/python/sglang/srt/mem_cache/unified_cache_components/README.md index ab8847031..b98b21628 100644 --- a/python/sglang/srt/mem_cache/unified_cache_components/README.md +++ b/python/sglang/srt/mem_cache/unified_cache_components/README.md @@ -199,14 +199,14 @@ Unlock a previously locked node path. --- -### `cache_finished_req(req: Req, is_insert: bool = True)` +### `cache_finished_req(req: Req, is_insert: bool = True, *, kv_len_to_handle: int)` Cache a completed request's KV data into the tree. | Aspect | Detail | |--------|--------| | **Purpose** | After a request finishes, insert its token/KV data into the tree for future reuse | -| **Inputs** | `req` — the finished request; `is_insert` — whether to insert (True) or just release locks (False) | +| **Inputs** | `req` — the finished request; `is_insert` — whether to insert (True) or just release locks (False); `kv_len_to_handle` — committed KV length supplied by the caller | | **Output** | `None` | | **Mutation** | Calls component hooks → `insert` → `dec_lock_ref` → component cleanup. Frees unaligned tail KV indices; frees non-inserted KV indices when `is_insert=False`. | | **Complexity** | **O(K + D·C)** — insert O(K + D·C) + lock release O(D). Simplifies to **O(K)**. | diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index 9677de9e6..12699930a 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -712,24 +712,24 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache): self._update_evictable_leaf_sets(node) return DecLockRefResult() - def cache_finished_req(self, req: Req, is_insert: bool = True, **kwargs) -> None: + def cache_finished_req( + self, req: Req, is_insert: bool = True, *, kv_len_to_handle: int, **kwargs + ) -> None: if self.session.try_cache_finished_req(req, is_insert=is_insert, **kwargs): return - kv_committed_len = req.pop_committed_kv_cache() - if self.disable: kv_indices = self.req_to_token_pool.req_to_token[ - req.req_pool_idx, :kv_committed_len + req.req_pool_idx, :kv_len_to_handle ] self.token_to_kv_pool_allocator.free(kv_indices) for comp in self._components_tuple: comp.cleanup_after_caching_req(req, is_finished=True) return - token_ids = (req.origin_input_ids + req.output_ids)[:kv_committed_len] + token_ids = (req.origin_input_ids + req.output_ids)[:kv_len_to_handle] kv_indices = self.req_to_token_pool.req_to_token[ - req.req_pool_idx, :kv_committed_len + req.req_pool_idx, :kv_len_to_handle ] result = None diff --git a/python/sglang/srt/session/streaming_session.py b/python/sglang/srt/session/streaming_session.py index a818aeb7d..bc819cbca 100644 --- a/python/sglang/srt/session/streaming_session.py +++ b/python/sglang/srt/session/streaming_session.py @@ -325,7 +325,6 @@ class StreamingSession(BasePrefixCache): self.release_session(session_id) req.req_pool_idx = None req.session.abort_req() - self._mark_kv_freed(req) return True if is_first: @@ -348,7 +347,6 @@ class StreamingSession(BasePrefixCache): # Update req_nodes to this successfully finished request. req.session.finish_req(req) - self._mark_kv_freed(req) return True def try_cache_unfinished_req( @@ -570,14 +568,6 @@ class StreamingSession(BasePrefixCache): 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) -> None: - """Set bookkeeping flags so busy check skips this finished req.""" - if not req.kv_committed_freed: - req.pop_committed_kv_cache() - if not req.kv_overallocated_freed: - req.pop_overallocated_kv_cache() - # -- Pass-through methods -- def evictable_size(self): diff --git a/test/registered/disaggregation/test_specv2_kvcache_offloading.py b/test/registered/disaggregation/test_specv2_kvcache_offloading.py index bcb0e096e..6f08e1da8 100644 --- a/test/registered/disaggregation/test_specv2_kvcache_offloading.py +++ b/test/registered/disaggregation/test_specv2_kvcache_offloading.py @@ -36,22 +36,8 @@ def _make_mock_req( req.req_pool_idx = req_pool_idx req.kv_committed_len = kv_committed_len req.kv = SimpleNamespace(kv_allocated_len=kv_allocated_len) - req.kv_committed_freed = False - req.kv_overallocated_freed = False req.prefix_indices = list(range(prefix_indices_len)) - - def pop_committed(): - assert not req.kv_committed_freed - req.kv_committed_freed = True - return req.kv_committed_len - - def pop_overallocated(): - assert not req.kv_overallocated_freed - req.kv_overallocated_freed = True - return req.kv_committed_len, req.kv.kv_allocated_len - - req.pop_committed_kv_cache = pop_committed - req.pop_overallocated_kv_cache = pop_overallocated + req.effective_kv_committed_len = lambda: req.kv_committed_len return req diff --git a/test/registered/unit/mem_cache/test_decode_radix_lock_ref.py b/test/registered/unit/mem_cache/test_decode_radix_lock_ref.py index 8336b7a74..7cba0af20 100644 --- a/test/registered/unit/mem_cache/test_decode_radix_lock_ref.py +++ b/test/registered/unit/mem_cache/test_decode_radix_lock_ref.py @@ -83,18 +83,10 @@ class MockReq: self.priority = 0 self.kv_committed_len = len(fill_ids) self.kv = SimpleNamespace(kv_allocated_len=len(fill_ids)) - self.kv_committed_freed = False def get_fill_ids(self): return self.full_untruncated_fill_ids[: self.extend_range.end] - def pop_committed_kv_cache(self): - self.kv_committed_freed = True - return self.kv_committed_len - - def pop_overallocated_kv_cache(self): - return (self.kv_committed_len, self.kv.kv_allocated_len) - def _make_req(fill_ids, req_pool_idx=0, cache_protected_len=0, last_node=None): return MockReq(fill_ids, req_pool_idx, cache_protected_len, last_node) @@ -152,7 +144,7 @@ class TestDecodeLockRefScenarios(unittest.TestCase): cache.cache_unfinished_req(req) # Step 3: cache_finished_req with is_insert=True (dec lock) - cache.cache_finished_req(req) + cache.cache_finished_req(req, kv_len_to_handle=req.kv_committed_len) # Verify: all non-root nodes should have lock_ref == 0 # (root always has lock_ref == 1) @@ -201,7 +193,7 @@ class TestDecodeLockRefScenarios(unittest.TestCase): cache.cache_unfinished_req(req) # Step 3: cache_finished_req (dec leaf) - cache.cache_finished_req(req) + cache.cache_finished_req(req, kv_len_to_handle=req.kv_committed_len) # Root lock unchanged, all nodes unlocked self.assertEqual(cache.root_node.lock_ref, root_lock_before) @@ -244,7 +236,9 @@ class TestDecodeLockRefScenarios(unittest.TestCase): # Transfer fails -> cache_finished_req with is_insert=False # This frees delta tokens and dec_lock_ref on last_node - cache.cache_finished_req(req, is_insert=False) + cache.cache_finished_req( + req, is_insert=False, kv_len_to_handle=req.kv_committed_len + ) # The prefix node should be unlocked (back to evictable) self.assertEqual(cache.root_node.lock_ref, 1) @@ -289,7 +283,9 @@ class TestDecodeLockRefScenarios(unittest.TestCase): # Transfer fails -> cache_finished_req with is_insert=False # dec_lock_ref(root) is a no-op - cache.cache_finished_req(req, is_insert=False) + cache.cache_finished_req( + req, is_insert=False, kv_len_to_handle=req.kv_committed_len + ) # Root lock unchanged, nothing protected or evictable self.assertEqual(cache.root_node.lock_ref, root_lock_before) @@ -400,7 +396,7 @@ class TestDecodeLockRefScenarios(unittest.TestCase): ) cache.cache_unfinished_req(req) - cache.cache_finished_req(req) + cache.cache_finished_req(req, kv_len_to_handle=req.kv_committed_len) # After all iterations, root lock should be 1, no protected nodes self.assertEqual(cache.root_node.lock_ref, 1) diff --git a/test/registered/unit/mem_cache/test_pure_swa_chunk_cache.py b/test/registered/unit/mem_cache/test_pure_swa_chunk_cache.py index a4cd9c129..aff31561e 100644 --- a/test/registered/unit/mem_cache/test_pure_swa_chunk_cache.py +++ b/test/registered/unit/mem_cache/test_pure_swa_chunk_cache.py @@ -37,7 +37,7 @@ class TestPureSWAChunkCache(CustomTestCase): ) cache.token_to_kv_pool_allocator = _FakeAllocator() - cache.cache_finished_req(_FakeReq()) + cache.cache_finished_req(_FakeReq(), kv_len_to_handle=8) self.assertEqual(len(cache.token_to_kv_pool_allocator.freed), 1) freed = cache.token_to_kv_pool_allocator.freed[0] diff --git a/test/registered/unit/mem_cache/test_streaming_session_unit.py b/test/registered/unit/mem_cache/test_streaming_session_unit.py index cab86c332..dbaa775d2 100644 --- a/test/registered/unit/mem_cache/test_streaming_session_unit.py +++ b/test/registered/unit/mem_cache/test_streaming_session_unit.py @@ -61,8 +61,6 @@ class _FakeReq: kv_allocated_len=allocated, swa_evicted_seqlen=0, ) - self.kv_committed_freed = False - self.kv_overallocated_freed = False self.origin_input_ids = list(range(committed)) self.output_ids = [] self.extra_key = None @@ -74,22 +72,10 @@ class _FakeReq: self.mamba_next_track_idx = None self.mamba_last_track_seqlen = None self.mamba_branching_seqlen = None - self.pop_overallocated_calls = 0 self.to_finish = None self.finished_reason = None self.finished_len = None - def pop_committed_kv_cache(self): - assert not self.kv_committed_freed - self.kv_committed_freed = True - return self.kv_committed_len - - def pop_overallocated_kv_cache(self): - assert not self.kv_overallocated_freed - self.pop_overallocated_calls += 1 - self.kv_overallocated_freed = True - return self.kv_committed_len, self.kv.kv_allocated_len - def test_preabort_detaches_session_and_preserves_slot(): """Pre-aborted req (to_finish set before match_prefix) is detached from @@ -161,9 +147,6 @@ def test_first_mid_abort_nukes_ephemeral_slot(): assert req_to_token_pool.free_slots == [0] assert len(allocator.freed) == 1 assert allocator.freed[0].tolist() == list(range(20)) - # Bookkeeping flags set. - assert req.kv_committed_freed is True - assert req.kv_overallocated_freed is True def test_nth_mid_abort_nukes_session_slot(): @@ -200,9 +183,6 @@ def test_nth_mid_abort_nukes_session_slot(): # Pool slot returned. assert req_to_token_pool.free_slots == [0] assert req.req_pool_idx is None - # Bookkeeping flags set. - assert req.kv_committed_freed is True - assert req.kv_overallocated_freed is True # Shrink tests removed: streaming sessions are append-only after the diff --git a/test/registered/unit/mem_cache/test_swa_eviction_boundary.py b/test/registered/unit/mem_cache/test_swa_eviction_boundary.py index b1f554d5f..cd8114237 100644 --- a/test/registered/unit/mem_cache/test_swa_eviction_boundary.py +++ b/test/registered/unit/mem_cache/test_swa_eviction_boundary.py @@ -116,7 +116,6 @@ def _make_req(req_pool_idx, token_ids, cache_protected_len, tree): prefix_indices=torch.tensor([], dtype=torch.int64, device=tree.device), _kv_committed_len=len(token_ids), ) - req.pop_committed_kv_cache = lambda: req._kv_committed_len return req @@ -191,7 +190,9 @@ class TestSWAEvictionBoundary(unittest.TestCase): insert_len = seq_len // page_size * page_size self.assertLess(req.kv.swa_evicted_seqlen, insert_len) - tree.cache_finished_req(req, is_insert=True) + tree.cache_finished_req( + req, is_insert=True, kv_len_to_handle=req._kv_committed_len + ) tree.sanity_check() # -- Eviction formula: page_size == 1 -- @@ -216,7 +217,9 @@ class TestSWAEvictionBoundary(unittest.TestCase): req.kv.swa_evicted_seqlen, max(0, seq_len - 1 - max(window, page_size)) ) - tree.cache_finished_req(req, is_insert=True) + tree.cache_finished_req( + req, is_insert=True, kv_len_to_handle=req._kv_committed_len + ) tree.sanity_check() # -- Eviction formula: no-op when seq too short -- @@ -253,7 +256,9 @@ class TestSWAEvictionBoundary(unittest.TestCase): kv1 = _swa_alloc(allocator, first_len) pool.write((0, slice(0, first_len)), kv1) req1 = _make_req(0, list(range(first_len)), 0, tree) - tree.cache_finished_req(req1, is_insert=True) + tree.cache_finished_req( + req1, is_insert=True, kv_len_to_handle=req1._kv_committed_len + ) tree.sanity_check() # Second request: 24 tokens, first 16 overlap with tree @@ -269,7 +274,9 @@ class TestSWAEvictionBoundary(unittest.TestCase): self.assertLessEqual(req2.kv.swa_evicted_seqlen, first_len) swa_evictable_before = tree.swa_evictable_size_ - tree.cache_finished_req(req2, is_insert=True) + tree.cache_finished_req( + req2, is_insert=True, kv_len_to_handle=req2._kv_committed_len + ) # New tokens [16, 24) should all be non-tombstone new_tokens = second_len // page_size * page_size - first_len @@ -300,7 +307,9 @@ class TestSWAEvictionBoundary(unittest.TestCase): self.assertGreater(req.kv.swa_evicted_seqlen, 0, "Should have some eviction") self.assertLess(req.kv.swa_evicted_seqlen, insert_len, "Should be partial") - tree.cache_finished_req(req, is_insert=True) + tree.cache_finished_req( + req, is_insert=True, kv_len_to_handle=req._kv_committed_len + ) non_tombstone = insert_len - req.kv.swa_evicted_seqlen self.assertEqual(tree.swa_evictable_size_, swa_evictable_before + non_tombstone) @@ -338,7 +347,9 @@ class TestSWAEvictionBoundary(unittest.TestCase): req.kv.swa_evicted_seqlen = old_evicted swa_evictable_before = tree.swa_evictable_size_ - tree.cache_finished_req(req, is_insert=True) + tree.cache_finished_req( + req, is_insert=True, kv_len_to_handle=req._kv_committed_len + ) self.assertEqual(tree.swa_evictable_size_, swa_evictable_before) @@ -367,7 +378,9 @@ class TestSWAEvictionBoundary(unittest.TestCase): insert_len = seq_len // page_size * page_size self.assertLess(req.kv.swa_evicted_seqlen, insert_len, f"turn {turn}") - tree.cache_finished_req(req, is_insert=True) + tree.cache_finished_req( + req, is_insert=True, kv_len_to_handle=req._kv_committed_len + ) tree.sanity_check() # -- Integration: page_size=1 full flow -- @@ -391,7 +404,9 @@ class TestSWAEvictionBoundary(unittest.TestCase): req.kv.swa_evicted_seqlen, max(0, seq_len - 1 - max(window, page_size)) ) - tree.cache_finished_req(req, is_insert=True) + tree.cache_finished_req( + req, is_insert=True, kv_len_to_handle=req._kv_committed_len + ) tree.sanity_check() diff --git a/test/registered/unit/mem_cache/test_swa_unittest.py b/test/registered/unit/mem_cache/test_swa_unittest.py index 446277153..534303900 100644 --- a/test/registered/unit/mem_cache/test_swa_unittest.py +++ b/test/registered/unit/mem_cache/test_swa_unittest.py @@ -38,9 +38,6 @@ class _DummyReq: self.swa_prefix_lock_released = False self.kv = SimpleNamespace(swa_evicted_seqlen=0) - def pop_committed_kv_cache(self): - return self._kv_committed_len - def _build_swa_tree( is_eagle: bool, @@ -592,7 +589,9 @@ class TestSWA(unittest.TestCase): return original_insert(params) tree.insert = wrapped_insert - tree.cache_finished_req(req, is_insert=True) + tree.cache_finished_req( + req, is_insert=True, kv_len_to_handle=req._kv_committed_len + ) self.assertEqual(captured["prev_prefix_len"], req.cache_protected_len) self.assertTrue(captured["is_bigram"]) @@ -624,7 +623,9 @@ class TestSWA(unittest.TestCase): return original_free(indices) allocator.free = wrapped_free - tree.cache_finished_req(req2, is_insert=False) + tree.cache_finished_req( + req2, is_insert=False, kv_len_to_handle=req2._kv_committed_len + ) # EAGLE + page_size=1 => page_aligned_len = committed_len - 1 = 5 # Expected frees: diff --git a/test/registered/unit/mem_cache/test_unified_radix_cache_bench.py b/test/registered/unit/mem_cache/test_unified_radix_cache_bench.py index 9202da96a..94e3c9403 100644 --- a/test/registered/unit/mem_cache/test_unified_radix_cache_bench.py +++ b/test/registered/unit/mem_cache/test_unified_radix_cache_bench.py @@ -643,7 +643,6 @@ def bench_cache_finished( req.last_node = node req.cache_protected_len = matched_len req.kv_committed_len = len(seq) - req.kv_committed_freed = False if hasattr(lr, "swa_uuid_for_lock"): req.swa_uuid_for_lock = lr.swa_uuid_for_lock env.rtp.req_to_token[req.req_pool_idx, : len(kv_indices)] = kv_indices @@ -656,7 +655,9 @@ def bench_cache_finished( return bench_api( "cache_finished", lambda: req_items, - lambda req: env.tree.cache_finished_req(req, is_insert=True), + lambda req: env.tree.cache_finished_req( + req, is_insert=True, kv_len_to_handle=req.kv_committed_len + ), len(req_items) - warmup, env.avg_tokens, warmup, 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 b0651ab1d..1e6a614ff 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 @@ -854,7 +854,9 @@ class UnifiedRadixCacheSuite: if self.cfg.has_mamba: req.mamba_last_track_seqlen = kv_len - cache.cache_finished_req(req, is_insert=True) + cache.cache_finished_req( + req, is_insert=True, kv_len_to_handle=req.effective_kv_committed_len() + ) all_ids = input_ids + output_ids aligned_len = (len(all_ids) // ps) * ps @@ -893,8 +895,10 @@ class UnifiedRadixCacheSuite: get_server_args().strip_thinking_cache = True try: avail_before = allocator.available_size() - cache.cache_finished_req(req, is_insert=True) - start_p, end_p = req.pop_overallocated_kv_cache() + cache.cache_finished_req( + req, is_insert=True, kv_len_to_handle=req.effective_kv_committed_len() + ) + start_p, end_p = req.effective_kv_committed_len(), req.kv.kv_allocated_len finally: get_server_args().strip_thinking_cache = False if ps > 1: @@ -936,7 +940,9 @@ class UnifiedRadixCacheSuite: ) avail_before = allocator.available_size() - cache.cache_finished_req(req, is_insert=False) + cache.cache_finished_req( + req, is_insert=False, kv_len_to_handle=req.effective_kv_committed_len() + ) self.assertEqual(allocator.available_size(), avail_before + kv_len) m = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", tokens)))) @@ -1108,7 +1114,9 @@ class UnifiedRadixCacheSuite: req.mamba_last_track_seqlen = kv_len avail_before = allocator.available_size() - cache.cache_finished_req(req, is_insert=True) + cache.cache_finished_req( + req, is_insert=True, kv_len_to_handle=req.effective_kv_committed_len() + ) self.assertEqual(allocator.available_size(), avail_before + tail_extra) aligned = input_ids[: (len(input_ids) // ps) * ps]