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

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