Unify memory management across (overlap, non-overlap) x (page>=1) x (spec, non-spec, spec v2) x (retract, finished) (#12224)

2025-11-11 02:56:22 +08:00
parent 838bcb0d93
commit 665416f6dd
24 changed files with 193 additions and 156 deletions
@@ -52,10 +52,10 @@ class ChunkCache(BasePrefixCache):
        )

    def cache_finished_req(self, req: Req, is_insert: bool = True):
+        kv_committed_len = req.pop_committed_kv_cache()
+        # 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,
-            # For decode server: if req.output_ids is empty, we want to free all req.origin_input_ids
-            : len(req.origin_input_ids) + max(len(req.output_ids) - 1, 0),
+            req.req_pool_idx, :kv_committed_len
        ]
        self.req_to_token_pool.free(req.req_pool_idx)
        self.token_to_kv_pool_allocator.free(kv_indices)