[mem_cache] Add free_full to release the full side of a tombstoned SWA node (#36637)

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

@@ -123,6 +123,14 @@ class BaseTokenToKVPoolAllocator(abc.ABC):
     def free(self, free_index: torch.Tensor):
         raise NotImplementedError()
 
+    def free_full(self, free_index: torch.Tensor):
+        """Free slots whose SWA peers the caller already released.
+
+        A hybrid SWA allocator pairs each full-attention slot with an SWA slot
+        that can die first; this releases the full side alone. A single pool has
+        no peer, so it is a plain free()."""
+        self.free(free_index)
+
     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

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

@@ -316,6 +316,7 @@ class DeepSeekV4HiSparseTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
         self.free_pages = None
         self.release_pages = None
         self.free_group = None
+        self.full_free_group = []
         self.clear()
 
         self.hisparse_kvcache.register_mapping(
@@ -364,6 +365,15 @@ class DeepSeekV4HiSparseTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
     def free_swa(self, free_indices: torch.Tensor):
         self.logical_attn_allocator.free_swa(free_indices)
 
+    def free_full(self, free_indices: torch.Tensor):
+        if free_indices.numel() == 0:
+            return
+
+        if self.free_group is None:
+            self.logical_attn_allocator.free_full(free_indices)
+        else:
+            self.full_free_group.append(self._copy_for_free_group(free_indices))
+
     def available_size(self) -> int:
         return min(
             self.logical_attn_allocator.available_size(),
@@ -567,6 +577,7 @@ class DeepSeekV4HiSparseTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
 
         self.full_to_hisparse_device_index_mapping[:-1].fill_(0)
         self.free_group = None
+        self.full_free_group = []
 
     def free(self, free_index: torch.Tensor):
         if free_index.numel() == 0:
@@ -576,3 +587,14 @@ class DeepSeekV4HiSparseTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
             self.logical_attn_allocator.free(free_index)
         else:
             self.free_group.append(self._copy_for_free_group(free_index))
+
+    def free_group_begin(self):
+        super().free_group_begin()
+        self.full_free_group = []
+
+    def free_group_end(self):
+        super().free_group_end()
+        if self.full_free_group:
+            full_free_group = self.full_free_group
+            self.full_free_group = []
+            self.free_full(torch.cat(full_free_group))

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

@@ -95,6 +95,7 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
         self.release_pages = None
         self.free_group = None
         self.swa_free_group = []
+        self.full_free_group = []
 
         self._kvcache = kvcache
         self.clear()
@@ -370,9 +371,24 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
 
         self.swa_attn_allocator.free(swa_indices)
 
+    def free_full(self, free_index: torch.Tensor):
+        if free_index.numel() == 0:
+            return
+
+        if self.free_group is None:
+            # Full side only: a tombstoned range's mapping entries read as the
+            # padding slot, so `free` would push slot 0 into the SWA free list.
+            self.full_attn_allocator.free(free_index)
+        else:
+            self.full_free_group.append(self._copy_for_free_group(free_index))
+        assert (
+            self.full_attn_allocator.available_size() <= self.full_attn_allocator.size
+        )
+
     def free_group_begin(self):
         super().free_group_begin()
         self.swa_free_group = []
+        self.full_free_group = []
 
     def free_group_end(self):
         super().free_group_end()
@@ -380,6 +396,10 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
             swa_free_group = self.swa_free_group
             self.swa_free_group = []
             self.swa_attn_allocator.free(torch.cat(swa_free_group))
+        if self.full_free_group:
+            full_free_group = self.full_free_group
+            self.full_free_group = []
+            self.free_full(torch.cat(full_free_group))
 
     def _expand_to_full_pages(self, indices: torch.Tensor) -> torch.Tensor:
         pages = torch.unique(indices // self.page_size)
@@ -409,6 +429,7 @@ 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)
@@ -516,8 +537,13 @@ class PureSWATokenToKVPoolAllocator(SWATokenToKVPoolAllocator):
         else:
             self.free_group.append(self._copy_for_free_group(free_index))
 
-    # Not inherited: the SWA parent's hooks drive swa_free_group,
-    # which this pure-SWA variant does not have.
+    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.
+        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_group_begin(self):
         self.free_group = []
 

python/sglang/srt/mem_cache/multi_ended_allocator.py

@@ -2448,6 +2448,17 @@ class UnifiedSWATokenToKVPoolAllocator(SWATokenToKVPoolAllocator):
         self.swa_attn_allocator.free(live)
         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`)."""
+        if free_index is None or free_index.numel() == 0:
+            return
+        if self.free_group is not None:
+            self.full_free_group.append(self._copy_for_free_group(free_index))
+            return
+        self.full_attn_allocator.free(free_index.detach().to(torch.int64))
+        self.full_attn_allocator.clear_inverse_history()
+
     def set_full_to_swa_mapping(
         self, full_indices: torch.Tensor, swa_indices: torch.Tensor
     ) -> None:

python/sglang/srt/mem_cache/swa_radix_cache.py

@@ -591,6 +591,21 @@ class SWARadixCache(BasePrefixCache):
     def total_size(self) -> Tuple[int, int]:
         return self._total_size_helper()
 
+    def _free_node_value(
+        self, node: TreeNode, value: Optional[torch.Tensor] = None
+    ) -> Tuple[int, int]:
+        if value is None:
+            value = node.value
+        num_tokens = len(value)
+        if node.swa_tombstone:
+            # 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)
+            return num_tokens, 0
+        self.token_to_kv_pool_allocator.free(value)
+        return num_tokens, num_tokens
+
     def evict(self, params: EvictParams) -> EvictResult:
         if self.disable:
             return EvictResult()
@@ -611,11 +626,9 @@ class SWARadixCache(BasePrefixCache):
 
                 # 1. free node kv indices, evict full and swa tokens
                 self.kv_events.record_remove(x)
-                self.token_to_kv_pool_allocator.free(x.value)
-                full_num_evicted += len(x.value)
-                # Tombstoned leaves had their SWA freed earlier in `dec_swa_lock_only`
-                if not x.swa_tombstone:
-                    swa_num_evicted += len(x.value)
+                node_full_evicted, node_swa_evicted = self._free_node_value(x)
+                full_num_evicted += node_full_evicted
+                swa_num_evicted += node_swa_evicted
 
                 # 2. get the next leaf, update the lru lists
                 x_next = self.full_lru_list.get_prev_leaf_no_lock(x)
@@ -676,9 +689,9 @@ class SWARadixCache(BasePrefixCache):
                     ), f"leaf node with full lock must also have swa lock, {x.id=}"
                     # 1. a leaf node, free full and swa tokens
                     self.kv_events.record_remove(x)
-                    self.token_to_kv_pool_allocator.free(x.value)
-                    full_num_evicted += len(x.value)
-                    swa_num_evicted += len(x.value)
+                    node_full_evicted, node_swa_evicted = self._free_node_value(x)
+                    full_num_evicted += node_full_evicted
+                    swa_num_evicted += node_swa_evicted
 
                     # 2. get the next node, update the lru lists
                     x_next = self.swa_lru_list.get_prev_no_lock(x)
@@ -1302,7 +1315,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.full_attn_allocator.free(incoming_full)
+        allocator.free_full(incoming_full)
 
         node.swa_tombstone = False
         self.swa_lru_list.insert_mru(node)
@@ -1346,8 +1359,8 @@ class SWARadixCache(BasePrefixCache):
             ), f"tombstone swa_lock_ref should always be 0, {node.parent.full_lock_ref=}, {node.parent.swa_lock_ref=}, {node.parent.id=}"
             # delete tombstone node evicts full tokens
             self.kv_events.record_remove(node.parent)
-            self.token_to_kv_pool_allocator.free(node.parent.value)
-            full_num_evicted += len(node.parent.value)
+            node_full_evicted, _ = self._free_node_value(node.parent)
+            full_num_evicted += node_full_evicted
             self.full_lru_list.remove_node(node.parent)
             self._delete_tombstone_leaf(node.parent)
             node = node.parent

python/sglang/srt/mem_cache/unified_cache/cache_action.py

@@ -36,6 +36,13 @@ class FreeDeviceKV(msgspec.Struct, frozen=True):
     indices: list[torch.Tensor]
 
 
+class FreeDeviceKVFullOnly(msgspec.Struct, frozen=True):
+    """Free the full side only, for a tombstoned node whose SWA peers are gone;
+    FreeDeviceKV would release the SWA side twice."""
+
+    indices: list[torch.Tensor]
+
+
 class ComponentAction(msgspec.Struct, frozen=True):
     """Base for component-routed actions; the cache dispatches each one to
     ``component_type``'s class-level ``apply_component_action``; every subclass
@@ -100,4 +107,6 @@ class SWARebuild(ComponentAction, frozen=True):
 
 
 # Cache-owned actions, applied by UnifiedRadixCache itself.
-CacheAction = ReplaceWriteThroughOnNodeSplit | FreeDeviceKV | BackupKV
+CacheAction = (
+    ReplaceWriteThroughOnNodeSplit | FreeDeviceKV | FreeDeviceKVFullOnly | BackupKV
+)

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

@@ -22,7 +22,7 @@ from sglang.srt.mem_cache.hicache_storage import (
 from sglang.srt.mem_cache.unified_cache.cache_action import (
     FreeComponentDeviceSlot,
     FreeComponentHostSlot,
-    FreeDeviceKV,
+    FreeDeviceKVFullOnly,
     RebuildFullToSWAMapping,
     RecoverSWAWithLockedFull,
     SWARebuild,
@@ -295,7 +295,7 @@ class SWAComponent(TreeComponent):
                 )
                 return 0
             full_cd.value = value_slice.clone()
-            cache_actions.append(FreeDeviceKV([old_full]))
+            cache_actions.append(FreeDeviceKVFullOnly([old_full]))
             cache_actions.append(SWARebuild(node.id, value_slice))
             return 0
         elif swa_evicted_seqlen < total_prefix_len + prefix_len:
@@ -313,7 +313,7 @@ class SWAComponent(TreeComponent):
                 )
                 return start_idx
             node.component_data[BASE_COMPONENT_TYPE].value = new_full.clone()
-            cache_actions.append(FreeDeviceKV([old_full]))
+            cache_actions.append(FreeDeviceKVFullOnly([old_full]))
             cache_actions.append(SWARebuild(node.id, new_full))
             return start_idx
         else:
@@ -1156,7 +1156,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.full_attn_allocator.free(action.incoming_full)
+            alloc.free_full(action.incoming_full)
             self.tree_core.set_component_device_value(
                 action.node_id, self.component_type, swa_value
             )

python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py

@@ -48,6 +48,7 @@ from sglang.srt.mem_cache.unified_cache.cache_action import (
     CacheAction,
     ComponentAction,
     FreeDeviceKV,
+    FreeDeviceKVFullOnly,
     ReplaceWriteThroughOnNodeSplit,
 )
 from sglang.srt.mem_cache.unified_cache.components import (
@@ -939,7 +940,10 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
     @staticmethod
     def _is_deferrable_action(action: CacheAction | ComponentAction) -> bool:
         """Fire-and-forget actions safe to batch until the next barrier."""
-        return isinstance(action, (FreeDeviceKV, ReplaceWriteThroughOnNodeSplit))
+        return isinstance(
+            action,
+            (FreeDeviceKV, FreeDeviceKVFullOnly, ReplaceWriteThroughOnNodeSplit),
+        )
 
     def _insert_walk_step(self, state: _InsertWalkState) -> None:
         """Process one walked node, appending its barrier actions to the state."""

python/sglang/srt/mem_cache/unified_radix_cache.py

@@ -51,6 +51,7 @@ from sglang.srt.mem_cache.unified_cache.cache_action import (
     ComponentAction,
     FreeComponentDeviceSlot,
     FreeDeviceKV,
+    FreeDeviceKVFullOnly,
     ReplaceWriteThroughOnNodeSplit,
 )
 
@@ -1052,6 +1053,9 @@ class UnifiedRadixCache(BasePrefixCache):
             # 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_segment(indices, start_pos=0)
+        elif isinstance(action, FreeDeviceKVFullOnly):
+            for indices in action.indices:
+                self.token_to_kv_pool_allocator.free_full(indices)
         elif isinstance(action, BackupKV):
             self._execute_and_commit_kv_backup(action)
         else:

test/registered/unit/mem_cache/test_multi_ended_allocator.py

@@ -744,6 +744,33 @@ class TestUnifiedSWATokenToKVPoolAllocator(unittest.TestCase):
         )
         self.assertIn(tgt, free_full)
 
+    def test_swa_free_full_defers_inside_a_free_group(self):
+        """The full-only release joins the barrier, like `free`."""
+        _, allocator, kvcache = self._build()
+        v = self._alloc(allocator, kvcache, 3)
+        target = v[1:2]
+        tgt = int(target.item())
+        # Tombstone the swa side, erasing each marker before its release
+        # (compaction runs inside both).
+        target_swa = allocator.swa_attn_allocator.virtual_to_physical[target]
+        kvcache.swa_kv_pool.buf[target_swa] = -1
+        allocator.free_swa(target)
+        full_phys = int(allocator.full_attn_allocator.virtual_to_physical[tgt].item())
+        kvcache.full_kv_pool.buf[full_phys] = -1
+
+        allocator.free_group_begin()
+        allocator.free_full(target)
+        deferred = set(
+            int(x) for x in allocator.full_attn_allocator.free_virtual_ids.tolist()
+        )
+        self.assertNotIn(tgt, deferred)
+
+        allocator.free_group_end()
+        drained = set(
+            int(x) for x in allocator.full_attn_allocator.free_virtual_ids.tolist()
+        )
+        self.assertIn(tgt, drained)
+
     # 4. Compaction diverges between the two sub-pools (each runs its own).
     def test_swa_compaction_diverges_physical_layout(self):
         _, allocator, kvcache = self._build()

test/registered/unit/mem_cache/test_swa_unittest.py

@@ -851,5 +851,48 @@ class TestSWASplitLeafOnInsert(CustomTestCase):
         tree.sanity_check()
 
 
+class TestFreeFullPartition(CustomTestCase):
+    """`free_full` releases only the full side of a hybrid SWA allocator."""
+
+    def setUp(self):
+        _, self.allocator, _ = _build_swa_tree(is_eagle=False)
+        self.full_baseline = self.allocator.full_available_size()
+        self.swa_baseline = self.allocator.swa_available_size()
+
+    def _sizes(self):
+        return (
+            self.allocator.full_available_size(),
+            self.allocator.swa_available_size(),
+        )
+
+    def test_free_full_keeps_the_swa_peers_allocated(self):
+        indices = _swa_alloc(self.allocator, 4)
+        self.allocator.free_full(indices)
+
+        full_avail, swa_avail = self._sizes()
+        self.assertEqual(full_avail, self.full_baseline)
+        self.assertEqual(swa_avail, self.swa_baseline - 4)
+
+    def test_free_full_leaves_the_mapping_intact(self):
+        indices = _swa_alloc(self.allocator, 4)
+        before = self.allocator.full_to_swa_index_mapping[indices].clone()
+        self.allocator.free_full(indices)
+
+        self.assertTrue(bool((before > 0).all()))
+        self.assertTrue(
+            torch.equal(self.allocator.full_to_swa_index_mapping[indices], before)
+        )
+
+    def test_free_full_is_deferred_inside_a_free_group(self):
+        indices = _swa_alloc(self.allocator, 4)
+
+        self.allocator.free_group_begin()
+        self.allocator.free_full(indices)
+        self.assertEqual(self.allocator.full_available_size(), self.full_baseline - 4)
+        self.allocator.free_group_end()
+
+        self.assertEqual(self.allocator.full_available_size(), self.full_baseline)
+
+
 if __name__ == "__main__":
     unittest.main()

test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py

@@ -6733,6 +6733,7 @@ class TestUnifiedRadixCacheActionRouting(CustomTestCase):
         # translate the source full to SWA and store it on the node (no free)
         alloc.translate_loc_from_full_to_swa.assert_called_once_with(source_value)
         alloc.free.assert_not_called()
+        alloc.free_full.assert_not_called()
         cache.tree_core.set_component_device_value.assert_called_once_with(
             5, ComponentType.SWA, swa_value
         )
@@ -6767,7 +6768,9 @@ class TestUnifiedRadixCacheActionRouting(CustomTestCase):
         alloc.set_full_to_swa_mapping.assert_called_once_with(kept_full, swa_value)
         # the incoming full's stale mapping is cleared, then its slot freed (full-only)
         alloc.clear_full_to_swa_mapping.assert_called_once_with(incoming_full)
-        alloc.full_attn_allocator.free.assert_called_once_with(incoming_full)
+        alloc.free_full.assert_called_once_with(incoming_full)
+        # not the inner allocator (skips the free-group defer) and not both halves
+        alloc.full_attn_allocator.free.assert_not_called()
         alloc.free.assert_not_called()
         cache.tree_core.set_component_device_value.assert_called_once_with(
             5, ComponentType.SWA, swa_value