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

python/sglang/srt/disaggregation/decode.py

@@ -51,6 +51,7 @@ from sglang.srt.managers.schedule_batch import FINISH_ABORT, RequestStage, Sched
 from sglang.srt.managers.utils import GenerationBatchResult
 from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
 from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache
+from sglang.srt.mem_cache.common import release_kv_cache
 from sglang.srt.mem_cache.memory_pool import (
     HybridLinearKVPool,
     HybridReqToTokenPool,
@@ -618,37 +619,21 @@ class DecodePreallocQueue:
 
         req.req_pool_idx = req_pool_indices[0]
 
+        # Alloc all tokens for the prebuilt req (except for the reserved input token for decoding)
+        fill_len = len(req.origin_input_ids) + max(len(req.output_ids) - 1, 0)
+        req.kv_allocated_len = fill_len
+        req.kv_committed_len = fill_len
         if self.token_to_kv_pool_allocator.page_size == 1:
-            kv_loc = self.token_to_kv_pool_allocator.alloc(
+            kv_loc = self.token_to_kv_pool_allocator.alloc(fill_len)
-                len(req.origin_input_ids) + max(len(req.output_ids) - 1, 0)
-            )
         else:
-            num_tokens = len(req.origin_input_ids) + max(len(req.output_ids) - 1, 0)
+            device = self.token_to_kv_pool_allocator.device
             kv_loc = self.token_to_kv_pool_allocator.alloc_extend(
-                prefix_lens=torch.tensor(
+                prefix_lens=torch.tensor([0], dtype=torch.int64, device=device),
-                    [0],
+                prefix_lens_cpu=torch.tensor([0], dtype=torch.int64),
-                    dtype=torch.int64,
+                seq_lens=torch.tensor([fill_len], dtype=torch.int64, device=device),
-                    device=self.token_to_kv_pool_allocator.device,
+                seq_lens_cpu=torch.tensor([fill_len], dtype=torch.int64),
-                ),
+                last_loc=torch.tensor([-1], dtype=torch.int64, device=device),
-                prefix_lens_cpu=torch.tensor(
+                extend_num_tokens=fill_len,
-                    [0],
-                    dtype=torch.int64,
-                ),
-                seq_lens=torch.tensor(
-                    [num_tokens],
-                    dtype=torch.int64,
-                    device=self.token_to_kv_pool_allocator.device,
-                ),
-                seq_lens_cpu=torch.tensor(
-                    [num_tokens],
-                    dtype=torch.int64,
-                ),
-                last_loc=torch.tensor(
-                    [-1],
-                    dtype=torch.int64,
-                    device=self.token_to_kv_pool_allocator.device,
-                ),
-                extend_num_tokens=num_tokens,
             )
 
         assert (
@@ -763,7 +748,7 @@ class DecodeTransferQueue:
                     [decode_req.req], decode_req.req.return_logprob
                 )
                 # release pre-allocated kv cache, but don't insert into the tree since it's failed
-                self.tree_cache.cache_finished_req(decode_req.req, is_insert=False)
+                release_kv_cache(decode_req.req, self.tree_cache, is_insert=False)
                 indices_to_remove.add(i)
                 if self.scheduler.enable_metrics:
                     self.scheduler.metrics_collector.increment_transfer_failed_reqs()

python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py

@@ -1,6 +1,9 @@
+from __future__ import annotations
+
 import logging
 import threading
 import time
+from typing import TYPE_CHECKING
 
 import torch
 
@@ -18,6 +21,9 @@ from sglang.srt.mem_cache.memory_pool_host import (
 )
 from sglang.srt.server_args import ServerArgs
 
+if TYPE_CHECKING:
+    from sglang.srt.managers.schedule_batch import Req
+
 logger = logging.getLogger(__name__)
 
 
@@ -177,14 +183,14 @@ class DecodeKVCacheOffloadManager:
                 )
             finish_count -= 1
 
-    def _release_finished_req(self, req, prefill_offloaded_len):
+    def _release_finished_req(self, req: Req, prefill_offloaded_len: int):
+        # FIXME: not sure which length to use here: kv_allocated_len or kv_committed_len
         kv_indices = self.req_to_token_pool.req_to_token[
-            req.req_pool_idx,
+            req.req_pool_idx, prefill_offloaded_len : req.kv_allocated_len
-            : len(req.origin_input_ids) + max(len(req.output_ids) - 1, 0),
         ]
 
         # Free the incremental part of the request
-        self.token_to_kv_pool_allocator.free(kv_indices[prefill_offloaded_len:])
+        self.token_to_kv_pool_allocator.free(kv_indices)
         self.req_to_token_pool.free(req.req_pool_idx)
 
     def _check_backup_progress(self, finish_count):

python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py

