diff --git a/python/sglang/srt/mem_cache/allocator/base.py b/python/sglang/srt/mem_cache/allocator/base.py index 1bdd99975..1a8b08560 100644 --- a/python/sglang/srt/mem_cache/allocator/base.py +++ b/python/sglang/srt/mem_cache/allocator/base.py @@ -108,3 +108,27 @@ class BaseTokenToKVPoolAllocator(abc.ABC): @abc.abstractmethod def free(self, free_index: torch.Tensor): raise NotImplementedError() + + def free_segment(self, free_index: torch.Tensor, *, start_pos: int): + """Free ``kv_row[start_pos : start_pos + n]`` of one request (or a + page-aligned copy); subclasses may use ``start_pos`` to skip the + data-dependent dedup. Default: plain free().""" + self.free(free_index) + + def free_segments(self, segments): + """Free disjoint ascending ``(free_index, start_pos)`` segments of one + request's kv row; a boundary page shared by consecutive segments is + emitted once (the later segment's head is trimmed).""" + ps = self.page_size + prev_end = None + for free_index, start_pos in segments: + n = free_index.numel() + if n == 0: + continue + seg_end = start_pos + n + if prev_end is not None and start_pos // ps == (prev_end - 1) // ps: + boundary = (start_pos // ps + 1) * ps + free_index = free_index[boundary - start_pos :] + start_pos = boundary + prev_end = seg_end + self.free_segment(free_index, start_pos=start_pos) diff --git a/python/sglang/srt/mem_cache/allocator/paged.py b/python/sglang/srt/mem_cache/allocator/paged.py index eb34cf7a5..9cfdfa069 100755 --- a/python/sglang/srt/mem_cache/allocator/paged.py +++ b/python/sglang/srt/mem_cache/allocator/paged.py @@ -263,16 +263,68 @@ class PagedTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): return if self.is_not_in_free_group: - free_page_indices = torch.unique(free_index // self.page_size) - if self.need_sort: - self.release_pages = torch.cat((free_page_indices, self.release_pages)) - else: - self.free_pages = torch.cat((free_page_indices, self.free_pages)) + self._release_page_ids(torch.unique(free_index // self.page_size)) else: self.free_group.append(free_index) if self.debug_mode: - assert len(torch.unique(self.free_pages)) == len(self.free_pages) + self._debug_check_no_duplicate_pages() + + def free_segment(self, free_index: torch.Tensor, *, start_pos: int): + """Fixed-shape counterpart of free(): a page's tokens sit consecutively + in the kv row, so page representatives are stride slices -- no + torch.unique, whose data-dependent output shape forces a device sync. + Contract: see base; a page must be freed by only one call per group.""" + if free_index.numel() == 0: + return + + ps = self.page_size + offset = start_pos % ps + if offset == 0: + pieces = (free_index[::ps],) + else: + pieces = (free_index[:1], free_index[ps - offset :: ps]) + + if self.debug_mode: + # reference unique on CPU: the NPU subclass deliberately avoids device unique + page_ids = torch.cat([p // ps for p in pieces]) + assert torch.equal( + torch.sort(page_ids.cpu())[0], + torch.unique(free_index.cpu() // ps), + ) + + if self.is_not_in_free_group: + self._release_page_ids(*(p // ps for p in pieces)) + if self.debug_mode: + self._debug_check_no_duplicate_pages() + else: + self.free_page_reps_group.extend(pieces) + + def _debug_check_no_duplicate_pages(self): + # span both containers: need_sort (PD disagg) routes frees into release_pages + pages = torch.cat((self.free_pages, self.release_pages)) + assert len(torch.unique(pages)) == len(pages) + + def _release_page_ids(self, *page_ids: torch.Tensor): + if self.need_sort: + self.release_pages = torch.cat((*page_ids, self.release_pages)) + else: + self.free_pages = torch.cat((*page_ids, self.free_pages)) + + def free_group_begin(self): + super().free_group_begin() + self.free_page_reps_group = [] + + def free_group_end(self): + super().free_group_end() + if self.free_page_reps_group: + self._release_page_ids( + torch.cat(self.free_page_reps_group) // self.page_size + ) + self.free_page_reps_group = [] + if self.debug_mode: + # the no-double-free contract can only break across a group's calls + self._debug_check_no_duplicate_pages() def clear(self): # The padded slot 0 is used for writing dummy outputs from padded tokens. @@ -281,6 +333,7 @@ class PagedTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): ) self.is_not_in_free_group = True self.free_group = [] + self.free_page_reps_group = [] self.release_pages = torch.empty((0,), dtype=torch.int64, device=self.device) def get_cpu_copy(self, indices, mamba_indices=None): diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index 25d3f96e8..d50c77677 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/python/sglang/srt/mem_cache/common.py @@ -196,7 +196,11 @@ def _release_overallocated_kv_indices( 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) + # start_p is ceil-aligned above: never shares a page with + # cache_finished_req's tail frees in the same group. + tree_cache.token_to_kv_pool_allocator.free_segment( + indices_to_free, start_pos=start_p + ) def available_and_evictable_str(tree_cache: BasePrefixCache) -> str: diff --git a/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py b/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py index efcd42366..81a6802c5 100644 --- a/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py +++ b/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py @@ -898,7 +898,12 @@ class HiMambaRadixCache(MambaRadixCache): else: if prev_prefix_len < total_prefix_length + prefix_len: start = max(0, prev_prefix_len - total_prefix_length) - self.token_to_kv_pool_allocator.free(value[start:prefix_len]) + # same as MambaRadixCache._insert_helper: page-exact segment + # at offset total_prefix_length of the kv row + self.token_to_kv_pool_allocator.free_segment( + value[start:prefix_len], + start_pos=total_prefix_length + start, + ) total_prefix_length += prefix_len self._inc_hit_count(node, chunked) diff --git a/python/sglang/srt/mem_cache/mamba_radix_cache.py b/python/sglang/srt/mem_cache/mamba_radix_cache.py index 20905dd87..6daf9b16c 100644 --- a/python/sglang/srt/mem_cache/mamba_radix_cache.py +++ b/python/sglang/srt/mem_cache/mamba_radix_cache.py @@ -546,7 +546,7 @@ class MambaRadixCache(KVCacheEventMixin, BasePrefixCache): kv_indices = self.req_to_token_pool.req_to_token[ req.req_pool_idx, :kv_len_to_handle ] - self.token_to_kv_pool_allocator.free(kv_indices) + self.token_to_kv_pool_allocator.free_segment(kv_indices, start_pos=0) self.req_to_token_pool.free_mamba_cache(req) return @@ -573,7 +573,9 @@ class MambaRadixCache(KVCacheEventMixin, BasePrefixCache): cache_len = 0 if cache_len != len(token_ids): cache_end_idx = max(cache_len, req.cache_protected_len) - self.token_to_kv_pool_allocator.free(kv_indices[cache_end_idx:]) + self.token_to_kv_pool_allocator.free_segment( + kv_indices[cache_end_idx:], start_pos=cache_end_idx + ) token_ids = token_ids[:cache_len] kv_indices = kv_indices[:cache_len] @@ -636,7 +638,10 @@ class MambaRadixCache(KVCacheEventMixin, BasePrefixCache): # state already cached -> the int8 slot we just allocated is a duplicate self.int8_ckpt_pool.free(mamba_value) else: - self.token_to_kv_pool_allocator.free(kv_indices[req.cache_protected_len :]) + self.token_to_kv_pool_allocator.free_segment( + kv_indices[req.cache_protected_len :], + start_pos=req.cache_protected_len, + ) mamba_exist = True if mamba_exist: @@ -797,7 +802,8 @@ class MambaRadixCache(KVCacheEventMixin, BasePrefixCache): assert x.mamba_value is not None, f"leaf node mamba value is not None, {x.id=}" # 1. a leaf node, free full tokens and mamba self._record_remove_event(x) - self.token_to_kv_pool_allocator.free(x.value) + # Tree values are page-aligned copies of a kv row: page-exact segment. + self.token_to_kv_pool_allocator.free_segment(x.value, start_pos=0) full_num_evicted = len(x.value) self._free_mamba_value(x.mamba_value) mamba_num_evicted = len(x.mamba_value) @@ -1244,7 +1250,12 @@ class MambaRadixCache(KVCacheEventMixin, BasePrefixCache): if prev_prefix_len < total_prefix_length + prefix_len: start = max(0, prev_prefix_len - total_prefix_length) - self.token_to_kv_pool_allocator.free(value[start:prefix_len]) + # value sits at offset total_prefix_length of the kv row; match() + # rounds prefix_len to page multiples, so frees never share a page. + self.token_to_kv_pool_allocator.free_segment( + value[start:prefix_len], + start_pos=total_prefix_length + start, + ) total_prefix_length += prefix_len key = key[prefix_len:] @@ -1299,7 +1310,7 @@ class MambaRadixCache(KVCacheEventMixin, BasePrefixCache): ), f"tombstone mamba_lock_ref should always be 0, {node.parent.full_lock_ref=}, {node.parent.mamba_lock_ref=}, {node.parent.id=}" # delete tombstone node evicts full tokens self._record_remove_event(node.parent) - self.token_to_kv_pool_allocator.free(node.parent.value) + self.token_to_kv_pool_allocator.free_segment(node.parent.value, start_pos=0) full_num_evicted += len(node.parent.value) self.full_lru_list.remove_node(node.parent) self._delete_tombstone_leaf(node.parent) diff --git a/python/sglang/srt/mem_cache/radix_cache.py b/python/sglang/srt/mem_cache/radix_cache.py index b766be298..4ac746313 100644 --- a/python/sglang/srt/mem_cache/radix_cache.py +++ b/python/sglang/srt/mem_cache/radix_cache.py @@ -447,7 +447,9 @@ class RadixCache(SessionRadixCacheMixin, KVCacheEventMixin, BasePrefixCache): kv_indices = self.req_to_token_pool.req_to_token[ req.req_pool_idx, req.cache_protected_len : kv_len_to_handle ] - self.token_to_kv_pool_allocator.free(kv_indices) + self.token_to_kv_pool_allocator.free_segment( + kv_indices, start_pos=req.cache_protected_len + ) return token_ids = (req.origin_input_ids + req.output_ids)[:kv_len_to_handle] @@ -468,18 +470,21 @@ class RadixCache(SessionRadixCacheMixin, KVCacheEventMixin, BasePrefixCache): InsertParams(key=radix_key, value=values, priority=priority) ) session_leaf = result.last_device_node - # Free the duplicates that were already in the tree - self.token_to_kv_pool_allocator.free( - kv_indices[req.cache_protected_len : result.prefix_len] - ) + freed_end = result.prefix_len else: session_leaf = None - self.token_to_kv_pool_allocator.free( - kv_indices[req.cache_protected_len : key_len] - ) + freed_end = key_len - # free the unaligned tail - self.token_to_kv_pool_allocator.free(kv_indices[key_len:]) + # duplicates / uninserted range, then the unaligned tail + self.token_to_kv_pool_allocator.free_segments( + [ + ( + kv_indices[req.cache_protected_len : freed_end], + req.cache_protected_len, + ), + (kv_indices[key_len:], key_len), + ] + ) self._tag_session_leaf(req, radix_key, node=session_leaf) @@ -513,8 +518,9 @@ class RadixCache(SessionRadixCacheMixin, KVCacheEventMixin, BasePrefixCache): ) new_prefix_len = result.prefix_len - self.token_to_kv_pool_allocator.free( - kv_indices[req.cache_protected_len : new_prefix_len] + self.token_to_kv_pool_allocator.free_segment( + kv_indices[req.cache_protected_len : new_prefix_len], + start_pos=req.cache_protected_len, ) # The prefix indices could be updated, reuse it @@ -578,7 +584,8 @@ class RadixCache(SessionRadixCacheMixin, KVCacheEventMixin, BasePrefixCache): while num_evicted < num_tokens and len(eviction_heap): _priority, x = heapq.heappop(eviction_heap) - self.token_to_kv_pool_allocator.free(x.value) + # Tree values are page-aligned copies of a kv row: page-exact segment. + self.token_to_kv_pool_allocator.free_segment(x.value, start_pos=0) num_evicted += len(x.value) self._delete_leaf(x) diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index d699ebd1f..234225c89 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -588,7 +588,7 @@ class UnifiedRadixCache(BasePrefixCache): kv_indices = self.req_to_token_pool.req_to_token[ req.req_pool_idx, :kv_len_to_handle ] - self.token_to_kv_pool_allocator.free(kv_indices) + self.token_to_kv_pool_allocator.free_segment(kv_indices, start_pos=0) for comp in self._components_tuple: comp.cleanup_after_caching_req(req, is_finished=True) return @@ -619,10 +619,12 @@ class UnifiedRadixCache(BasePrefixCache): if cl is not None: effective_cache_len = min(effective_cache_len, cl) - # Truncate if needed + # Truncate if needed; the tail free is deferred and batched with + # the unaligned tail below so a shared boundary page is emitted once. + kv_indices_full = kv_indices + tail_free_start = None if effective_cache_len < len(token_ids): - free_start = max(effective_cache_len, req.cache_protected_len) - self.token_to_kv_pool_allocator.free(kv_indices[free_start:]) + tail_free_start = max(effective_cache_len, req.cache_protected_len) token_ids = token_ids[:effective_cache_len] kv_indices = kv_indices[:effective_cache_len] @@ -636,10 +638,16 @@ class UnifiedRadixCache(BasePrefixCache): insert_params.value = values result = self.insert(insert_params) - # Free unaligned tail - self.token_to_kv_pool_allocator.free(kv_indices[page_aligned_len:]) + # Free unaligned tail (+ deferred truncation tail) + segments = [(kv_indices[page_aligned_len:], page_aligned_len)] + 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) else: - self.token_to_kv_pool_allocator.free(kv_indices[req.cache_protected_len :]) + self.token_to_kv_pool_allocator.free_segment( + kv_indices[req.cache_protected_len :], + start_pos=req.cache_protected_len, + ) self.dec_lock_ref( req.last_node, @@ -784,8 +792,9 @@ class UnifiedRadixCache(BasePrefixCache): [action.new_node_id, action.new_child_node_id], ) elif isinstance(action, FreeDeviceKV): + # tree values are page-aligned copies of a kv row: page-exact segments for indices in action.indices: - self.token_to_kv_pool_allocator.free(indices) + self.token_to_kv_pool_allocator.free_segment(indices, start_pos=0) elif isinstance(action, BackupKV): self._execute_and_commit_kv_backup(action) else: diff --git a/test/registered/unit/mem_cache/test_paged_free_segment.py b/test/registered/unit/mem_cache/test_paged_free_segment.py new file mode 100644 index 000000000..6e35fda3e --- /dev/null +++ b/test/registered/unit/mem_cache/test_paged_free_segment.py @@ -0,0 +1,179 @@ +"""free_segment / free_segments vs the torch.unique reference: stride page +extraction over all segment alignments, plus boundary-page dedup and free-group +deferral. See PagedTokenToKVPoolAllocator.free_segment for why unique is avoided. + + python -m pytest test/registered/unit/mem_cache/test_paged_free_segment.py -v +""" + +import unittest + +import torch + +from sglang.srt.mem_cache.allocator.base import BaseTokenToKVPoolAllocator +from sglang.srt.mem_cache.allocator.paged import PagedTokenToKVPoolAllocator +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=15, suite="base-a-test-cpu") + +PAGE_SIZE = 4 +NUM_PAGES = 64 + + +def _make_allocator(need_sort=False): + return PagedTokenToKVPoolAllocator( + size=NUM_PAGES * PAGE_SIZE, + page_size=PAGE_SIZE, + dtype=torch.float16, + device="cpu", + kvcache=None, + need_sort=need_sort, + ) + + +def _make_kv_row(alloc, num_tokens): + # Page-aligned allocation, then trim to num_tokens: mirrors a request's + # req_to_token row (token position t lives at page*page_size + t%page_size). + num_pages = -(num_tokens // -PAGE_SIZE) + indices = alloc.alloc(num_pages * PAGE_SIZE) + return indices[:num_tokens] + + +class TestFreeSegment(unittest.TestCase): + def test_matches_unique_over_alignments(self): + # Sweep (start, end) so segments cover: aligned/unaligned head and + # tail, single partial page, full row. + for num_tokens in (1, PAGE_SIZE, PAGE_SIZE + 1, 3 * PAGE_SIZE - 1): + for start in range(num_tokens): + for end in range(start + 1, num_tokens + 1): + alloc = _make_allocator() + row = _make_kv_row(alloc, num_tokens) + expected = torch.unique(row[start:end] // PAGE_SIZE) + before = len(alloc.free_pages) + alloc.free_segment(row[start:end], start_pos=start) + freed = alloc.free_pages[: len(alloc.free_pages) - before] + self.assertTrue( + torch.equal(torch.sort(freed)[0], expected), + f"{num_tokens=} {start=} {end=}", + ) + + def test_empty_segment_is_noop(self): + alloc = _make_allocator() + row = _make_kv_row(alloc, PAGE_SIZE) + before = len(alloc.free_pages) + alloc.free_segment(row[:0], start_pos=0) + self.assertEqual(len(alloc.free_pages), before) + + def test_need_sort_routes_to_release_pages(self): + alloc = _make_allocator(need_sort=True) + row = _make_kv_row(alloc, 2 * PAGE_SIZE) + alloc.free_segment(row, start_pos=0) + self.assertEqual(len(alloc.release_pages), 2) + + def test_group_defers_until_group_end(self): + alloc = _make_allocator() + row = _make_kv_row(alloc, 2 * PAGE_SIZE) + before = len(alloc.free_pages) + alloc.free_group_begin() + alloc.free_segment(row, start_pos=0) + self.assertEqual(len(alloc.free_pages), before) + alloc.free_group_end() + self.assertEqual(len(alloc.free_pages), before + 2) + + def test_group_end_debug_assert_catches_cross_call_double_free(self): + # legacy free() + free_segment() on the same page in one group must + # trip free_group_end's debug assert + alloc = _make_allocator() + alloc.debug_mode = True + row = _make_kv_row(alloc, PAGE_SIZE) + alloc.free_group_begin() + alloc.free(row) + alloc.free_segment(row, start_pos=0) + with self.assertRaises(AssertionError): + alloc.free_group_end() + + def test_group_end_debug_assert_covers_release_pages(self): + # need_sort routes frees into release_pages; the duplicate check must + # not go vacuous there (PD disaggregation runs with need_sort=True). + alloc = _make_allocator(need_sort=True) + alloc.debug_mode = True + row = _make_kv_row(alloc, PAGE_SIZE) + alloc.free_group_begin() + alloc.free(row) + alloc.free_segment(row, start_pos=0) + with self.assertRaises(AssertionError): + alloc.free_group_end() + + +class TestFreeSegments(unittest.TestCase): + def _freed_by_segments(self, num_tokens, spans): + alloc = _make_allocator() + row = _make_kv_row(alloc, num_tokens) + before = len(alloc.free_pages) + alloc.free_segments([(row[a:b], a) for a, b in spans]) + freed = alloc.free_pages[: len(alloc.free_pages) - before] + reference = torch.unique(torch.cat([row[a:b] for a, b in spans]) // PAGE_SIZE) + return freed, reference + + def test_adjacent_segments_share_boundary_page(self): + # [0, 6) and [6, 11) with page_size 4: page 1 spans both segments and + # must be freed exactly once. + freed, reference = self._freed_by_segments(11, [(0, 6), (6, 11)]) + self.assertTrue(torch.equal(torch.sort(freed)[0], reference)) + + def test_disjoint_segments_share_boundary_page(self): + # [0, 5) and [7, 11): gap [5, 7) stays within page 1, which both + # segments touch. + freed, reference = self._freed_by_segments(11, [(0, 5), (7, 11)]) + self.assertTrue(torch.equal(torch.sort(freed)[0], reference)) + + def test_second_segment_inside_shared_page_is_skipped(self): + # [0, 5) and [5, 7): the second segment lies entirely in page 1, + # already emitted by the first. + freed, reference = self._freed_by_segments(7, [(0, 5), (5, 7)]) + self.assertTrue(torch.equal(torch.sort(freed)[0], reference)) + + def test_page_aligned_segments_no_trim(self): + freed, reference = self._freed_by_segments( + 3 * PAGE_SIZE, [(0, PAGE_SIZE), (PAGE_SIZE, 3 * PAGE_SIZE)] + ) + self.assertTrue(torch.equal(torch.sort(freed)[0], reference)) + + +class _RecordingBaseAllocator(BaseTokenToKVPoolAllocator): + """Base-fallback allocator: free_segment inherits the default (ignore + start_pos, call free()), free() records what it received.""" + + def __init__(self): + super().__init__( + size=NUM_PAGES * PAGE_SIZE, + page_size=PAGE_SIZE, + dtype=torch.float16, + device="cpu", + kvcache=None, + need_sort=False, + ) + self.freed = [] + + def alloc(self, need_size: int): + raise NotImplementedError + + def clear(self): + pass + + def free(self, free_index: torch.Tensor): + self.freed.append(free_index) + + +class TestBaseFallbackFreeSegments(unittest.TestCase): + def test_trim_dedups_boundary_page_before_fallback_free(self): + # fallback allocators (UnifiedMamba/SWA) dedup per free() call at best; + # the shared boundary page must reach free() in exactly one call + alloc = _RecordingBaseAllocator() + row = torch.arange(11) # position i lives on page i // PAGE_SIZE + alloc.free_segments([(row[0:6], 0), (row[6:11], 6)]) + per_call_pages = [set((t // PAGE_SIZE).tolist()) for t in alloc.freed] + self.assertEqual(per_call_pages, [{0, 1}, {2}]) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/mem_cache/test_radix_cache_unit.py b/test/registered/unit/mem_cache/test_radix_cache_unit.py index 5bbe92b5f..778afa06c 100644 --- a/test/registered/unit/mem_cache/test_radix_cache_unit.py +++ b/test/registered/unit/mem_cache/test_radix_cache_unit.py @@ -590,8 +590,8 @@ class TestRadixCache(unittest.TestCase): f"evicted {result.num_tokens_evicted} tokens, expected at least 2", ) - # Should have called free and reduced size - mock_allocator.free.assert_called() + # Should have called free_segment and reduced size + mock_allocator.free_segment.assert_called() self.assertLess(cache.total_size(), initial_size) def test_page_alignment_boundary(self): diff --git a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py index 9b971612e..09f594213 100644 --- a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py +++ b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py @@ -5363,8 +5363,8 @@ class TestUnifiedRadixCacheActionRouting(CustomTestCase): first, second = torch.tensor([4, 5]), torch.tensor([6]) action = FreeDeviceKV([first, second]) UnifiedRadixCache._apply_cache_action(cache, action) - cache.token_to_kv_pool_allocator.free.assert_has_calls( - [mock.call(first), mock.call(second)] + cache.token_to_kv_pool_allocator.free_segment.assert_has_calls( + [mock.call(first, start_pos=0), mock.call(second, start_pos=0)] ) def test_apply_cache_action_routes_free_component_device_kv(self):