minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[mem_cache] Add free_kv_row to release a request's kv row by row range (#36721)

Browse Source
This commit is contained in:
Liangsheng Yin
2026-09-01 01:14:43 -07:00
committed by GitHub
parent 1c3ad92438
commit 3484f7f836
12 changed files with 323 additions and 84 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py
+8 -15
View File
@@ -250,27 +250,20 @@ class DecodeKVCacheOffloadManager:
# Prefill-aligned slots are freed only here, at request finish; freeing
# them mid-decode races with concurrent admission over live slots.
prefill_len = self._prefill_offloaded_len(req)
ranges = []
if prefill_len > 0:
prefill_indices = self.req_to_token_pool.req_to_token[
req.kv.req_pool_idx, :prefill_len
]
self.token_to_kv_pool_allocator.free(prefill_indices)
start = prefill_len
end = kv_committed_len
# Free the incremental part of the request (DSA-aware)
kv_indices = self.req_to_token_pool.req_to_token[req.kv.req_pool_idx, start:end]
self.token_to_kv_pool_allocator.free(kv_indices)
ranges.append((0, prefill_len))
# The incremental part of the request (DSA-aware)
ranges.append((prefill_len, kv_committed_len))
# 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.
# Over-allocated KV cache slots (e.g. from speculative decoding v2).
# Without spec v2, start_p == end_p so this contributes nothing.
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:
overalloc_indices = self.req_to_token_pool.req_to_token[
req.kv.req_pool_idx, start_p:end_p
]
self.token_to_kv_pool_allocator.free(overalloc_indices)
ranges.append((start_p, end_p))
self.tree_cache.free_kv_row(req.kv, ranges)
self.req_to_token_pool.free(req)
req.kv.mark_kv_released()
python/sglang/srt/mem_cache/base_prefix_cache.py
+13
View File
@@ -369,6 +369,19 @@ class BasePrefixCache(ABC, PrefixCacheTrait):
def cache_unfinished_req(self, req: Req, **kwargs):
pass
def free_kv_row(self, kv: Any, ranges: list[tuple[int, int]]) -> None:
"""Give back ascending, disjoint, half-open row-position ranges
of the ``kv`` record's row; one call keeps a shared page freed once.
"""
from sglang.srt.mem_cache.common import free_kv_row_segments
row = self.req_to_token_pool.req_to_token[kv.req_pool_idx]
free_kv_row_segments(
self.token_to_kv_pool_allocator,
[(row[start:end], start) for start, end in ranges],
swa_evicted_seqlen=kv.swa_evicted_seqlen,
)
@abstractmethod
def evict(self, params: EvictParams) -> EvictResult:
pass
python/sglang/srt/mem_cache/chunk_cache.py
+1 -4
View File
@@ -81,10 +81,7 @@ class ChunkCache(BasePrefixCache):
):
# For decode server: if req.output_ids is empty, we want to free all req.origin_input_ids
# The protected prefix is not this req's to free.
kv_indices = self.req_to_token_pool.req_to_token[
req.kv.req_pool_idx, req.kv.cache_protected_len : kv_len_to_handle
]
self.token_to_kv_pool_allocator.free(kv_indices)
self.free_kv_row(req.kv, [(req.kv.cache_protected_len, kv_len_to_handle)])
def cache_unfinished_req(self, req: Req, chunked=False):
kv_indices = self.req_to_token_pool.req_to_token[
python/sglang/srt/mem_cache/common.py
+40 -4
View File
@@ -119,6 +119,45 @@ def free_swa_out_of_window_slots(
req.kv.swa_evicted_seqlen = new_swa_evicted_seqlen
def free_kv_row_segments(
allocator: BaseTokenToKVPoolAllocator,
segments: list[tuple[torch.Tensor, int]],
*,
swa_evicted_seqlen: int,
) -> None:
"""Free ascending disjoint ``(kv_indices, start_pos)`` segments of one
request's kv row, split at the SWA eviction floor."""
swa_dead: list[torch.Tensor] = []
swa_alive: list[tuple[torch.Tensor, int]] = []
for kv_indices, start_pos in segments:
num_indices = kv_indices.numel()
if num_indices == 0:
continue
# Below the floor the SWA peers are already gone -- window eviction, or
# the deliberately unmapped prefix of a PD decode SWA-tail prealloc.
num_dead = min(max(swa_evicted_seqlen - start_pos, 0), num_indices)
if num_dead > 0:
swa_dead.append(kv_indices[:num_dead])
if num_dead < num_indices:
swa_alive.append((kv_indices[num_dead:], start_pos + num_dead))
if swa_dead and swa_alive:
# A mid-page floor would send a page shared by the dead and alive
# sides back twice.
assert swa_evicted_seqlen % allocator.page_size == 0, (
f"SWA eviction floor {swa_evicted_seqlen} splits a page "
f"(page_size {allocator.page_size})"
)
if len(swa_dead) == 1:
allocator.free_full(swa_dead[0])
elif swa_dead:
# Two dead pieces can share a boundary page, and only free_full's own
# page dedup covers that -- free_segments trims the alive side alone.
allocator.free_full(torch.cat(swa_dead))
if swa_alive:
allocator.free_segments(swa_alive)
def maybe_cache_unfinished_req(req: Req, tree_cache: BasePrefixCache, **kwargs):
if getattr(req, "skip_radix_cache_insert", False):
return
@@ -275,12 +314,9 @@ def _release_overallocated_kv_indices(
start_p = ceil_align(start_p, page_size)
if start_p < end_p:
indices_to_free = tree_cache.req_to_token_pool.req_to_token[
req.kv.req_pool_idx
][start_p:end_p]
# start_p is aligned to the allocator's physical page size above, so it
# never shares a page with cache_finished_req's tail free in this group.
allocator.free_segment(indices_to_free, start_pos=start_p)
tree_cache.free_kv_row(req.kv, [(start_p, end_p)])
def available_and_evictable_str(tree_cache: BasePrefixCache) -> str:
python/sglang/srt/mem_cache/swa_radix_cache.py
+17 -15
View File
@@ -42,6 +42,7 @@ from sglang.srt.mem_cache.base_prefix_cache import (
MatchResult,
)
from sglang.srt.mem_cache.cache_init_params import CacheInitParams
from sglang.srt.mem_cache.common import free_kv_row_segments
from sglang.srt.mem_cache.events import KVCacheEventRecorder
from sglang.srt.mem_cache.radix_cache import RadixKey
from sglang.srt.mem_cache.utils import split_node_hash_value
@@ -464,10 +465,7 @@ class SWARadixCache(BasePrefixCache):
) -> None:
"""Cache request when it finishes."""
if self.disable:
kv_indices = self.req_to_token_pool.req_to_token[
req.kv.req_pool_idx, :kv_len_to_handle
]
self.token_to_kv_pool_allocator.free(kv_indices)
self.free_kv_row(req.kv, [(0, kv_len_to_handle)])
return
token_ids = (req.origin_input_ids + req.output_ids)[:kv_len_to_handle]
@@ -497,12 +495,10 @@ class SWARadixCache(BasePrefixCache):
)
)
else:
self.token_to_kv_pool_allocator.free(
kv_indices[old_prefix_len:page_aligned_len]
)
self.free_kv_row(req.kv, [(old_prefix_len, page_aligned_len)])
# free the unaligned tail
self.token_to_kv_pool_allocator.free(kv_indices[page_aligned_len:])
self.free_kv_row(req.kv, [(page_aligned_len, kv_len_to_handle)])
# Remove req slot release the cache lock
self.dec_lock_ref(
@@ -1207,7 +1203,7 @@ class SWARadixCache(BasePrefixCache):
)
else:
# Free full tokens in the original tree node.
self.token_to_kv_pool_allocator.free(
self.token_to_kv_pool_allocator.free_full(
node.value[:prefix_len]
)
# Overwrite the new value in request to the tree node.
@@ -1225,18 +1221,18 @@ class SWARadixCache(BasePrefixCache):
self._recover_tombstone_keeping_locked_full(
node, value[start_update_idx:prefix_len]
)
self.token_to_kv_pool_allocator.free(
self.token_to_kv_pool_allocator.free_full(
value[:start_update_idx]
)
else:
self.token_to_kv_pool_allocator.free(
self.token_to_kv_pool_allocator.free_full(
node.value[start_update_idx:prefix_len]
)
self._split_node(node.key, node, start_update_idx)
# Here node is the new node after split, so we can overwrite the value to the new node.
# The old node is still swa tombstone and the full token is not freed.
node.value = value[start_update_idx:prefix_len].clone()
self.token_to_kv_pool_allocator.free(
self.token_to_kv_pool_allocator.free_full(
value[:start_update_idx]
)
node.swa_tombstone = False
@@ -1244,10 +1240,16 @@ class SWARadixCache(BasePrefixCache):
self.swa_evictable_size_ += len(node.value)
else:
# Branch 3: all swa tokens of value[:prefix_len] are evicted, so we don't need to update the node.
self.token_to_kv_pool_allocator.free(value[:prefix_len])
self.token_to_kv_pool_allocator.free_full(value[:prefix_len])
else:
# The node is not tombstone, so we don't need to update the node.
self.token_to_kv_pool_allocator.free(value[:prefix_len])
# The incoming slice can still straddle this request's own
# eviction floor, so split it there.
free_kv_row_segments(
self.token_to_kv_pool_allocator,
[(value[:prefix_len], total_prefix_length)],
swa_evicted_seqlen=swa_evicted_seqlen,
)
total_prefix_length += prefix_len
key = key[prefix_len:]
@@ -1274,7 +1276,7 @@ class SWARadixCache(BasePrefixCache):
# occurring in normal operation. This check is a defensive guard
# against unexpected eviction states from other code paths.
if swa_evicted_seqlen == total_prefix_length + len(key):
self.token_to_kv_pool_allocator.free(value)
self.token_to_kv_pool_allocator.free_full(value)
return total_prefix_length
if (
python/sglang/srt/mem_cache/unified_radix_cache.py
+5 -11
View File
@@ -842,10 +842,7 @@ class UnifiedRadixCache(BasePrefixCache):
return
if self.disable:
kv_indices = self.req_to_token_pool.req_to_token[
req.kv.req_pool_idx, :kv_len_to_handle
]
self.token_to_kv_pool_allocator.free_segment(kv_indices, start_pos=0)
self.free_kv_row(req.kv, [(0, kv_len_to_handle)])
for comp in self._components_tuple:
comp.cleanup_after_caching_req(req, is_finished=True)
return
@@ -899,15 +896,12 @@ class UnifiedRadixCache(BasePrefixCache):
result = self.insert(insert_params)
# Free unaligned tail (+ deferred truncation tail)
segments = [(kv_indices[page_aligned_len:], page_aligned_len)]
ranges = [(page_aligned_len, len(kv_indices))]
if tail_free_start is not None:
segments.append((kv_indices_full[tail_free_start:], tail_free_start))
self.token_to_kv_pool_allocator.free_segments(segments)
ranges.append((tail_free_start, len(kv_indices_full)))
self.free_kv_row(req.kv, ranges)
else:
self.token_to_kv_pool_allocator.free_segment(
kv_indices[req.kv.cache_protected_len :],
start_pos=req.kv.cache_protected_len,
)
self.free_kv_row(req.kv, [(req.kv.cache_protected_len, kv_len_to_handle)])
self._dec_req_lock(req, skip_swa=req.swa_prefix_lock_released)
python/sglang/srt/session/streaming_session.py
+14 -16
View File
@@ -232,8 +232,10 @@ class StreamingSession(BasePrefixCache):
f"{slot.kv.cache_protected_len=}"
)
# Floor-align prefix_len to page boundary (NPU workaround).
if is_npu() and self.page_size > 1:
# NPU requires page-aligned KV reuse; a rewind below the SWA eviction
# cursor must also land on a page boundary -- free_kv_row_segments
# splits dead/alive at the cursor, and a mid-page cut frees a page twice.
if self.page_size > 1 and (is_npu() or req.kv.swa_evicted_seqlen > prefix_len):
prefix_len = (prefix_len // self.page_size) * self.page_size
req.kv.kv_committed_len = min(req.kv.kv_committed_len, prefix_len)
@@ -401,13 +403,7 @@ class StreamingSession(BasePrefixCache):
)
if slot.kv.holds_kv:
start = protected_len
end = slot.kv.kv_allocated_len
if start < end:
kv_indices = self.req_to_token_pool.req_to_token[
slot.kv.req_pool_idx, start:end
]
self.token_to_kv_pool_allocator.free(kv_indices)
self.free_kv_row(slot.kv, [(protected_len, slot.kv.kv_allocated_len)])
self.req_to_token_pool.free(slot)
self._free_slot_mamba(slot)
@@ -504,7 +500,7 @@ class StreamingSession(BasePrefixCache):
decoding pushes allocated above committed, or when retract retry's
logit-reserve pulls prefix_len below committed.
"""
self._free_kv_aligned(kv.req_pool_idx, prefix_len, kv.kv_allocated_len)
self._free_kv_aligned(kv, prefix_len, kv.kv_allocated_len)
kv.kv_allocated_len = prefix_len
kv.kv_committed_len = min(kv.kv_committed_len, prefix_len)
kv.swa_evicted_seqlen = min(kv.swa_evicted_seqlen, prefix_len)
@@ -516,14 +512,18 @@ class StreamingSession(BasePrefixCache):
be released to avoid token/KV mismatch.
"""
target = len(req.origin_input_ids) + finished_len
self._free_kv_aligned(req.kv.req_pool_idx, target, req.kv.kv_allocated_len)
if self.page_size > 1 and req.kv.swa_evicted_seqlen > target:
# Same hazard as the match-path rewind: the cursor must stay
# page-aligned; the partial page is re-prefilled next turn.
target = (target // self.page_size) * self.page_size
self._free_kv_aligned(req.kv, target, req.kv.kv_allocated_len)
req.kv.kv_allocated_len = min(req.kv.kv_allocated_len, target)
req.kv.kv_committed_len = min(req.kv.kv_committed_len, target)
req.kv.swa_evicted_seqlen = min(req.kv.swa_evicted_seqlen, target)
req.output_ids = req.output_ids[:finished_len]
def _free_kv_aligned(self, pool_idx: int, target: int, end: int) -> None:
"""Free req_to_token[pool_idx, ceil_align(target):end). Page-aligned
def _free_kv_aligned(self, kv: ReqKvInfo, target: int, end: int) -> None:
"""Free the record's kv row over [ceil_align(target), end). Page-aligned
because PagedTokenToKVPoolAllocator.free returns whole pages
(free_index // page_size), so partial-page free would corrupt pages
still holding committed tokens. The range [target, ceil_align(target))
@@ -534,9 +534,7 @@ class StreamingSession(BasePrefixCache):
start = target
if self.page_size > 1:
start = ceil_align(start, self.page_size)
if start < end:
tail = self.req_to_token_pool.req_to_token[pool_idx, start:end]
self.token_to_kv_pool_allocator.free(tail)
self.free_kv_row(kv, [(start, end)])
# -- Pass-through methods --