[mem_cache] Route hybrid SWA full-side kv-row frees through free_segment (#37876)

Co-authored-by: weireweire <20922698+weireweire@users.noreply.github.com>
Co-authored-by: Sam Shleifer <sshleifer@gmail.com>

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

@@ -192,6 +192,23 @@ class BaseTokenToKVPoolAllocator(abc.ABC):
         Starts sit on page boundaries, ends may fall mid-page, and the page
         ranges of consecutive segments do not overlap -- so in page units the
         segments are aligned and disjoint, and every page is released once."""
+        for free_index, start_pos in self._page_disjoint(segments):
+            self.free_segment(free_index, start_pos=start_pos)
+
+    def free_full_segment(self, free_index: torch.Tensor, *, start_pos: int):
+        """free_full() for a kv-row segment; same start-alignment contract as
+        free_segment(). Default: plain free_full()."""
+        assert start_pos % self.page_size == 0, (
+            f"segment start {start_pos} is not page-aligned"
+        )
+        self.free_full(free_index)
+
+    def free_full_segments(self, segments):
+        """free_segments() for the full side alone; see free_full()."""
+        for free_index, start_pos in self._page_disjoint(segments):
+            self.free_full_segment(free_index, start_pos=start_pos)
+
+    def _page_disjoint(self, segments):
         ps = self.page_size
         prev_end = None
         for free_index, start_pos in segments:
@@ -202,4 +219,4 @@ class BaseTokenToKVPoolAllocator(abc.ABC):
                 f"segment at {start_pos} shares a page with the one ending at {prev_end}"
             )
             prev_end = start_pos + n
-            self.free_segment(free_index, start_pos=start_pos)
+            yield free_index, start_pos

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

@@ -103,7 +103,6 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
         self.release_pages = None
         self.free_group = None
         self.swa_free_group = []
-        self.full_free_group = []
 
         self._kvcache = kvcache
         self.clear()
@@ -402,18 +401,34 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
             self.full_to_swa_index_mapping[free_index] == 0,
             msg="caller wants free",
         )
-        if self.free_group is None:
-            self.full_attn_allocator.free(free_index)
-        else:
-            self.full_free_group.append(self._copy_for_free_group(free_index))
+        self.full_attn_allocator.free(free_index)
         assert (
             self.full_attn_allocator.available_size() <= self.full_attn_allocator.size
         )
 
+    def free_segment(self, free_index: torch.Tensor, *, start_pos: int):
+        if free_index.numel() == 0:
+            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.full_attn_allocator.free_segment(free_index, start_pos=start_pos)
+
+    def free_full_segment(self, free_index: torch.Tensor, *, start_pos: int):
+        if free_index.numel() == 0:
+            return
+        expect(
+            _SWA_PEER_RELEASED,
+            self.full_to_swa_index_mapping[free_index] == 0,
+            msg="caller wants free_segment",
+        )
+        self.full_attn_allocator.free_segment(free_index, start_pos=start_pos)
+
     def free_group_begin(self):
         super().free_group_begin()
         self.swa_free_group = []
-        self.full_free_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()
@@ -421,10 +436,7 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
             swa_free_group = self.swa_free_group
             self.swa_free_group = []
             self._release_swa(torch.cat(swa_free_group))
-        if self.full_free_group:
-            full_free_group = self.full_free_group
-            self.full_free_group = []
-            self.full_attn_allocator.free(torch.cat(full_free_group))
+        self.full_attn_allocator.free_group_end()
         assert (
             self.full_attn_allocator.available_size() <= self.full_attn_allocator.size
         )
@@ -468,7 +480,6 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
         self.full_to_swa_index_mapping[:-1].fill_(0)
         self.free_group = None
         self.swa_free_group = []
-        self.full_free_group = []
 
     def get_cpu_copy(self, indices, mamba_indices=None):
         return self._kvcache.get_cpu_copy(indices, mamba_indices=mamba_indices)
@@ -594,8 +605,16 @@ class PureSWATokenToKVPoolAllocator(SWATokenToKVPoolAllocator):
         # release once the SWA side is gone.
         return
 
-    # Not inherited: the SWA parent's hooks drive swa_free_group and
-    # full_free_group, which this pure-SWA variant does not have.
+    def free_segment(self, free_index: torch.Tensor, *, start_pos: int):
+        # Single pool: the parent's split into an SWA and a full half would
+        # release the same slots twice.
+        self.free(free_index)
+
+    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
+    # 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/common.py

@@ -127,7 +127,7 @@ def free_kv_row_segments(
 ) -> None:
     """Free ascending disjoint ``(kv_indices, start_pos)`` segments of one
     request's kv row, split at the SWA eviction floor."""
-    swa_dead: list[torch.Tensor] = []
+    swa_dead: list[tuple[torch.Tensor, int]] = []
     swa_alive: list[tuple[torch.Tensor, int]] = []
     for kv_indices, start_pos in segments:
         num_indices = kv_indices.numel()
@@ -137,23 +137,19 @@ def free_kv_row_segments(
         # the deliberately unmapped prefix of a PD decode SWA-tail prealloc.
         num_dead = min(max(swa_evicted_seqlen - start_pos, 0), num_indices)
         if num_dead > 0:
-            swa_dead.append(kv_indices[:num_dead])
+            swa_dead.append((kv_indices[:num_dead], start_pos))
         if num_dead < num_indices:
             swa_alive.append((kv_indices[num_dead:], start_pos + num_dead))
 
     if swa_dead and swa_alive:
-        # A mid-page floor would send a page shared by the dead and alive
-        # sides back twice.
+        # The two sides are separate calls, so neither one's page-disjointness
+        # check sees a floor that splits a page between them.
         assert swa_evicted_seqlen % allocator.page_size == 0, (
             f"SWA eviction floor {swa_evicted_seqlen} splits a page "
             f"(page_size {allocator.page_size})"
         )
-    if len(swa_dead) == 1:
-        allocator.free_full(swa_dead[0])
-    elif swa_dead:
-        # Two dead pieces can share a boundary page, and only free_full's own
-        # page dedup covers that -- free_segments trims the alive side alone.
-        allocator.free_full(torch.cat(swa_dead))
+    if swa_dead:
+        allocator.free_full_segments(swa_dead)
     if swa_alive:
         allocator.free_segments(swa_alive)
 

python/sglang/srt/mem_cache/multi_ended_allocator.py

@@ -3227,6 +3227,7 @@ class UnifiedSWATokenToKVPoolAllocator(SWATokenToKVPoolAllocator):
 
         self.free_group = None
         self.free_page_reps_group: Optional[List[torch.Tensor]] = None
+        self.full_free_group: List[torch.Tensor] = []
         # Empty (not None) for the leak checker.
         self.free_pages = torch.empty(0, dtype=torch.int64, device=device)
         self.release_pages = torch.empty(0, dtype=torch.int64, device=device)
@@ -3661,6 +3662,17 @@ class UnifiedSWATokenToKVPoolAllocator(SWATokenToKVPoolAllocator):
         self.full_attn_allocator.free(free_index.detach().to(torch.int64))
         self.full_attn_allocator.clear_inverse_history()
 
+    def free_full_segment(self, free_index: torch.Tensor, *, start_pos: int) -> None:
+        if free_index is None or free_index.numel() == 0:
+            return
+        if self.page_size == 1:
+            # token == page: free_full already frees by exact ids, no dedup.
+            self.free_full(free_index)
+            return
+        # The swa v2p is the mapping, so a tombstoned swa page drops out of the
+        # two-sided segment path by itself; full-only is the same call.
+        self.free_segment(free_index, start_pos=start_pos)
+
     def set_full_to_swa_mapping(
         self, full_indices: torch.Tensor, swa_indices: torch.Tensor
     ) -> None:
@@ -3676,13 +3688,20 @@ class UnifiedSWATokenToKVPoolAllocator(SWATokenToKVPoolAllocator):
 
     # -- free-group --
 
+    # Not the SWA parent's hooks: those open the parent's paged full allocator
+    # as a free group, and this composite's sub-pools defer on their own.
     def free_group_begin(self) -> None:
-        super().free_group_begin()
+        BaseTokenToKVPoolAllocator.free_group_begin(self)
         self.free_page_reps_group = []
+        self.full_free_group = []
 
     def free_group_end(self) -> None:
         pending, self.free_page_reps_group = self.free_page_reps_group, None
-        super().free_group_end()
+        full_free_group, self.full_free_group = self.full_free_group, []
+        BaseTokenToKVPoolAllocator.free_group_end(self)
+        if full_free_group:
+            self.full_attn_allocator.free(torch.cat(full_free_group))
+            self.full_attn_allocator.clear_inverse_history()
         if pending:
             self._release_page_reps(pending)
 
@@ -3746,6 +3765,7 @@ class UnifiedSWATokenToKVPoolAllocator(SWATokenToKVPoolAllocator):
         self.swa_attn_allocator.clear()
         self.free_group = None
         self.free_page_reps_group = None
+        self.full_free_group = []
 
     # -- Lazy compaction hooks --
 

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

@@ -503,10 +503,11 @@ class FullComponent(TreeComponent):
         if isinstance(action, FreeComponentDeviceSlot):
             alloc = self.cache.token_to_kv_pool_allocator
             for indices in action.indices:
+                # tree values are page-aligned copies of a kv row: page-exact segments
                 if self.cache.is_swa_enabled:
-                    alloc.full_attn_allocator.free(indices)
+                    alloc.full_attn_allocator.free_segment(indices, start_pos=0)
                 else:
-                    alloc.free(indices)
+                    alloc.free_segment(indices, start_pos=0)
             return
         raise AssertionError(
             f"FullComponent: unhandled ComponentAction {type(action).__name__}"

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

@@ -1371,7 +1371,7 @@ class SWAComponent(TreeComponent):
             swa_value = self._translate_full_to_swa(action.incoming_full)
             alloc.set_full_to_swa_mapping(action.kept_full, swa_value)
             alloc.clear_full_to_swa_mapping(action.incoming_full)
-            alloc.free_full(action.incoming_full)
+            alloc.free_full_segment(action.incoming_full, start_pos=0)
             self.tree_core.set_component_device_value(
                 action.node_id, self.component_type, swa_value
             )

python/sglang/srt/mem_cache/unified_radix_cache.py

@@ -1097,7 +1097,7 @@ class UnifiedRadixCache(BasePrefixCache):
                 self.token_to_kv_pool_allocator.free_segment(indices, start_pos=0)
         elif isinstance(action, FreeDeviceKVFullOnly):
             for indices in action.indices:
-                self.token_to_kv_pool_allocator.free_full(indices)
+                self.token_to_kv_pool_allocator.free_full_segment(indices, start_pos=0)
         elif isinstance(action, BackupKV):
             if self.linker is not None:
                 self.linker.offload_nodes(action.node_ids)