@@ -7,6 +7,7 @@ from typing import TYPE_CHECKING
 import torch
 
 from sglang.srt.disaggregation.utils import prepare_abort
+from sglang.srt.mem_cache.common import release_kv_cache
 from sglang.srt.model_executor.forward_batch_info import CaptureHiddenMode, ForwardMode
 from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo
 
@@ -119,7 +120,7 @@ class ScheduleBatchDisaggregationDecodeMixin:
                     # Grammar accept_token can raise ValueError if the token is not in the grammar.
                     # This can happen if the grammar is not set correctly or the token is invalid.
                     error_message = f"Grammar accept_token failed for req {req.rid} with token {req.output_ids[-1]}: {e}"
-                    self.tree_cache.cache_finished_req(req)
+                    release_kv_cache(req, self.tree_cache)
                     prepare_abort(
                         req, error_message, status_code=HTTPStatus.INTERNAL_SERVER_ERROR
                     )

python/sglang/srt/disaggregation/prefill.py

@@ -48,6 +48,7 @@ from sglang.srt.managers.schedule_batch import (
     RequestStage,
     ScheduleBatch,
 )
+from sglang.srt.mem_cache.common import release_kv_cache
 from sglang.srt.mem_cache.memory_pool import (
     HybridLinearKVPool,
     NSATokenToKVPool,
@@ -468,7 +469,7 @@ class SchedulerDisaggregationPrefillMixin:
                         # Grammar accept_token can raise ValueError if the token is not in the grammar.
                         # This can happen if the grammar is not set correctly or the token is invalid.
                         error_message = f"Grammar accept_token failed for req {req.rid} with token {next_token_id}: {e}"
-                        self.tree_cache.cache_finished_req(req)
+                        release_kv_cache(req, self.tree_cache)
                         prepare_abort(
                             req,
                             error_message,
@@ -534,7 +535,7 @@ class SchedulerDisaggregationPrefillMixin:
             if poll in [KVPoll.WaitingForInput, KVPoll.Transferring]:
                 undone_reqs.append(req)
             elif poll == KVPoll.Success:  # transfer done
-                self.tree_cache.cache_finished_req(req)  # unlock the tree
+                release_kv_cache(req, self.tree_cache)  # unlock the tree
                 req.finished_reason = FINISH_LENGTH(length=0)
                 # FIXME: clean up req's data in transfer engine
                 if hasattr(req.disagg_kv_sender, "clear"):
@@ -547,7 +548,7 @@ class SchedulerDisaggregationPrefillMixin:
                 except Exception as e:
                     error_message += f" with exception {e}"
                 logger.warning(error_message)
-                self.tree_cache.cache_finished_req(req)  # unlock the tree
+                release_kv_cache(req, self.tree_cache)  # unlock the tree
                 prepare_abort(
                     req, error_message, status_code=HTTPStatus.INTERNAL_SERVER_ERROR
                 )

python/sglang/srt/layers/attention/nsa/nsa_indexer.py

@@ -8,7 +8,7 @@ from einops import rearrange
 
 from sglang.srt.custom_op import CustomOp
 from sglang.srt.layers.layernorm import LayerNorm
-from sglang.srt.utils import add_prefix, align, is_cuda, is_hip, is_npu
+from sglang.srt.utils import add_prefix, ceil_align, is_cuda, is_hip, is_npu
 
 if is_cuda():
     try:
@@ -114,7 +114,7 @@ class Indexer(CustomOp):
         self.fuse_wk_and_weights_proj = fuse_wk_and_weights_proj
         if is_cuda():
             self.sm_count = deep_gemm.get_num_sms()
-            self.half_device_sm_count = align(self.sm_count // 2, 8)
+            self.half_device_sm_count = ceil_align(self.sm_count // 2, 8)
 
         self.wq_b = ReplicatedLinear(
             self.q_lora_rank,
@@ -511,7 +511,7 @@ class Indexer(CustomOp):
             end_pos = seq_len
             topk_indices = index_score.topk(min(topk, end_pos), dim=-1)[1].squeeze(0)
 
-            pad_len = align(topk_indices.shape[-1], 2048) - topk_indices.shape[-1]
+            pad_len = ceil_align(topk_indices.shape[-1], 2048) - topk_indices.shape[-1]
             topk_indices = torch.nn.functional.pad(
                 topk_indices, (0, pad_len), "constant", -1
             )

python/sglang/srt/layers/quantization/fp8_kernel.py

@@ -25,7 +25,7 @@ import triton.language as tl
 
 from sglang.srt.layers import deep_gemm_wrapper
 from sglang.srt.utils import (
-    align,
+    ceil_align,
     direct_register_custom_op,
     get_bool_env_var,
     get_device_core_count,
@@ -442,8 +442,8 @@ def create_per_token_group_quant_fp8_output_scale(
         assert column_major_scales and scale_tma_aligned
         *x_batch, x_q_mn, x_q_k = x_shape
         x_s_mn, x_s_k = x_q_mn, x_q_k // 128
-        aligned_mn = align(x_s_mn, 4)
+        aligned_mn = ceil_align(x_s_mn, 4)
-        aligned_k = align(x_s_k, 4)
+        aligned_k = ceil_align(x_s_k, 4)
         # TODO(FIXME): Fix cuda kernel and recover here to empty.
         return torch.empty(
             (*x_batch, aligned_k // 4, aligned_mn),

python/sglang/srt/layers/quantization/fp8_utils.py

@@ -27,7 +27,7 @@ from sglang.srt.layers.quantization.fp8_kernel import (
     w8a8_block_fp8_matmul_triton,
 )
 from sglang.srt.utils import (
-    align,
+    ceil_align,
     get_bool_env_var,
     get_cuda_version,
     get_device_capability,
@@ -495,7 +495,7 @@ def per_block_cast_to_fp8(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
     assert x.dim() == 2
     m, n = x.shape
     x_padded = torch.zeros(
-        (align(m, 128), align(n, 128)), dtype=x.dtype, device=x.device
+        (ceil_align(m, 128), ceil_align(n, 128)), dtype=x.dtype, device=x.device
     )
     x_padded[:m, :n] = x
     x_view = x_padded.view(-1, 128, x_padded.size(1) // 128, 128)

python/sglang/srt/managers/schedule_batch.py

@@ -66,6 +66,7 @@ from sglang.srt.mem_cache.common import (
     alloc_for_decode,
     alloc_for_extend,
     evict_from_tree_cache,
+    release_kv_cache,
 )
 from sglang.srt.mem_cache.mamba_radix_cache import MambaRadixCache
 from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
@@ -479,6 +480,12 @@ class Req:
         self.session_id = session_id
         self.input_embeds = input_embeds
 
+        # For req-level memory management
+        self.kv_committed_len = 0
+        self.kv_allocated_len = 0
+        self.kv_committed_freed = False
+        self.kv_overallocated_freed = False
+
         # for corss-endoder model
         self.token_type_ids = token_type_ids
 
@@ -693,6 +700,35 @@ class Req:
             return self.output_ids[: self.finished_len]
         return self.output_ids
 
+    def pop_committed_kv_cache(self) -> int:
+        """Return the length of committed KV cache and mark them as freed."""
+
+        # NOTE: This function is called exactly once after the request is finished.
+        global_server_args = get_global_server_args()
+        topk = global_server_args.speculative_eagle_topk
+
+        enable_kv_committed_len = topk is None or topk == 1
+        if enable_kv_committed_len:
+            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
+        else:
+            return len(self.origin_input_ids) + max(len(self.output_ids) - 1, 0)
+
+    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_allocated_len=}"
+        self.kv_overallocated_freed = True
+        return self.kv_committed_len, self.kv_allocated_len
+
     def add_latency(self, stage: RequestStage):
         if self.metrics_collector is None:
             return
@@ -918,6 +954,10 @@ class Req:
         self.is_chunked = 0
         self.mamba_pool_idx = None
         self.already_computed = 0
+        self.kv_allocated_len = 0
+        self.kv_committed_len = 0
+        self.kv_committed_freed = False
+        self.kv_overallocated_freed = False
 
     def offload_kv_cache(self, req_to_token_pool, token_to_kv_pool_allocator):
         token_indices = req_to_token_pool.req_to_token[
@@ -1262,6 +1302,10 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
             req.req_pool_idx = req_pool_indices[i]
             assert seq_len - pre_len == req.extend_input_len
 
+            # update req-level memory management fields
+            req.kv_committed_len = seq_len
+            req.kv_allocated_len = seq_len
+
             # If input_embeds are available, store them
             if req.input_embeds is not None:
                 # If req.input_embeds is already a list, append its content directly
@@ -1536,7 +1580,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
                 self.req_to_token_pool, self.token_to_kv_pool_allocator
             )
         # TODO (csy): for preempted requests, we may want to insert into the tree
-        self.tree_cache.cache_finished_req(req, is_insert=False)
+        release_kv_cache(req, self.tree_cache, is_insert=False)
         # NOTE(lsyin): we should use the newly evictable memory instantly.
         num_tokens = remaing_req_count * envs.SGLANG_RETRACT_DECODE_STEPS.get()
         evict_from_tree_cache(self.tree_cache, num_tokens)
@@ -1614,6 +1658,11 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
         # Allocate memory
         self.out_cache_loc = alloc_for_decode(self, token_per_req=1)
 
+        # Update req-level memory management fields
+        for req in self.reqs:
+            req.kv_committed_len += 1
+            req.kv_allocated_len += 1
+
         # Update seq_lens after allocation
         if self.enable_overlap:
             # Do not use in-place operations in the overlap mode

python/sglang/srt/managers/scheduler.py

@@ -149,6 +149,7 @@ from sglang.srt.managers.scheduler_update_weights_mixin import (
 from sglang.srt.managers.session_controller import Session
 from sglang.srt.managers.utils import GenerationBatchResult, validate_input_length
 from sglang.srt.mem_cache.chunk_cache import ChunkCache, SWAChunkCache
+from sglang.srt.mem_cache.common import release_kv_cache
 from sglang.srt.mem_cache.hiradix_cache import HiRadixCache
 from sglang.srt.mem_cache.mamba_radix_cache import MambaRadixCache
 from sglang.srt.mem_cache.radix_cache import RadixCache
@@ -2410,11 +2411,11 @@ class Scheduler(
             self.send_to_tokenizer.send_output(AbortReq(rid=req.rid), req)
             # For disaggregation decode mode, the request in the waiting queue has KV cache allocated.
             if self.disaggregation_mode == DisaggregationMode.DECODE:
-                self.tree_cache.cache_finished_req(req)
+                release_kv_cache(req, self.tree_cache)
 
             # For mamba radix cache
             if req.mamba_pool_idx is not None:
-                self.tree_cache.cache_finished_req(req, is_insert=False)
+                release_kv_cache(req, self.tree_cache, is_insert=False)
             logger.debug(f"Abort queued request. {req.rid=}")
 
         # Delete the requests in the grammar queue

python/sglang/srt/managers/scheduler_output_processor_mixin.py

@@ -20,8 +20,8 @@ from sglang.srt.managers.schedule_batch import (
     RequestStage,
     ScheduleBatch,
 )
+from sglang.srt.mem_cache.common import release_kv_cache
 from sglang.srt.tracing.trace import trace_slice, trace_slice_batch, trace_slice_end
-from sglang.srt.utils.common import ceil_div
 
 if TYPE_CHECKING:
     from sglang.srt.managers.scheduler import (
@@ -55,7 +55,7 @@ class SchedulerOutputProcessorMixin:
                     req.rid,
                     thread_finish_flag=True,
                 )
-                self.tree_cache.cache_finished_req(req)
+                release_kv_cache(req, self.tree_cache)
 
         # Note: Logprobs should be handled on the prefill engine.
         trace_slice_batch(RequestStage.DECODE_FAKE_OUTPUT, batch.reqs)
@@ -102,26 +102,8 @@ class SchedulerOutputProcessorMixin:
             logprob_pt = 0
 
             for i, (req, next_token_id) in enumerate(zip(batch.reqs, next_token_ids)):
-                if self.enable_overlap and req.is_retracted and len(req.output_ids) > 0:
+                if req.finished() or req.is_retracted:
-                    req_idx = batch.req_pool_indices[i]
+                    # decode req in mixed batch or retracted req
-                    seq_len = len(req.origin_input_ids) + len(req.output_ids)
-                    pos = batch.req_to_token_pool.req_to_token[req_idx][
-                        seq_len - 1 : seq_len
-                    ]
-                    self.token_to_kv_pool_allocator.free(pos)
-                    continue
-
-                if (
-                    self.is_mixed_chunk
-                    and self.enable_overlap
-                    and (req.finished() or req.is_retracted)
-                ):
-                    # Free the one delayed token for the mixed decode batch
-                    j = len(batch.out_cache_loc) - len(batch.reqs) + i
-                    self.token_to_kv_pool_allocator.free(batch.out_cache_loc[j : j + 1])
-                    continue
-
-                if req.is_retracted:
                     continue
 
                 if req.is_chunked <= 0:
@@ -130,7 +112,7 @@ class SchedulerOutputProcessorMixin:
                     req.check_finished()
 
                     if req.finished():
-                        self.tree_cache.cache_finished_req(req)
+                        release_kv_cache(req, self.tree_cache)
                         req.time_stats.completion_time = time.perf_counter()
                     elif not batch.decoding_reqs or req not in batch.decoding_reqs:
                         # This updates radix so others can match
@@ -259,7 +241,7 @@ class SchedulerOutputProcessorMixin:
                     req.check_finished()
 
                     if req.finished():
-                        self.tree_cache.cache_finished_req(req)
+                        release_kv_cache(req, self.tree_cache)
                     else:
                         self.tree_cache.cache_unfinished_req(req)
                 else:
@@ -327,42 +309,17 @@ class SchedulerOutputProcessorMixin:
 
         self.token_to_kv_pool_allocator.free_group_begin()
 
+        # NOTE: in any case, we should check finish here
+        # if finished, also clean up committed kv cache and over-allocated kv cache here
+
         # Check finish condition
-        # NOTE: the length of reqs and next_token_ids don't match if it is spec decoding.
-        # We should ignore using next_token_ids for spec decoding cases.
         for i, (req, next_token_id) in enumerate(zip(batch.reqs, next_token_ids)):
             req: Req
 
             if self.enable_overlap and (req.finished() or req.is_retracted):
-                indices_to_free = None
+                # NOTE: This (req.finished() or req.is_retracted) should only happen when overlap scheduling is enabled.
-                if batch.spec_algorithm.is_eagle():
+                # (currently not, e.g. Eagle V1 still check finish during forward)
-                    from sglang.srt.speculative.eagle_info import EagleDraftInput
+                # And all the over-allocated tokens will be freed in `release_kv_cache`.
-
-                    end_p = allocate_lens_list[i]
-                    start_p = end_p - EagleDraftInput.ALLOC_LEN_PER_DECODE
-                    if self.page_size > 1:
-                        start_p = ceil_div(start_p, self.page_size) * self.page_size
-
-                    indices_to_free = self.req_to_token_pool.req_to_token[
-                        req.req_pool_idx
-                    ][start_p:end_p]
-
-                else:
-                    if self.page_size == 1:
-                        # Free the one extra delayed token
-                        indices_to_free = batch.out_cache_loc[i : i + 1]
-                    else:
-                        if (
-                            len(req.origin_input_ids) + len(req.output_ids) - 1
-                        ) % self.page_size == 0:
-                            # Only free when the extra token is in a new page
-                            indices_to_free = batch.out_cache_loc[i : i + 1]
-
-                if indices_to_free is not None:
-                    self.token_to_kv_pool_allocator.free(indices_to_free)
-                continue
-
-            if req.is_retracted:
                 continue
 
             new_accepted_len = 1
@@ -376,27 +333,12 @@ class SchedulerOutputProcessorMixin:
             req.check_finished(new_accepted_len)
 
             if req.finished():
-                if batch.is_v2_eagle and self.cur_batch.forward_mode.is_extend():
-                    # FIXME(lsyin): fix the messy logic here
-                    # 1) when not overlap (v2 impl), we free the extra tokens in the req
-                    # 2) overlap eagle and the current batch is prefill. This seq will not run extra iteration.
-                    start_p = batch.seq_lens_cpu[i] + accept_lens_list[i]
-                    end_p = allocate_lens_list[i]
-
-                    if self.page_size > 1:
-                        start_p = ceil_div(start_p, self.page_size) * self.page_size
-
-                    indices_to_free = self.req_to_token_pool.req_to_token[
-                        req.req_pool_idx
-                    ][start_p:end_p]
-                    self.token_to_kv_pool_allocator.free(indices_to_free)
-
                 if self.server_args.disaggregation_decode_enable_offload_kvcache:
-                    # Asynchronously offload KV cache; cache_finished_req will be called after Device->Host transfer completes
+                    # Asynchronously offload KV cache; release_kv_cache will be called after Device->Host transfer completes
                     if not self.decode_offload_manager.offload_kv_cache(req):
-                        self.tree_cache.cache_finished_req(req)
+                        release_kv_cache(req, self.tree_cache)
                 else:
-                    self.tree_cache.cache_finished_req(req)
+                    release_kv_cache(req, self.tree_cache)
 
                 req.time_stats.completion_time = time.perf_counter()
 

python/sglang/srt/mem_cache/chunk_cache.py

@@ -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,
+            req.req_pool_idx, :kv_committed_len
-            # 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),
         ]
         self.req_to_token_pool.free(req.req_pool_idx)
         self.token_to_kv_pool_allocator.free(kv_indices)

python/sglang/srt/mem_cache/common.py

@@ -14,6 +14,7 @@ from sglang.srt.mem_cache.mamba_radix_cache import MambaRadixCache
 from sglang.srt.mem_cache.memory_pool import HybridReqToTokenPool, ReqToTokenPool
 from sglang.srt.server_args import get_global_server_args
 from sglang.srt.utils import support_triton
+from sglang.srt.utils.common import ceil_align
 
 if TYPE_CHECKING:
     from sglang.srt.managers.schedule_batch import Req, ScheduleBatch
@@ -462,6 +463,31 @@ def alloc_for_decode(batch: ScheduleBatch, token_per_req: int) -> torch.Tensor:
     return out_cache_loc
 
 
+def release_kv_cache(req: Req, tree_cache: BasePrefixCache, is_insert: bool = True):
+    tree_cache.cache_finished_req(req, is_insert=is_insert)
+    start_p, end_p = req.pop_overallocated_kv_cache()
+
+    global_server_args = get_global_server_args()
+    page_size = global_server_args.page_size
+    spec_algo = global_server_args.speculative_algorithm
+
+    if spec_algo is None:
+        assert (
+            start_p == end_p
+        ), f"Unexpected overallocated KV cache, {req.kv_committed_len=}, {req.kv_allocated_len=}"
+
+    if page_size > 1:
+        start_p = ceil_align(start_p, page_size)
+
+    if start_p >= end_p:
+        return
+
+    indices_to_free = tree_cache.req_to_token_pool.req_to_token[req.req_pool_idx][
+        start_p:end_p
+    ]
+    tree_cache.token_to_kv_pool_allocator.free(indices_to_free)
+
+
 def available_and_evictable_str(tree_cache) -> str:
     token_to_kv_pool_allocator = tree_cache.token_to_kv_pool_allocator
     if isinstance(token_to_kv_pool_allocator, SWATokenToKVPoolAllocator):

python/sglang/srt/mem_cache/mamba_radix_cache.py

@@ -430,21 +430,21 @@ class MambaRadixCache(BasePrefixCache):
             value = torch.tensor([x for x in key.token_ids], dtype=torch.int64)
         return self._insert_helper(self.root_node, key, value, mamba_value)
 
-    def cache_finished_req(self, req: Req, is_insert=True) -> None:
+    def cache_finished_req(self, req: Req, is_insert: bool = True):
         """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,
+                req.req_pool_idx, :kv_committed_len
-                : len(req.origin_input_ids) + max(len(req.output_ids) - 1, 0),
             ]
             self.token_to_kv_pool_allocator.free(kv_indices)
             self.req_to_token_pool.free(req.req_pool_idx)
             return
 
-        cache_len = len(req.origin_input_ids) + max(len(req.output_ids) - 1, 0)
+        token_ids = (req.origin_input_ids + req.output_ids)[:kv_committed_len]
-        token_ids = (req.origin_input_ids + req.output_ids)[:cache_len]
         kv_indices = self.req_to_token_pool.req_to_token[
-            req.req_pool_idx, : len(token_ids)
+            req.req_pool_idx, :kv_committed_len
         ]
 
         page_aligned_len = len(kv_indices)

python/sglang/srt/mem_cache/radix_cache.py

@@ -341,21 +341,21 @@ class RadixCache(BasePrefixCache):
 
     def cache_finished_req(self, req: Req, is_insert: bool = True):
         """Cache request when it finishes."""
-        all_token_len = len(req.origin_input_ids) + max(len(req.output_ids) - 1, 0)
+        committed_kv_len = req.pop_committed_kv_cache()
         if self.disable:
             kv_indices = self.req_to_token_pool.req_to_token[
-                req.req_pool_idx, :all_token_len
+                req.req_pool_idx, :committed_kv_len
             ]
             self.token_to_kv_pool_allocator.free(kv_indices)
             self.req_to_token_pool.free(req.req_pool_idx)
             return
 
-        token_ids = (req.origin_input_ids + req.output_ids)[:all_token_len]
+        token_ids = (req.origin_input_ids + req.output_ids)[:committed_kv_len]
         # For EAGLE radix cache, we will convert the key to bigram key, e.g. [1,2,3,4] -> [(1,2), (2,3), (3,4)], the length will -1. ((len([(1,2), (2,3), (3,4)]) = len([1,2,3,4]) - 1))
         # So for the corresponding kv length should also -1. Then we get the actual_kv_len, and use it to do later calculation and slicing.
-        actual_kv_len = all_token_len - 1 if self.is_eagle else all_token_len
+        actual_kv_len = committed_kv_len - 1 if self.is_eagle else committed_kv_len
         kv_indices = self.req_to_token_pool.req_to_token[
-            req.req_pool_idx, :all_token_len
+            req.req_pool_idx, :committed_kv_len
         ]
 
         if self.page_size != 1:

python/sglang/srt/mem_cache/radix_cache_cpp.py

@@ -165,15 +165,16 @@ class RadixCacheCpp(BasePrefixCache):
     def cache_finished_req(self, req: Req, is_insert: bool = True):
         """Cache request when it finishes."""
         assert req.req_pool_idx is not None
-        all_token_len = len(req.origin_input_ids) + max(len(req.output_ids) - 1, 0)
+        kv_committed_len = req.pop_committed_kv_cache()
-        token_ids = (req.origin_input_ids + req.output_ids)[:all_token_len]
+        token_ids = (req.origin_input_ids + req.output_ids)[:kv_committed_len]
-        overall_len = len(token_ids)  # prefill + decode
+        kv_indices = self.req_to_token_pool.req_to_token[
-        kv_indices = self.req_to_token_pool.req_to_token[req.req_pool_idx, :overall_len]
+            req.req_pool_idx, :kv_committed_len
+        ]
 
         # 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 = overall_len // self.page_size * self.page_size
+        page_aligned_overall_len = kv_committed_len // self.page_size * self.page_size
 
         if is_insert:
             new_prefix_len = self._insert(
@@ -190,7 +191,7 @@ class RadixCacheCpp(BasePrefixCache):
             )
 
         # need to free the unaligned part, since it cannot be inserted into the radix tree
-        if page_aligned_overall_len < overall_len:
+        if page_aligned_overall_len < kv_committed_len:
             # NOTE: sglang PagedAllocator support unaligned free (which will automatically align it)
             self.token_to_kv_pool.free(kv_indices[page_aligned_overall_len:])
 

python/sglang/srt/mem_cache/storage/lmcache/lmc_radix_cache.py

@@ -219,16 +219,17 @@ class LMCRadixCache(RadixCache):
 
         return base_res
 
-    def cache_finished_req(self, req: "Req", is_insert: bool = True) -> None:  # type: ignore[override]
+    def cache_finished_req(self, req: Req, is_insert: bool = True) -> None:  # type: ignore[override]
         """On request completion, insert device KV into radix and store to LMCache."""
 
         super().cache_finished_req(req, is_insert=is_insert)
         if not is_insert:
             return
 
-        token_ids = (req.origin_input_ids + req.output_ids)[:-1]
+        kv_committed_len = req.pop_committed_kv_cache()
+        token_ids = (req.origin_input_ids + req.output_ids)[:kv_committed_len]
         kv_indices = self.req_to_token_pool.req_to_token[
-            req.req_pool_idx, : len(token_ids)
+            req.req_pool_idx, :kv_committed_len
         ]
 
         _, new_last_node, _, _ = self.match_prefix(RadixKey(token_ids, req.extra_key))

python/sglang/srt/mem_cache/swa_radix_cache.py

@@ -441,21 +441,21 @@ class SWARadixCache(BasePrefixCache):
 
     def cache_finished_req(self, req: Req, is_insert: bool = True) -> None:
         """Cache request when it finishes."""
-        all_token_len = len(req.origin_input_ids) + max(len(req.output_ids) - 1, 0)
+        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, :all_token_len
+                req.req_pool_idx, :kv_committed_len
             ]
             self.token_to_kv_pool_allocator.free(kv_indices)
             self.req_to_token_pool.free(req.req_pool_idx)
             return
 
-        token_ids = (req.origin_input_ids + req.output_ids)[:all_token_len]
+        token_ids = (req.origin_input_ids + req.output_ids)[:kv_committed_len]
         # For EAGLE radix cache, we will convert the key to bigram key, e.g. [1,2,3,4] -> [(1,2), (2,3), (3,4)], the length will -1. ((len([(1,2), (2,3), (3,4)]) = len([1,2,3,4]) - 1))
         # So for the corresponding kv length should also -1. Then we get the actual_kv_len, and use it to do later calculation and slicing.
-        actual_kv_len = all_token_len - 1 if self.is_eagle else all_token_len
+        actual_kv_len = kv_committed_len - 1 if self.is_eagle else kv_committed_len
         kv_indices = self.req_to_token_pool.req_to_token[
-            req.req_pool_idx, :all_token_len
+            req.req_pool_idx, :kv_committed_len
         ]
 
         if self.page_size != 1:

python/sglang/srt/speculative/eagle_info.py

@@ -116,6 +116,8 @@ class EagleVerifyInput(SpecInput, EagleVerifyInputV2Mixin):
                 len(batch.input_ids),
             )
             end_offset = batch.seq_lens + self.draft_token_num
+            for req in batch.reqs:
+                req.kv_allocated_len += 1
         else:
             prefix_lens = batch.seq_lens
             prefix_lens_cpu = batch.seq_lens_cpu
@@ -415,6 +417,9 @@ class EagleVerifyInput(SpecInput, EagleVerifyInputV2Mixin):
         if page_size == 1:
             # TODO: boolean array index leads to a device sync. Remove it.
             token_to_kv_pool_allocator.free(batch.out_cache_loc[evict_mask])
+            for i, req in enumerate(batch.reqs):
+                req.kv_committed_len += accept_length_list[i] + 1
+                req.kv_allocated_len = req.kv_committed_len
         else:
             if self.topk == 1:
                 # Only evict full empty page. Do not evict partial empty page
@@ -426,6 +431,9 @@ class EagleVerifyInput(SpecInput, EagleVerifyInputV2Mixin):
                     next_power_of_2(self.draft_token_num),
                 )
                 token_to_kv_pool_allocator.free(batch.out_cache_loc[evict_mask])
+                for i, req in enumerate(batch.reqs):
+                    req.kv_committed_len += accept_length_list[i] + 1
+                    req.kv_allocated_len = req.kv_committed_len
             else:
                 # Shift the accepted tokens to the beginning.
                 # Only evict the last part

python/sglang/srt/speculative/eagle_info_v2.py

@@ -129,6 +129,10 @@ class EagleDraftInputV2Mixin:
         batch.seq_lens_cpu = batch.seq_lens.cpu()
         batch.seq_lens_sum = batch.seq_lens_cpu.sum().item()
 
+        for i, req in enumerate(batch.reqs):
+            req.kv_committed_len = batch.seq_lens_cpu[i].item()
+            req.kv_allocated_len = req.kv_committed_len + self.ALLOC_LEN_PER_DECODE
+
     def prepare_for_v2_draft(
         self: EagleDraftInput,
         req_to_token_pool: ReqToTokenPool,

python/sglang/srt/speculative/eagle_worker.py

@@ -364,6 +364,8 @@ class EAGLEWorker(TpModelWorker):
         # [       topk 0         ] [       topk 1         ]
         # [iter=0, iter=1, iter=2] [iter=0, iter=1, iter=2]
         if self.page_size == 1:
+            for req in batch.reqs:
+                req.kv_allocated_len += self.speculative_num_steps * self.topk
             out_cache_loc, token_to_kv_pool_state_backup = alloc_token_slots(
                 batch.tree_cache,
                 num_seqs * self.speculative_num_steps * self.topk,

python/sglang/srt/speculative/ngram_info.py

@@ -195,7 +195,9 @@ class NgramVerifyInput(SpecInput):
             logits_output.hidden_states = logits_output.hidden_states[self.accept_index]
         self.verified_id = self.predict[self.accept_index]
 
-    def _free_cache(self, batch: ScheduleBatch, page_size: int):
+    def _free_cache(
+        self, batch: ScheduleBatch, page_size: int, accept_length_cpu: torch.Tensor
+    ):
         bs = batch.batch_size()
         # Free the KV cache for unaccepted tokens
         if page_size == 1:
@@ -250,6 +252,11 @@ class NgramVerifyInput(SpecInput):
             )
             batch.out_cache_loc = tgt_cache_loc
 
+        accept_length_list = accept_length_cpu.tolist()
+        for i, req in enumerate(batch.reqs):
+            req.kv_committed_len += accept_length_list[i] + 1
+            req.kv_allocated_len = req.kv_committed_len
+
         assign_req_to_token_pool[(bs,)](
             batch.req_pool_indices,
             batch.req_to_token_pool.req_to_token,
@@ -416,11 +423,12 @@ class NgramVerifyInput(SpecInput):
             # self._sampling_verify(batch, logits_output, sampling_info)
 
         self._fill_requests(batch, logits_output)
-        self._free_cache(batch, page_size)
 
         accept_length_cpu = self.accept_length.cpu()
         num_accepted_tokens = accept_length_cpu.sum().item()
 
+        self._free_cache(batch, page_size, accept_length_cpu)
+
         batch.seq_lens.add_(self.accept_length + 1)
         batch.seq_lens_cpu.add_(accept_length_cpu + 1)
 

python/sglang/srt/utils/common.py

@@ -2981,7 +2981,7 @@ def configure_gc_logger():
 
 
 # COPIED FROM DeepGEMM
-def align(x: int, y: int) -> int:
+def ceil_align(x: int, y: int) -> int:
     return ceil_div(x, y) * y
 
 

test/srt/hicache/test_hicache_storage_mooncake_backend.py

@@ -4,18 +4,14 @@ Usage:
     python3.10 -m pytest test/srt/hicache/test_hicache_storage_mooncake_backend.py -v
 """
 
-import json
 import os
 import subprocess
 import time
 import unittest
-from types import SimpleNamespace
 
 import requests
 from test_hicache_storage_file_backend import HiCacheStorageBaseMixin
 
-from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k
-from sglang.test.test_disaggregation_utils import get_rdma_devices_args
 from sglang.test.test_utils import (
     DEFAULT_MLA_MODEL_NAME_FOR_TEST,
     CustomTestCase,

test/srt/test_eagle_infer_beta.py

@@ -68,6 +68,7 @@ class TestEagleServerBase(CustomTestCase, MatchedStopMixin):
 
     def test_radix_attention(self):
         run_radix_attention_test(self.base_url)
+        assert self.process.poll() is None
 
     def test_gsm8k(self):
         args = SimpleNamespace(
@@ -84,6 +85,11 @@ class TestEagleServerBase(CustomTestCase, MatchedStopMixin):
         self.assertGreater(
             metrics["accuracy"], 0.23
         )  # 0.3333 for 60 questions; 0.234 for 1319 questions
+        assert self.process.poll() is None
+
+
+class TestEagleServerPage(TestEagleServerBase):
+    other_launch_args = ["--page-size", "64"]
 
 
 if __name__ == "__main__":