From 3484f7f8361f8bc32177ce3f1073cfc3a2988d66 Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Tue, 1 Sep 2026 01:14:43 -0700 Subject: [PATCH] [mem_cache] Add `free_kv_row` to release a request's kv row by row range (#36721) --- .../decode_kvcache_offload_manager.py | 23 ++-- .../sglang/srt/mem_cache/base_prefix_cache.py | 13 ++ python/sglang/srt/mem_cache/chunk_cache.py | 5 +- python/sglang/srt/mem_cache/common.py | 44 ++++++- .../sglang/srt/mem_cache/swa_radix_cache.py | 32 ++--- .../srt/mem_cache/unified_radix_cache.py | 16 +-- .../sglang/srt/session/streaming_session.py | 30 ++--- .../test_specv2_kvcache_offloading.py | 40 +++++- .../unit/layers/test_minicpm_sparse_cache.py | 19 ++- .../mem_cache/test_decode_radix_lock_ref.py | 1 + .../mem_cache/test_streaming_session_unit.py | 62 ++++++++- .../unit/mem_cache/test_swa_unittest.py | 122 +++++++++++++++++- 12 files changed, 323 insertions(+), 84 deletions(-) diff --git a/python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py b/python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py index 354ae4039..eed819e08 100644 --- a/python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py +++ b/python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py @@ -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() diff --git a/python/sglang/srt/mem_cache/base_prefix_cache.py b/python/sglang/srt/mem_cache/base_prefix_cache.py index 8338aeffd..2998b18e4 100644 --- a/python/sglang/srt/mem_cache/base_prefix_cache.py +++ b/python/sglang/srt/mem_cache/base_prefix_cache.py @@ -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 diff --git a/python/sglang/srt/mem_cache/chunk_cache.py b/python/sglang/srt/mem_cache/chunk_cache.py index cdc79ce49..c524a974e 100644 --- a/python/sglang/srt/mem_cache/chunk_cache.py +++ b/python/sglang/srt/mem_cache/chunk_cache.py @@ -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[ diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index 0915b5bb3..b1497c6f5 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/python/sglang/srt/mem_cache/common.py @@ -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: diff --git a/python/sglang/srt/mem_cache/swa_radix_cache.py b/python/sglang/srt/mem_cache/swa_radix_cache.py index 86374c907..5b6862d48 100644 --- a/python/sglang/srt/mem_cache/swa_radix_cache.py +++ b/python/sglang/srt/mem_cache/swa_radix_cache.py @@ -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 ( diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index 925eac13f..5cb53fac7 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -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) diff --git a/python/sglang/srt/session/streaming_session.py b/python/sglang/srt/session/streaming_session.py index be01edf68..da18238ea 100644 --- a/python/sglang/srt/session/streaming_session.py +++ b/python/sglang/srt/session/streaming_session.py @@ -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 -- diff --git a/test/registered/unit/disaggregation/test_specv2_kvcache_offloading.py b/test/registered/unit/disaggregation/test_specv2_kvcache_offloading.py index 0268eef4f..f915b68ec 100644 --- a/test/registered/unit/disaggregation/test_specv2_kvcache_offloading.py +++ b/test/registered/unit/disaggregation/test_specv2_kvcache_offloading.py @@ -20,6 +20,8 @@ from sglang.srt.disaggregation.decode_kvcache_offload_manager import ( from sglang.srt.disaggregation.kv_events import OffloadedState from sglang.srt.managers.cache_controller import HiCacheAck from sglang.srt.managers.schedule_batch import ReqKvInfo +from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator +from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache from sglang.test.ci.ci_register import register_cpu_ci register_cpu_ci(est_time=8, suite="base-a-test-cpu") @@ -47,6 +49,31 @@ def _make_mock_req( return req +class _RecordingAllocator(BaseTokenToKVPoolAllocator): + """Single-pool double. Subclassing the base routes free_full / free_segment / + free_segments into free(), so a new free API cannot slip past the recorder.""" + + def __init__(self, page_size: int): + super().__init__( + size=1024, + page_size=page_size, + dtype=torch.bfloat16, + device="cpu", + kvcache=None, + need_sort=False, + ) + self.freed = [] + + def clear(self): + self.freed = [] + + def alloc(self, need_size: int): + raise NotImplementedError + + def free(self, free_index: torch.Tensor): + self.freed.append(free_index.clone()) + + def _make_manager(pool_size: int, page_size: int = 1): """Create a DecodeKVCacheOffloadManager with mock pools for testing.""" # Build a real req_to_token tensor so indexing works @@ -55,15 +82,16 @@ def _make_manager(pool_size: int, page_size: int = 1): req_to_token_pool = MagicMock() req_to_token_pool.req_to_token = req_to_token - freed_indices = [] - - allocator = MagicMock() - allocator.free = MagicMock( - side_effect=lambda idx: freed_indices.append(idx.clone()) - ) + allocator = _RecordingAllocator(page_size) + freed_indices = allocator.freed tree_cache = MagicMock() tree_cache.protected_size_ = 0 + tree_cache.req_to_token_pool = req_to_token_pool + tree_cache.token_to_kv_pool_allocator = allocator + tree_cache.free_kv_row = lambda owner, ranges: BasePrefixCache.free_kv_row( + tree_cache, owner, ranges + ) # Bypass __init__ entirely and set attributes directly manager = object.__new__(DecodeKVCacheOffloadManager) diff --git a/test/registered/unit/layers/test_minicpm_sparse_cache.py b/test/registered/unit/layers/test_minicpm_sparse_cache.py index 3dffb3729..d42fd0f72 100644 --- a/test/registered/unit/layers/test_minicpm_sparse_cache.py +++ b/test/registered/unit/layers/test_minicpm_sparse_cache.py @@ -14,6 +14,7 @@ from sglang.srt.managers.scheduler_components.invariant_checker import ( from sglang.srt.managers.scheduler_components.pool_stats_observer import ( SchedulerPoolStatsObserver, ) +from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator from sglang.srt.mem_cache.memory_pool import ReqToTokenPool from sglang.srt.session.streaming_session import SessionSlot, StreamingSession from sglang.test.ci.ci_register import register_cpu_ci @@ -21,17 +22,23 @@ from sglang.test.ci.ci_register import register_cpu_ci register_cpu_ci(est_time=5, suite="base-a-test-cpu") -class RecordingAllocator: +class RecordingAllocator(BaseTokenToKVPoolAllocator): + """Single-pool double. Subclassing the base routes free_full / free_segment / + free_segments into free(), so a new free API cannot slip past the recorder.""" + def __init__(self, capacity: int): + super().__init__( + size=capacity, + page_size=1, + dtype=torch.bfloat16, + device="cpu", + kvcache=None, + need_sort=False, + ) self.capacity = capacity - self.page_size = 1 self.next_slot = 1 self.live: set[int] = set() - @property - def size(self): - return self.capacity - def alloc(self, size: int): if size > self.available_size(): return None 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 0f9957b7c..73528e65b 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 @@ -85,6 +85,7 @@ class MockReq: kv_committed_len=len(fill_ids), kv_allocated_len=len(fill_ids), cache_protected_len=cache_protected_len, + swa_evicted_seqlen=0, ) def get_fill_ids(self): 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 c0add4ae0..47064be66 100644 --- a/test/registered/unit/mem_cache/test_streaming_session_unit.py +++ b/test/registered/unit/mem_cache/test_streaming_session_unit.py @@ -3,6 +3,7 @@ from types import SimpleNamespace import torch from sglang.srt.managers.schedule_batch import FINISH_ABORT, ReqKvInfo +from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator from sglang.srt.mem_cache.base_prefix_cache import MatchResult from sglang.srt.session.streaming_session import SessionSlot, StreamingSession from sglang.test.ci.ci_register import register_cpu_ci @@ -10,10 +11,27 @@ from sglang.test.ci.ci_register import register_cpu_ci register_cpu_ci(est_time=12, suite="base-a-test-cpu") -class _FakeAllocator: - def __init__(self): +class _FakeAllocator(BaseTokenToKVPoolAllocator): + """Single-pool double. Subclassing the base routes free_full / free_segment / + free_segments into free(), so a new free API cannot slip past the recorder.""" + + def __init__(self, page_size: int = 1): + super().__init__( + size=1024, + page_size=page_size, + dtype=torch.bfloat16, + device="cpu", + kvcache=None, + need_sort=False, + ) self.freed = [] + def clear(self): + self.freed = [] + + def alloc(self, need_size: int): + raise NotImplementedError + def free(self, free_index: torch.Tensor): self.freed.append(free_index.clone()) @@ -114,7 +132,7 @@ def test_preabort_detaches_session_and_preserves_slot(): the session: session=None, abort_req() called. Slot stays intact.""" req_to_token = torch.arange(256, dtype=torch.int32).reshape(2, 128) req_to_token_pool = _FakeReqToTokenPool(req_to_token) - allocator = _FakeAllocator() + allocator = _FakeAllocator(page_size=16) inner = _FakeInnerCache( req_to_token_pool, allocator, @@ -283,9 +301,41 @@ def test_trim_overshoot_postcondition(): assert req.kv.kv_allocated_len == target assert req.kv.swa_evicted_seqlen == target assert len(req.output_ids) == 12 - # Tail [38, 44) freed by _free_kv_aligned. - assert len(allocator.freed) == 1 - assert allocator.freed[0].tolist() == list(range(38, 44)) + # Tail [38, 44) freed by _free_kv_aligned, split at the pre-trim eviction + # floor 42: [38, 42) gave its SWA peers back already, so it goes back full-only. + assert [t.tolist() for t in allocator.freed] == [[38, 39, 40, 41], [42, 43]] + + +def test_trim_overshoot_keeps_cursor_page_aligned_on_paged(): + """A mid-page trim target must not become the SWA eviction cursor (the + dead/alive split there frees the shared page twice); rewind to the boundary.""" + page_size = 16 + req_to_token = torch.arange(128, dtype=torch.int32).reshape(1, 128) + req_to_token_pool = _FakeReqToTokenPool(req_to_token) + allocator = _FakeAllocator(page_size=page_size) + tree_cache = StreamingSession( + _FakeInnerCache(req_to_token_pool, allocator, page_size) + ) + + # origin=26, finished=12 -> raw target 38 (mid-page); cursor 48 > target. + req = _FakeReq("session-a", req_pool_idx=0, committed=52, allocated=64) + req.origin_input_ids = list(range(26)) + req.output_ids = list(range(14)) + req.kv.swa_evicted_seqlen = 48 + + tree_cache._trim_overshoot(req, finished_len=12) + + # Rewound to floor_align(38) = 32; every cursor lands page-aligned. + assert req.kv.kv_allocated_len == 32 + assert req.kv.kv_committed_len == 32 + assert req.kv.swa_evicted_seqlen == 32 + assert len(req.output_ids) == 12 + # Freed [32, 64): [32, 48) below the old cursor goes back full-only, + # [48, 64) both halves. + assert [t.tolist() for t in allocator.freed] == [ + list(range(32, 48)), + list(range(48, 64)), + ] if __name__ == "__main__": diff --git a/test/registered/unit/mem_cache/test_swa_unittest.py b/test/registered/unit/mem_cache/test_swa_unittest.py index b01e5ca86..3d41cba75 100644 --- a/test/registered/unit/mem_cache/test_swa_unittest.py +++ b/test/registered/unit/mem_cache/test_swa_unittest.py @@ -6,8 +6,10 @@ import torch from sglang.srt.disaggregation.kv_events import BlockRemoved, BlockStored from sglang.srt.environ import envs +from sglang.srt.mem_cache.allocator.base import BaseTokenToKVPoolAllocator from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator from sglang.srt.mem_cache.base_prefix_cache import ( + BasePrefixCache, DecLockRefParams, EvictParams, EvictResult, @@ -15,7 +17,10 @@ from sglang.srt.mem_cache.base_prefix_cache import ( MatchPrefixParams, ) from sglang.srt.mem_cache.cache_init_params import CacheInitParams -from sglang.srt.mem_cache.common import available_and_evictable_str +from sglang.srt.mem_cache.common import ( + available_and_evictable_str, + free_kv_row_segments, +) from sglang.srt.mem_cache.memory_pool import ReqToTokenPool from sglang.srt.mem_cache.radix_cache import RadixKey from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool @@ -851,6 +856,31 @@ class TestSWASplitLeafOnInsert(CustomTestCase): tree.sanity_check() +class _SinglePoolAllocator(BaseTokenToKVPoolAllocator): + """Minimal single-pool allocator: no SWA peer, so the whole range dies + together whatever the floor says.""" + + def __init__(self): + super().__init__( + size=16, + page_size=1, + dtype=torch.bfloat16, + device="cpu", + kvcache=None, + need_sort=False, + ) + self.freed = [] + + def clear(self): + self.freed = [] + + def alloc(self, need_size: int): + raise NotImplementedError + + def free(self, free_index: torch.Tensor): + self.freed.append(free_index) + + class TestFreeFullPartition(CustomTestCase): """`free_full` releases only the full side of a hybrid SWA allocator.""" @@ -894,6 +924,96 @@ class TestFreeFullPartition(CustomTestCase): self.assertEqual(self.allocator.full_available_size(), self.full_baseline) +class _RowCache: + """Minimal PrefixCacheTrait host, so free_kv_row can be exercised without + standing up a whole tree.""" + + free_kv_row = BasePrefixCache.free_kv_row + + def __init__(self, allocator, row): + self.req_to_token_pool = SimpleNamespace(req_to_token=row.unsqueeze(0)) + self.token_to_kv_pool_allocator = allocator + self.page_size = allocator.page_size + + +class TestFreeKvRow(CustomTestCase): + """A kv row is given back split at `swa_evicted_seqlen`: the full side + whole, the SWA side only from the floor up.""" + + def setUp(self): + _, self.allocator, _ = _build_swa_tree(is_eagle=False) + self.full_baseline = self.allocator.full_available_size() + self.swa_baseline = self.allocator.swa_available_size() + + def _sizes(self): + return ( + self.allocator.full_available_size(), + self.allocator.swa_available_size(), + ) + + def test_floor_decides_how_much_of_the_swa_side_stays_out(self): + # (start_pos, num_slots, floor, rows whose SWA peers are already gone) + cases = [ + (0, 4, 4, 4), + (8, 4, 8, 0), + (8, 4, 10, 2), + (8, 4, 4, 0), + ] + for start_pos, num_slots, floor, num_dead in cases: + with self.subTest(start_pos=start_pos, floor=floor): + indices = _swa_alloc(self.allocator, num_slots) + free_kv_row_segments( + self.allocator, [(indices, start_pos)], swa_evicted_seqlen=floor + ) + self.assertEqual( + self._sizes(), + (self.full_baseline, self.swa_baseline - num_dead), + ) + # Give the held-back SWA peers back, so the next case starts clean. + if num_dead: + self.allocator.free_swa(indices[:num_dead]) + self.assertEqual(self._sizes(), (self.full_baseline, self.swa_baseline)) + + def test_adjacent_below_floor_pieces_release_their_shared_page_once(self): + _, allocator, _ = _build_swa_tree(is_eagle=False, page_size=4) + indices = _swa_alloc(allocator, 8) + after_alloc = allocator.full_available_size() + + # Rows [0, 6) and [6, 8) both sit below the floor and share page 1. + free_kv_row_segments( + allocator, + [(indices[:6], 0), (indices[6:], 6)], + swa_evicted_seqlen=8, + ) + + self.assertEqual(allocator.full_available_size(), after_alloc + 8) + + def test_free_kv_row_reads_the_record_row_and_its_floor(self): + indices = _swa_alloc(self.allocator, 8) + cache = _RowCache(self.allocator, indices) + kv = SimpleNamespace(req_pool_idx=0, swa_evicted_seqlen=3) + + cache.free_kv_row(kv, [(1, 5)]) + + # Rows [1, 5) go back on the full side; of those, [1, 3) lost their SWA + # peers already, so 6 of the 8 SWA slots are still out. + self.assertEqual(self._sizes(), (self.full_baseline - 4, self.swa_baseline - 6)) + + def test_single_pool_free_kv_row_still_frees_the_whole_range(self): + allocator = _SinglePoolAllocator() + cache = _RowCache(allocator, torch.arange(16, dtype=torch.int64)) + kv = SimpleNamespace(req_pool_idx=0, swa_evicted_seqlen=4) + + cache.free_kv_row(kv, [(2, 6)]) + + self.assertEqual([t.tolist() for t in allocator.freed], [[2, 3], [4, 5]]) + + # release_session and _free_kv_aligned dropped their own emptiness + # guards, so an empty range has to stay a no-op here. + cache.free_kv_row(kv, [(6, 6)]) + self.assertEqual(len(allocator.freed), 2) + + class TestCacheUnfinishedReqEvictedPrefix(CustomTestCase): """An unfinished request whose SWA prefix is already gone must insert that prefix as a tombstone, not as live SWA KV."""