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