[Perf] Free KV pages by segment in the paged allocator without a device sync (#32701)

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)

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):

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:

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)
 

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)

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)
 

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:

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()

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):

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):