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

Let cache backend do not couple with owned committed kv details and avoid kv_committed_freed/kv_overallocated_freed fields (#29428)

Browse Source
This commit is contained in:
fzyzcjy
2026-07-15 14:43:10 +08:00
committed by GitHub
parent d8d76c4d12
commit 27256aee5b
24 changed files with 121 additions and 154 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py
+2 -2
View File
@@ -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:
python/sglang/srt/managers/schedule_batch.py
+1 -25
View File
@@ -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
python/sglang/srt/managers/scheduler.py
+1 -3
View File
@@ -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
python/sglang/srt/managers/scheduler_components/invariant_checker.py
+1 -2
View File
@@ -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
python/sglang/srt/mem_cache/chunk_cache.py
+8 -5
View File
@@ -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
]
python/sglang/srt/mem_cache/common.py
+4 -2
View File
@@ -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
python/sglang/srt/mem_cache/mamba_radix_cache.py
+7 -6
View File
@@ -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
python/sglang/srt/mem_cache/pure_swa_radix_cache.py
+4 -2
View File
@@ -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
python/sglang/srt/mem_cache/radix_cache.py
+5 -4
View File
@@ -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)
]
python/sglang/srt/mem_cache/radix_cache_cpp.py
+7 -6
View File
@@ -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:])
python/sglang/srt/mem_cache/storage/flexkv/flexkv_radix_cache.py
+4 -2
View File
@@ -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
python/sglang/srt/mem_cache/storage/lmcache/lmc_radix_cache.py
+6 -2
View File
@@ -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)
python/sglang/srt/mem_cache/swa_radix_cache.py
+6 -5
View File
@@ -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(
python/sglang/srt/mem_cache/unified_cache_components/README.md
+2 -2
View File
@@ -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)**. |
python/sglang/srt/mem_cache/unified_radix_cache.py
+6 -6
View File
@@ -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
python/sglang/srt/session/streaming_session.py
-10
View File
@@ -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):
test/registered/disaggregation/test_specv2_kvcache_offloading.py
+1 -15
View File
@@ -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
test/registered/unit/mem_cache/test_decode_radix_lock_ref.py
+9 -13
View File
@@ -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)
test/registered/unit/mem_cache/test_pure_swa_chunk_cache.py
+1 -1
View File
@@ -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]
test/registered/unit/mem_cache/test_streaming_session_unit.py
-20
View File
@@ -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
test/registered/unit/mem_cache/test_swa_eviction_boundary.py
+24 -9
View File
@@ -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()
test/registered/unit/mem_cache/test_swa_unittest.py
+6 -5
View File
@@ -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:
test/registered/unit/mem_cache/test_unified_radix_cache_bench.py
+3 -2
View File
@@ -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,
test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py
+13 -5
View File
@@ -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]