[mem_cache] Free hybrid SWA pages by one representative per page on page_size > 1 (#38159)

python/sglang/srt/mem_cache/allocator/hisparse.py

@@ -358,6 +358,9 @@ class DeepSeekV4HiSparseTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
     def free_swa(self, free_indices: torch.Tensor):
         self.logical_attn_allocator.free_swa(free_indices)
 
+    def free_swa_segment(self, free_indices: torch.Tensor, *, start_pos: int):
+        self.logical_attn_allocator.free_swa_segment(free_indices, start_pos=start_pos)
+
     def free_full(self, free_indices: torch.Tensor):
         if free_indices.numel() == 0:
             return

python/sglang/srt/mem_cache/allocator/paged.py

@@ -296,12 +296,19 @@ class PagedTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
                 torch.unique(free_index.cpu() // ps),
             )
 
+        self.free_page_ids(reps // ps)
+
+    def free_page_ids(self, page_ids: torch.Tensor):
+        """Free exactly these pages; no page twice, no dedup."""
+        if page_ids.numel() == 0:
+            return
+
         if self.free_group is None:
-            self._release_page_ids(reps // ps)
+            self._release_page_ids(page_ids)
             if self.debug_mode:
                 self._debug_check_no_duplicate_pages()
         else:
-            self.free_page_reps_group.append(self._copy_for_free_group(reps))
+            self.free_page_ids_group.append(self._copy_for_free_group(page_ids))
 
     def _debug_check_no_duplicate_pages(self):
         pages = self.get_all_free_pages()
@@ -316,15 +323,13 @@ class PagedTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
 
     def free_group_begin(self):
         super().free_group_begin()
-        self.free_page_reps_group = []
+        self.free_page_ids_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.free_page_ids_group:
+            self._release_page_ids(torch.cat(self.free_page_ids_group))
+            self.free_page_ids_group = []
         if self.debug_mode:
             # the no-double-free contract can only break across a group's calls
             self._debug_check_no_duplicate_pages()
@@ -335,7 +340,7 @@ class PagedTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
             1, self.num_pages + 1, dtype=torch.int64, device=self.device
         )
         self.free_group = None
-        self.free_page_reps_group = []
+        self.free_page_ids_group = []
         # need_sort only: freed pages wait here, unsorted, until an alloc runs short.
         self.staged_pages: list[torch.Tensor] = []
         self.num_staged_pages = 0

python/sglang/srt/mem_cache/allocator/swa.py

@@ -111,6 +111,7 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
         self.release_pages = None
         self.free_group = None
         self.swa_free_group = []
+        self.swa_page_ids_group = []
 
         self._kvcache = kvcache
 
@@ -429,19 +430,19 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
             self.full_to_swa_index_mapping.index_fill_(0, full_indices, 0)
 
     def free_swa(self, free_index: torch.Tensor):
+        """Release the SWA peers of an arbitrary slot set and clear their mapping.
+        Synchronizes at page_size > 1; kv-row segments go through free_swa_segment()."""
         if free_index.numel() == 0:
             return
 
         if self.page_size == 1:
-            # A filter here would make the output shape data-dependent,
-            # which costs a device-to-host sync.
-            mapping_indices = free_index
-            swa_indices = self.full_to_swa_index_mapping[mapping_indices]
-            expect(_SWA_PEER_MAPPED, swa_indices > 0, msg="caller wants free_full")
-        else:
-            mapping_indices = self._expand_to_full_pages(free_index)
-            swa_indices = self.full_to_swa_index_mapping[mapping_indices]
+            self._free_swa_pages(free_index, start_pos=0)
+            return
 
+        # The expansion gathers slots this caller never allocated, which read as
+        # the padding slot; `_release_swa` filters them once per group.
+        mapping_indices = self._expand_to_full_pages(free_index)
+        swa_indices = self.full_to_swa_index_mapping[mapping_indices]
         self.clear_full_to_swa_mapping(mapping_indices)
 
         if self.free_group is not None:
@@ -452,11 +453,50 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
 
         self._release_swa(swa_indices)
 
+    def free_swa_segment(self, free_index: torch.Tensor, *, start_pos: int):
+        """free_swa() for a kv-row segment; same start-alignment contract as
+        free_segment(), and fixed-shape at every page size."""
+        if free_index.numel() == 0:
+            return
+        self._free_swa_pages(free_index, start_pos=start_pos)
+
+    def _free_swa_pages(self, free_index: torch.Tensor, *, start_pos: int):
+        ps = self.page_size
+        assert start_pos % ps == 0, f"segment start {start_pos} is not page-aligned"
+        # First token of every page the segment touches; the caller allocated
+        # each one, so a dead entry means the caller wanted free_full.
+        reps = free_index[::ps]
+        swa_tokens = self.full_to_swa_index_mapping[reps]
+        expect(_SWA_PEER_MAPPED, swa_tokens > 0, msg="caller wants free_full")
+
+        if ps == 1:
+            swa_pages = swa_tokens
+            mapping_indices = free_index
+        else:
+            swa_pages = swa_tokens // ps
+            # Both pools page in step (alloc_extend / alloc_decode drive them
+            # with one seq_lens), so a rep's peer page is the whole peer page.
+            mapping_indices = self._expand_to_full_pages(reps)
+            if self.swa_attn_allocator.debug_mode:
+                ref = self.full_to_swa_index_mapping[mapping_indices].cpu()
+                assert torch.equal(
+                    torch.sort(swa_pages.cpu())[0],
+                    torch.unique(ref[ref > 0] // ps),
+                ), "swa pages do not match the mapped pages"
+        self.clear_full_to_swa_mapping(mapping_indices)
+
+        if self.free_group is not None:
+            # Resolve ownership now, as above.
+            self.swa_page_ids_group.append(swa_pages)
+            return
+
+        self.swa_attn_allocator.free_page_ids(swa_pages)
+        assert self.swa_attn_allocator.available_size() <= self.swa_attn_allocator.size
+
     def _release_swa(self, swa_indices: torch.Tensor):
         if self.page_size > 1:
-            # HiCache LOAD_BACK re-pairs a page-aligned full chunk with an offset
-            # SWA one (commit_hicache_transfer advances by raw token count), so a
-            # page can hold unmapped slots; one filter per group, not per call.
+            # Set-shaped frees only (see free_swa): drop the padding-slot entries
+            # the page expansion picked up.
             swa_indices = swa_indices[swa_indices > 0]
         self.swa_attn_allocator.free(swa_indices)
         assert self.swa_attn_allocator.available_size() <= self.swa_attn_allocator.size
@@ -482,7 +522,7 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
             return
         # SWA first, as in free(): it reads the mapping that a later cache
         # action in this group may re-point.
-        self.free_swa(free_index)
+        self.free_swa_segment(free_index, start_pos=start_pos)
         self.full_attn_allocator.free_segment(free_index, start_pos=start_pos)
 
     def free_full_segment(self, free_index: torch.Tensor, *, start_pos: int):
@@ -498,11 +538,16 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
     def free_group_begin(self):
         super().free_group_begin()
         self.swa_free_group = []
+        self.swa_page_ids_group = []
         # No full-side pile here: the full allocator's own group defers those.
         self.full_attn_allocator.free_group_begin()
 
     def free_group_end(self):
         super().free_group_end()
+        if self.swa_page_ids_group:
+            swa_page_ids_group = self.swa_page_ids_group
+            self.swa_page_ids_group = []
+            self.swa_attn_allocator.free_page_ids(torch.cat(swa_page_ids_group))
         if self.swa_free_group:
             swa_free_group = self.swa_free_group
             self.swa_free_group = []
@@ -553,6 +598,7 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
         self.full_to_swa_index_mapping[:-1].fill_(0)
         self.free_group = None
         self.swa_free_group = []
+        self.swa_page_ids_group = []
 
     def get_cpu_copy(self, indices, mamba_indices=None, req_pool_index=None):
         return self._kvcache.get_cpu_copy(
@@ -682,6 +728,9 @@ class PureSWATokenToKVPoolAllocator(SWATokenToKVPoolAllocator):
         else:
             self.free_group.append(self._copy_for_free_group(free_index))
 
+    def free_swa_segment(self, free_index: torch.Tensor, *, start_pos: int):
+        self.free_swa(free_index)
+
     def free_full(self, free_index: torch.Tensor):
         # All-SWA models have no full-attention pool, so there is nothing to
         # release once the SWA side is gone.
@@ -695,7 +744,7 @@ class PureSWATokenToKVPoolAllocator(SWATokenToKVPoolAllocator):
     def free_full_segment(self, free_index: torch.Tensor, *, start_pos: int):
         return
 
-    # Not inherited: the SWA parent's hooks drive swa_free_group and the full
+    # Not inherited: the SWA parent's hooks drive the SWA piles and the full
     # allocator's group, which this pure-SWA variant does not have.
     def free_group_begin(self):
         BaseTokenToKVPoolAllocator.free_group_begin(self)

python/sglang/srt/mem_cache/allocator/token.py

@@ -74,6 +74,10 @@ class TokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
         else:
             self.free_group.append(self._copy_for_free_group(free_index))
 
+    def free_page_ids(self, page_ids: torch.Tensor):
+        # page_size == 1: page ids are token ids.
+        self.free(page_ids)
+
     def get_cpu_copy(self, indices, mamba_indices=None, req_pool_index=None):
         return self._kvcache.get_cpu_copy(
             indices,

python/sglang/srt/mem_cache/allocator/unified_hybrid_swa.py

@@ -475,29 +475,13 @@ class UnifiedSWATokenToKVPoolAllocator(SWATokenToKVPoolAllocator):
             self.full_attn_allocator.clear_inverse_history()
             self.swa_attn_allocator.clear_inverse_history()
 
-    def free_swa(
-        self, free_index: torch.Tensor, *, start_pos: Optional[int] = None
-    ) -> None:
+    def free_swa(self, free_index: torch.Tensor) -> None:
         """SWA tombstone path: release swa-physical, keep the virtual id and
         full-physical live; `swa.v2p_page[v_page] = -1` IS the tombstone."""
         if free_index is None or free_index.numel() == 0:
             return
         v = free_index.detach().to(torch.int64)
         ps = self.page_size
-        # `start_pos` promises a contiguous ascending range starting at that prefix
-        # position, so page reps come from stride arithmetic, not `torch.unique`.
-        if start_pos is not None and ps > 1:
-            reps = self.swa_attn_allocator._page_reps(v, start_pos)
-            # Keep only pages still bound on swa; freeing a tombstoned one would
-            # corrupt the hole list. `> 0` strict: -1 tombstoned, 0 padding sink.
-            rep_pages = reps // ps
-            swa_v2p_pages = self.swa_attn_allocator.virtual_to_physical[rep_pages]
-            live_reps = reps[swa_v2p_pages > 0]
-            if live_reps.numel() == 0:
-                return
-            self.swa_attn_allocator.free(live_reps, _pages=live_reps // ps)
-            self.swa_attn_allocator.clear_inverse_history()
-            return
         v_pages = v // ps
         # `> 0` strict: -1 = tombstoned, page 0 = padding sink (never freeable).
         swa_v2p_pages = self.swa_attn_allocator.virtual_to_physical[v_pages]
@@ -512,6 +496,28 @@ class UnifiedSWATokenToKVPoolAllocator(SWATokenToKVPoolAllocator):
             self.swa_attn_allocator.free(live)
         self.swa_attn_allocator.clear_inverse_history()
 
+    def free_swa_segment(self, free_index: torch.Tensor, *, start_pos: int) -> None:
+        """free_swa() for a kv-row segment: `start_pos` promises a contiguous
+        ascending range, so page reps come from stride arithmetic, not `torch.unique`."""
+        if free_index is None or free_index.numel() == 0:
+            return
+        if self.page_size == 1:
+            self.free_swa(free_index)
+            return
+        ps = self.page_size
+        reps = self.swa_attn_allocator._page_reps(
+            free_index.detach().to(torch.int64), start_pos
+        )
+        # Keep only pages still bound on swa; freeing a tombstoned one would
+        # corrupt the hole list. `> 0` strict: -1 tombstoned, 0 padding sink.
+        rep_pages = reps // ps
+        swa_v2p_pages = self.swa_attn_allocator.virtual_to_physical[rep_pages]
+        live_reps = reps[swa_v2p_pages > 0]
+        if live_reps.numel() == 0:
+            return
+        self.swa_attn_allocator.free(live_reps, _pages=live_reps // ps)
+        self.swa_attn_allocator.clear_inverse_history()
+
     def free_full(self, free_index: torch.Tensor) -> None:
         """Release the full-physical page and the virtual id, leaving the swa
         side alone -- the caller already tombstoned it (`swa.v2p_page == -1`)."""

python/sglang/srt/mem_cache/common.py

@@ -101,21 +101,9 @@ def free_swa_out_of_window_slots(
         free_slots = req_to_token_pool.req_to_token[
             req.kv.req_pool_idx, req.kv.swa_evicted_seqlen : new_swa_evicted_seqlen
         ]
-        # Local import: the unified allocators import this module lazily for
-        # eviction; a module-level import here would be a cycle hazard.
-        from sglang.srt.mem_cache.allocator.unified_hybrid_swa import (
-            UnifiedSWATokenToKVPoolAllocator,
+        token_to_kv_pool_allocator.free_swa_segment(
+            free_slots, start_pos=req.kv.swa_evicted_seqlen
         )
-
-        if isinstance(token_to_kv_pool_allocator, UnifiedSWATokenToKVPoolAllocator):
-            # Contiguous range with host-int bounds: hand the composite its
-            # start position so the free stays host-sync-free (`free_segment`
-            # derives page reps by stride math instead of `torch.unique`).
-            token_to_kv_pool_allocator.free_swa(
-                free_slots, start_pos=req.kv.swa_evicted_seqlen
-            )
-        else:
-            token_to_kv_pool_allocator.free_swa(free_slots)
         req.kv.swa_evicted_seqlen = new_swa_evicted_seqlen
 
 

python/sglang/srt/mem_cache/swa_radix_cache.py

@@ -599,9 +599,9 @@ class SWARadixCache(BasePrefixCache):
             # SWA peers went back in `dec_swa_lock_only` or an SWA evict, so
             # only the full side is still ours; `free` would hand the SWA pool
             # mapping entries that read as the padding slot.
-            self.token_to_kv_pool_allocator.free_full(value)
+            self.token_to_kv_pool_allocator.free_full_segment(value, start_pos=0)
             return num_tokens, 0
-        self.token_to_kv_pool_allocator.free(value)
+        self.token_to_kv_pool_allocator.free_segment(value, start_pos=0)
         return num_tokens, num_tokens
 
     def evict(self, params: EvictParams) -> EvictResult:
@@ -660,7 +660,9 @@ class SWARadixCache(BasePrefixCache):
 
                 if len(x.children) > 0:
                     # 1. an internal node, free swa tokens.
-                    self.token_to_kv_pool_allocator.free_swa(x.value)
+                    self.token_to_kv_pool_allocator.free_swa_segment(
+                        x.value, start_pos=0
+                    )
                     swa_num_evicted += len(x.value)
 
                     # 2. get the next node, update the lru lists
@@ -673,7 +675,9 @@ class SWARadixCache(BasePrefixCache):
                     # Leaf still holds a full-side lock (can happen when the
                     # SWA leaf-lock early-release optimization revived a
                     # tombstoned leaf. Treat it like an internal tombstone.
-                    self.token_to_kv_pool_allocator.free_swa(x.value)
+                    self.token_to_kv_pool_allocator.free_swa_segment(
+                        x.value, start_pos=0
+                    )
                     swa_num_evicted += len(x.value)
 
                     x_next = self.swa_lru_list.get_prev_no_lock(x)
@@ -845,7 +849,9 @@ class SWARadixCache(BasePrefixCache):
                     # swa_lru_list so SWA-eviction won't pick this tombstoned
                     # leaf (which still holds full_lock_ref > 0). The full kv
                     # stays alive until the request releases its full lock.
-                    self.token_to_kv_pool_allocator.free_swa(node.value)
+                    self.token_to_kv_pool_allocator.free_swa_segment(
+                        node.value, start_pos=0
+                    )
                     self.swa_lru_list.remove_node(node)
                     node.swa_tombstone = True
                 else:
@@ -1202,8 +1208,8 @@ class SWARadixCache(BasePrefixCache):
                             )
                         else:
                             # Free full tokens in the original tree node.
-                            self.token_to_kv_pool_allocator.free_full(
-                                node.value[:prefix_len]
+                            self.token_to_kv_pool_allocator.free_full_segment(
+                                node.value[:prefix_len], start_pos=0
                             )
                             # Overwrite the new value in request to the tree node.
                             node.value = value[:prefix_len].clone()
@@ -1220,26 +1226,28 @@ class SWARadixCache(BasePrefixCache):
                             self._recover_tombstone_keeping_locked_full(
                                 node, value[start_update_idx:prefix_len]
                             )
-                            self.token_to_kv_pool_allocator.free_full(
-                                value[:start_update_idx]
+                            self.token_to_kv_pool_allocator.free_full_segment(
+                                value[:start_update_idx], start_pos=0
                             )
                         else:
-                            self.token_to_kv_pool_allocator.free_full(
-                                node.value[start_update_idx:prefix_len]
+                            self.token_to_kv_pool_allocator.free_full_segment(
+                                node.value[start_update_idx:prefix_len], start_pos=0
                             )
                             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_full(
-                                value[:start_update_idx]
+                            self.token_to_kv_pool_allocator.free_full_segment(
+                                value[:start_update_idx], start_pos=0
                             )
                             node.swa_tombstone = False
                             self.swa_lru_list.insert_mru(node)
                             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_full(value[:prefix_len])
+                        self.token_to_kv_pool_allocator.free_full_segment(
+                            value[:prefix_len], start_pos=0
+                        )
                 else:
                     # The node is not tombstone, so we don't need to update the node.
                     # The incoming slice can still straddle this request's own
@@ -1275,7 +1283,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_full(value)
+                self.token_to_kv_pool_allocator.free_full_segment(value, start_pos=0)
                 return total_prefix_length
 
             if (
@@ -1319,7 +1327,7 @@ class SWARadixCache(BasePrefixCache):
         swa_value = allocator.translate_loc_from_full_to_swa(incoming_full)
         allocator.set_full_to_swa_mapping(node.value, swa_value)
         allocator.clear_full_to_swa_mapping(incoming_full)
-        allocator.free_full(incoming_full)
+        allocator.free_full_segment(incoming_full, start_pos=0)
 
         node.swa_tombstone = False
         self.swa_lru_list.insert_mru(node)

python/sglang/srt/mem_cache/unified_cache/components/swa_component.py

@@ -1436,7 +1436,8 @@ class SWAComponent(TreeComponent):
         alloc = self.cache.token_to_kv_pool_allocator
         if isinstance(action, FreeComponentDeviceSlot):
             for indices in action.indices:
-                alloc.free_swa(indices)
+                # Component values are page-aligned copies of a kv row.
+                alloc.free_swa_segment(indices, start_pos=0)
             return
         if isinstance(action, FreeComponentHostSlot):
             for host_indices in action.host_indices:

test/registered/unit/mem_cache/test_swa_unittest.py

@@ -299,14 +299,20 @@ class TestSWA(unittest.TestCase):
         )
 
     def test_free_swa_group_owns_deferred_indices(self):
+        for page_size in (1, 4):
+            with self.subTest(page_size=page_size):
+                self._free_swa_group_owns_deferred_indices(page_size)
+
+    def _free_swa_group_owns_deferred_indices(self, page_size):
         _, allocator, _ = _build_swa_tree(
             is_eagle=False,
-            kv_size=32,
-            kv_size_swa=32,
+            page_size=page_size,
+            kv_size=32 * page_size,
+            kv_size_swa=32 * page_size,
         )
         index_batches = []
         for size in (2, 3, 1, 4):
-            indices = _swa_alloc(allocator, size)
+            indices = _swa_alloc(allocator, size * page_size)
             assert indices is not None
             index_batches.append(indices)
         original_indices = torch.cat([indices.clone() for indices in index_batches])
@@ -314,9 +320,10 @@ class TestSWA(unittest.TestCase):
         available_before_free = allocator.swa_available_size()
         allocator.free_group_begin()
         for indices in index_batches:
-            allocator.free_swa(indices)
+            allocator.free_swa_segment(indices, start_pos=0)
 
-        self.assertEqual(len(allocator.swa_free_group), len(index_batches))
+        # The reps were gathered at enqueue time, not from these views.
+        self.assertEqual(len(allocator.swa_page_ids_group), len(index_batches))
         self.assertEqual(allocator.swa_available_size(), available_before_free)
         for indices in index_batches:
             indices.zero_()
@@ -1144,13 +1151,57 @@ class TestSWAPeerMappedContract(CustomTestCase):
     def _strict(self):
         return envs.SGLANG_INVARIANT_CHECK.override(int(InvariantCheckLevel.STRICT))
 
-    def _condition_checked_by(self, allocator, indices):
+    def _condition_checked_by(self, allocator, indices, start_pos=None):
         """The predicate free_swa hands the async assert, as a python bool."""
         with self._strict():
             with mock.patch.object(torch, "_assert_async") as assert_async:
-                allocator.free_swa(indices)
+                if start_pos is None:
+                    allocator.free_swa(indices)
+                else:
+                    allocator.free_swa_segment(indices, start_pos=start_pos)
         return bool(assert_async.call_args.args[0])
 
+    def test_segment_free_flags_a_page_whose_peer_is_already_gone(self):
+        _, allocator, _ = _build_swa_tree(is_eagle=False, page_size=4)
+        live = _swa_alloc(allocator, 8)
+        stale = _swa_alloc(allocator, 8)
+        allocator.clear_full_to_swa_mapping(stale)
+
+        self.assertTrue(self._condition_checked_by(allocator, live, start_pos=0))
+        self.assertFalse(self._condition_checked_by(allocator, stale, start_pos=0))
+
+    @unittest.skipUnless(torch.cuda.is_available(), "sync detection needs CUDA")
+    def test_segment_free_does_not_synchronize_on_pages(self):
+        """page_size > 1: page reps by stride replace the page expansion's
+        filter and the inner allocator's torch.unique, in and out of a group."""
+        ps = 4
+        _, allocator, _ = _build_swa_tree(is_eagle=False, page_size=ps)
+
+        def grouped(indices):
+            allocator.free_group_begin()
+            allocator.free_swa_segment(indices, start_pos=0)
+            allocator.free_group_end()
+
+        # Warm up both paths outside the window: a first-time cudaMalloc can
+        # synchronize on its own, which the detector would blame on this call.
+        allocator.free_swa_segment(_swa_alloc(allocator, 2 * ps), start_pos=0)
+        grouped(_swa_alloc(allocator, 2 * ps))
+        first = _swa_alloc(allocator, 3 * ps)
+        second = _swa_alloc(allocator, 2 * ps)
+
+        # Gate on the pre-fix form: a detector blind to this sync class would pass
+        # the asserts below no matter how free_swa derives the pages.
+        if _sync_error(lambda: torch.unique(first // ps)) is None:
+            self.skipTest("sync debug mode does not flag a data-dependent shape here")
+
+        with self._strict():
+            self.assertIsNone(
+                _sync_error(
+                    lambda: allocator.free_swa_segment(first[: 3 * ps - 1], start_pos=0)
+                )
+            )
+            self.assertIsNone(_sync_error(lambda: grouped(second[: 2 * ps - 1])))
+
     def test_free_swa_flags_a_slot_whose_peer_is_already_gone(self):
         _, allocator, _ = _build_swa_tree(is_eagle=False)
         live = _swa_alloc(allocator, 4)
@@ -1183,6 +1234,69 @@ class TestSWAPeerMappedContract(CustomTestCase):
             self.assertIsNone(_sync_error(lambda: allocator.free_swa(indices)))
 
 
+class TestSWAPageRepsFree(CustomTestCase):
+    """page_size > 1: with a start position the SWA side frees one representative
+    per page instead of expanding, filtering and dedup'ing through torch.unique."""
+
+    PS = 4
+
+    def _allocator(self):
+        _, allocator, _ = _build_swa_tree(is_eagle=False, page_size=self.PS)
+        return allocator
+
+    def _sizes(self, allocator):
+        return allocator.full_available_size(), allocator.swa_available_size()
+
+    def test_segment_free_releases_the_mapped_pages_for_every_tail(self):
+        ps = self.PS
+        for num_tokens in (1, ps, ps + 1, 3 * ps - 1, 3 * ps):
+            with self.subTest(num_tokens=num_tokens):
+                allocator = self._allocator()
+                indices = _swa_alloc(allocator, 3 * ps)
+                mapping = allocator.full_to_swa_index_mapping
+                expected = torch.unique(mapping[indices[:num_tokens]] // ps)
+                before = allocator.swa_attn_allocator.free_pages.numel()
+
+                allocator.free_swa_segment(indices[:num_tokens], start_pos=0)
+
+                free_pages = allocator.swa_attn_allocator.free_pages
+                freed = free_pages[: free_pages.numel() - before]
+                self.assertTrue(torch.equal(torch.sort(freed)[0], expected))
+                # The whole last page goes back, and its mapping with it.
+                touched = -(num_tokens // -ps) * ps
+                self.assertTrue(torch.all(mapping[indices[:touched]] == 0))
+                self.assertTrue(torch.all(mapping[indices[touched:]] > 0))
+
+    def test_node_frees_take_the_page_path_through_the_tree(self):
+        """Tree values are page-aligned copies of a kv row, so SWA eviction and
+        the full eviction of its tombstones both free by page reps."""
+        ps = self.PS
+        tree, allocator, _ = _build_swa_tree(
+            is_eagle=False, page_size=ps, sliding_window_size=ps
+        )
+        full_before, swa_before = self._sizes(allocator)
+        _insert(tree, allocator, list(range(1, 3 * ps + 1)))
+
+        # Either inner `free` is the torch.unique path a caller falls back to
+        # when it hands no start position.
+        with (
+            patch.object(
+                allocator.full_attn_allocator,
+                "free",
+                side_effect=AssertionError("full side took the unique path"),
+            ),
+            patch.object(
+                allocator.swa_attn_allocator,
+                "free",
+                side_effect=AssertionError("swa side took the unique path"),
+            ),
+        ):
+            tree.evict(EvictParams(num_tokens=0, swa_num_tokens=ps))
+            tree.evict(EvictParams(num_tokens=3 * ps, swa_num_tokens=0))
+
+        self.assertEqual(self._sizes(allocator), (full_before, swa_before))
+
+
 class TestCacheUnfinishedReqEvictedPrefix(CustomTestCase):
     """An unfinished request whose SWA prefix is already gone must insert that
     prefix as a tombstone, not as live SWA KV."""

test/registered/unit/mem_cache/test_unified_free_no_host_sync.py

@@ -373,7 +373,7 @@ class TestUnifiedSwaFullSideGroup(unittest.TestCase):
 
 class TestFreeSwaWindowRatchetNoHostSync(unittest.TestCase):
     """The per-decode-step SWA window ratchet frees a CONTIGUOUS row slice with
-    host-int, page-aligned bounds, so `free_swa(..., start_pos=)` must reach the
+    host-int, page-aligned bounds, so `free_swa_segment` must reach the
     swa side with caller-derived page ids: no `torch.unique` and no stale-slot
     `.item()` on the per-step path.
     """
@@ -429,15 +429,15 @@ class TestFreeSwaWindowRatchetNoHostSync(unittest.TestCase):
                 torch.Tensor, "item", side_effect=AssertionError("item = host sync")
             ),
         ):
-            alloc.free_swa(v[: 4 * self.PS], start_pos=0)
-            alloc.free_swa(v[4 * self.PS :], start_pos=4 * self.PS)
+            alloc.free_swa_segment(v[: 4 * self.PS], start_pos=0)
+            alloc.free_swa_segment(v[4 * self.PS :], start_pos=4 * self.PS)
 
     def test_full_only_segment_free_never_syncs(self):
         """Request-finish shape: the swa side is already tombstoned, so the
         full side must free by page reps rather than `free_full`'s dedup."""
         alloc = self._swa_composite(lazy=True)
         v = alloc.alloc(8 * self.PS)
-        alloc.free_swa(v, start_pos=0)
+        alloc.free_swa_segment(v, start_pos=0)
         before = alloc.full_available_size()
         with (
             mock.patch.object(
@@ -456,7 +456,7 @@ class TestFreeSwaWindowRatchetNoHostSync(unittest.TestCase):
         alloc = self._swa_composite(lazy=True)
         v = alloc.alloc(8 * self.PS)
         with self.assertRaises(AssertionError):
-            alloc.free_swa(v[1 : 5 * self.PS], start_pos=1)
+            alloc.free_swa_segment(v[1 : 5 * self.PS], start_pos=1)
 
     def test_start_pos_path_matches_the_fallback_end_state(self):
         """Derived property: the stride-rep path and the dedup fallback leave
@@ -468,7 +468,7 @@ class TestFreeSwaWindowRatchetNoHostSync(unittest.TestCase):
                 v1 = a1.alloc(6 * self.PS)
                 v2 = a2.alloc(6 * self.PS)
                 self.assertTrue(torch.equal(v1, v2))
-                a1.free_swa(v1[: 4 * self.PS], start_pos=0)
+                a1.free_swa_segment(v1[: 4 * self.PS], start_pos=0)
                 a2.free_swa(v2[: 4 * self.PS])  # fallback (radix shape)
                 self.assertTrue(
                     torch.equal(
@@ -487,8 +487,8 @@ class TestFreeSwaWindowRatchetNoHostSync(unittest.TestCase):
         liveness filter (radix eviction and the ratchet can overlap)."""
         alloc = self._swa_composite(lazy=True)
         v = alloc.alloc(4 * self.PS)
-        alloc.free_swa(v, start_pos=0)
-        alloc.free_swa(v, start_pos=0)  # all tombstoned -> filtered to empty
+        alloc.free_swa_segment(v, start_pos=0)
+        alloc.free_swa_segment(v, start_pos=0)  # all tombstoned -> filtered to empty
 
 
 @unittest.skipUnless(

test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py

@@ -7860,13 +7860,15 @@ class TestUnifiedRadixCacheActionRouting(CustomTestCase):
             indices, start_pos=0
         )
 
-    def test_apply_component_action_device_kv_swa_uses_free_swa(self):
+    def test_apply_component_action_device_kv_swa_uses_free_swa_segment(self):
         cache = mock.MagicMock()
         indices = torch.tensor([4, 5])
         _component_with_cache(ComponentType.SWA, cache).apply_component_action(
             FreeComponentDeviceSlot([indices], component_type=ComponentType.SWA)
         )
-        cache.token_to_kv_pool_allocator.free_swa.assert_called_once_with(indices)
+        cache.token_to_kv_pool_allocator.free_swa_segment.assert_called_once_with(
+            indices, start_pos=0
+        )
 
     def test_apply_component_action_device_kv_mamba_uses_mamba_allocator(self):
         cache = mock.MagicMock()

test/registered/unit/mem_cache/test_unified_tri_pool.py

@@ -407,7 +407,7 @@ class TestTriFreeSwaNoHostSync(unittest.TestCase):
                 torch.Tensor, "item", side_effect=AssertionError("item = host sync")
             ),
         ):
-            alloc.free_swa(v[: 4 * self.PS], start_pos=0)
+            alloc.free_swa_segment(v[: 4 * self.PS], start_pos=0)
         self.assertEqual(alloc.verify_byte_accounting(), [])
 
     def test_fallback_free_swa_still_correct_for_radix_shapes(self):
@@ -416,7 +416,7 @@ class TestTriFreeSwaNoHostSync(unittest.TestCase):
         a1, a2 = self._tri(), self._tri()
         v1, v2 = a1.alloc(6 * self.PS), a2.alloc(6 * self.PS)
         self.assertTrue(torch.equal(v1, v2))
-        a1.free_swa(v1[: 4 * self.PS], start_pos=0)
+        a1.free_swa_segment(v1[: 4 * self.PS], start_pos=0)
         a2.free_swa(v2[: 4 * self.PS])
         self.assertTrue(
             torch.equal(
@@ -715,7 +715,7 @@ class TestTriDeferredAbsorption(unittest.TestCase):
         v = alloc.alloc(8 * self.PS)
         sa = alloc.swa_attn_allocator
         span = sa._span_pages()
-        alloc.free_swa(v[6 * self.PS :], start_pos=6 * self.PS)  # high edge
+        alloc.free_swa_segment(v[6 * self.PS :], start_pos=6 * self.PS)  # high edge
         self.assertGreater(sa._hole_pages(), 0)  # deferred
         self.assertEqual(sa._span_pages(), span)
         moved = alloc.flush_opportunistic()
@@ -729,7 +729,7 @@ class TestTriDeferredAbsorption(unittest.TestCase):
         alloc = self._tri()
         v = alloc.alloc(8 * self.PS)
         sa = alloc.swa_attn_allocator
-        alloc.free_swa(v[6 * self.PS :], start_pos=6 * self.PS)
+        alloc.free_swa_segment(v[6 * self.PS :], start_pos=6 * self.PS)
         self.assertGreater(sa._hole_pages(), 0)
         moves_before = len(sa._inverse_history)
         from sglang.srt.mem_cache.allocator.unified_sub_pool import _relieve_for_alloc
@@ -743,7 +743,7 @@ class TestTriDeferredAbsorption(unittest.TestCase):
         value -- under-reporting is safe, over-reporting would over-admit."""
         alloc = self._tri()
         v = alloc.alloc(8 * self.PS)
-        alloc.free_swa(v[6 * self.PS :], start_pos=6 * self.PS)
+        alloc.free_swa_segment(v[6 * self.PS :], start_pos=6 * self.PS)
         deferred = alloc.available_size()
         alloc.swa_attn_allocator._flush(urgent=False)
         absorbed = alloc.available_size()
@@ -759,7 +759,7 @@ class TestTriDeferredAbsorption(unittest.TestCase):
 
         alloc = self._tri()
         v = alloc.alloc(8 * self.PS)
-        alloc.free_swa(v[2 * self.PS : 4 * self.PS], start_pos=2 * self.PS)
+        alloc.free_swa_segment(v[2 * self.PS : 4 * self.PS], start_pos=2 * self.PS)
         alloc.flush_opportunistic()  # consumes the dirty flag
         sa = alloc.swa_attn_allocator
         self.assertGreater(sa._hole_pages(), 0)  # interior holes remain
@@ -775,10 +775,10 @@ class TestTriDeferredAbsorption(unittest.TestCase):
         alloc = self._tri()
         v = alloc.alloc(8 * self.PS)
         sa = alloc.swa_attn_allocator
-        alloc.free_swa(v[: 2 * self.PS], start_pos=0)  # low-edge holes
+        alloc.free_swa_segment(v[: 2 * self.PS], start_pos=0)  # low-edge holes
         n_after_free = sa._hole_pages()
         alloc.alloc(2 * self.PS)  # drains them back to live
-        alloc.free_swa(v[6 * self.PS :], start_pos=6 * self.PS)  # high edge
+        alloc.free_swa_segment(v[6 * self.PS :], start_pos=6 * self.PS)  # high edge
         self.assertEqual(sa._hole_pages(), n_after_free)  # same COUNT as before
         span = sa._span_pages()
         self.assertGreater(alloc.flush_opportunistic(), 0)  # still absorbed
@@ -797,7 +797,7 @@ class TestTriDeferredAbsorption(unittest.TestCase):
         with mock.patch.object(
             torch.Tensor, "tolist", side_effect=AssertionError("tolist = D2H")
         ):
-            alloc.free_swa(v, start_pos=0)
+            alloc.free_swa_segment(v, start_pos=0)
         self.assertTrue(sa._is_frontier_transparent())
         self.assertEqual(sa._hole_pages(), 0)