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

test/registered/unit/mem_cache/test_swa_locked_full_recover_unified.py

@@ -141,7 +141,7 @@ class _StaticAllocRecorder:
         self.clear_calls.append(full)
         self.full_to_swa_index_mapping[full.to(torch.int64)] = 0
 
-    def free_full(self, indices):
+    def free_full_segment(self, indices, *, start_pos):
         self.freed_full.append(indices)
 
     def free(self, indices):
@@ -307,8 +307,8 @@ class TestStaticPoolPathUnchanged(unittest.TestCase):
             ),
             "incoming ids' mapping entries must be zeroed (static recipe)",
         )
-        # Through free_full, not the inner allocator: the latter skips the
-        # free-group defer.
+        # Through free_full_segment, not the inner allocator: the latter skips
+        # the free-group defer.
         self.assertEqual(len(static.freed_full), 1)
         self.assertEqual(static.freed_via_inner, [])
         ((node_id, ct, _),) = probe.tree_core.set_calls

test/registered/unit/mem_cache/test_swa_unittest.py

@@ -2,6 +2,7 @@ import unittest
 from array import array
 from types import SimpleNamespace
 from unittest import mock
+from unittest.mock import patch
 
 import torch
 
@@ -830,13 +831,13 @@ class TestSWA(unittest.TestCase):
         req2.prefix_indices = torch.tensor([21, 22, 23, 24, 25], device=tree.device)
 
         freed_lens = []
-        original_free = allocator.free
+        original_free_segment = allocator.free_segment
 
-        def wrapped_free(indices):
+        def wrapped_free_segment(indices, *, start_pos):
             freed_lens.append(int(indices.numel()))
-            return original_free(indices)
+            return original_free_segment(indices, start_pos=start_pos)
 
-        allocator.free = wrapped_free
+        allocator.free_segment = wrapped_free_segment
         tree.cache_finished_req(
             req2, is_insert=False, kv_len_to_handle=req2._kv_committed_len
         )
@@ -1070,21 +1071,45 @@ class TestFreeKvRow(CustomTestCase):
                 )
                 self.assertEqual(self._sizes(), (self.full_baseline, self.swa_baseline))
 
-    def test_adjacent_below_floor_pieces_release_their_shared_page_once(self):
+    def test_below_floor_pieces_go_back_through_the_full_side(self):
         _, allocator, _ = _build_swa_tree(is_eagle=False, page_size=4)
         indices = _swa_alloc(allocator, 8)
         allocator.free_swa(indices)
         after_alloc = allocator.full_available_size()
 
-        # Rows [0, 6) and [6, 8) both sit below the floor and share page 1.
-        free_kv_row_segments(
-            allocator,
-            [(indices[:6], 0), (indices[6:], 6)],
-            swa_evicted_seqlen=8,
-        )
+        # Both rows [0, 4) and [4, 8) sit below the floor: full side only.
+        with patch.object(
+            allocator.full_attn_allocator,
+            "free",
+            side_effect=AssertionError("full side took the unique path"),
+        ):
+            free_kv_row_segments(
+                allocator,
+                [(indices[:4], 0), (indices[4:], 4)],
+                swa_evicted_seqlen=8,
+            )
 
         self.assertEqual(allocator.full_available_size(), after_alloc + 8)
 
+    def test_grouped_full_side_frees_defer_and_skip_the_unique_path(self):
+        _, allocator, _ = _build_swa_tree(is_eagle=False, page_size=4)
+        indices = _swa_alloc(allocator, 12)
+        allocator.free_swa(indices[:8])
+        after_alloc = allocator.full_available_size()
+
+        with patch.object(
+            allocator.full_attn_allocator,
+            "free",
+            side_effect=AssertionError("full side took the unique path"),
+        ):
+            allocator.free_group_begin()
+            # dead rows [0, 8) and the alive row [8, 12) from one request
+            free_kv_row_segments(allocator, [(indices, 0)], swa_evicted_seqlen=8)
+            self.assertEqual(allocator.full_available_size(), after_alloc)
+            allocator.free_group_end()
+
+        self.assertEqual(allocator.full_available_size(), after_alloc + 12)
+
     def test_free_kv_row_reads_the_record_row_and_its_floor(self):
         indices = _swa_alloc(self.allocator, 8)
         cache = _RowCache(self.allocator, indices)

test/registered/unit/mem_cache/test_unified_free_no_host_sync.py

@@ -406,6 +406,26 @@ class TestEveryUnifiedAllocatorOverridesFreeSegment(unittest.TestCase):
                 self.assertIn("free_page_reps_group", inspect.getsource(cls))
 
 
+class TestUnifiedSwaFullSideGroup(unittest.TestCase):
+    """`free_full` inside a free group must defer and land at `free_group_end`;
+    the composite carries that pile itself, not through the SWA parent's hooks."""
+
+    def test_full_only_frees_defer_until_group_end(self):
+        from test_unified_swa_shared_virtual_ids import _build
+
+        alloc = _build(64, 32, 2, 4)
+        idx = alloc.alloc(8)
+        alloc.free_swa(idx)
+        before = alloc.full_available_size()
+
+        alloc.free_group_begin()
+        alloc.free_full(idx[:4])
+        alloc.free_full_segment(idx[4:], start_pos=4)
+        self.assertEqual(alloc.full_available_size(), before)
+        alloc.free_group_end()
+        self.assertEqual(alloc.full_available_size(), before + 8)
+
+
 class TestFreeSwaWindowRatchetNoHostSync(unittest.TestCase):
     """The per-decode-step SWA window ratchet frees a CONTIGUOUS row slice
     with host-int, page-aligned bounds — the same shape `free_segment` was
@@ -471,6 +491,25 @@ class TestFreeSwaWindowRatchetNoHostSync(unittest.TestCase):
             alloc.free_swa(v[: 4 * self.PS], start_pos=0)
             alloc.free_swa(v[4 * self.PS :], start_pos=4 * self.PS)
 
+    def test_full_only_segment_free_never_syncs(self):
+        """The request-finish dead half (swa already tombstoned) frees the full
+        side by page reps: no unique from free_full's token dedup."""
+        alloc = self._swa_composite(lazy=True)
+        v = alloc.alloc(8 * self.PS)
+        alloc.free_swa(v, start_pos=0)
+        before = alloc.full_available_size()
+        with (
+            mock.patch.object(
+                torch, "unique", side_effect=AssertionError("unique = host sync")
+            ),
+            mock.patch.object(
+                torch.Tensor, "item", side_effect=AssertionError("item = host sync")
+            ),
+        ):
+            alloc.free_full_segment(v[: 4 * self.PS], start_pos=0)
+            alloc.free_full_segment(v[4 * self.PS :], start_pos=4 * self.PS)
+        self.assertEqual(alloc.full_available_size(), before + 8 * self.PS)
+
     def test_unaligned_start_pos_is_rejected(self):
         """A mid-page start must fail loudly, not release the head page whole."""
         alloc = self._swa_composite(lazy=True)

test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py

@@ -516,15 +516,18 @@ def build_fixture(
             full_attention_layer_ids=cfg.full_attention_layer_ids,
             device=device,
         )
-        allocator = SWATokenToKVPoolAllocator(
-            size=cfg.kv_size,
-            size_swa=cfg.kv_size,
-            page_size=cfg.page_size,
-            dtype=cfg.dtype,
-            device=device,
-            kvcache=kv_pool,
-            need_sort=False,
-        )
+        # Tree values reach the allocator as start_pos=0 segments, so only the
+        # debug cross-check against torch.unique catches a mis-aligned value.
+        with envs.SGLANG_DEBUG_MEMORY_POOL.override(True):
+            allocator = SWATokenToKVPoolAllocator(
+                size=cfg.kv_size,
+                size_swa=cfg.kv_size,
+                page_size=cfg.page_size,
+                dtype=cfg.dtype,
+                device=device,
+                kvcache=kv_pool,
+                need_sort=False,
+            )
     elif cfg.has_mamba:
         kv_pool = HybridLinearKVPool(
             size=cfg.kv_size,
@@ -4070,11 +4073,8 @@ class UnifiedRadixCacheSuite:
         # for an un-charged gate to accept.
         extend_need = 2 * ps + 1
         max_new = 8
-        self.assertIsNotNone(
-            cons_alloc.swa_attn_allocator.alloc(
-                cons_alloc.swa_available_size() - (window + extend_need - 1)
-            )
-        )
+        hold = cons_alloc.swa_available_size() - (window + extend_need - 1)
+        self.assertIsNotNone(cons_alloc.swa_attn_allocator.alloc(hold // ps * ps))
 
         # The adder's SWA gate for this request (_swa_budget_for_req).
         surfaced_swa_hit = cons.staged_prefetch_swa_tokens(req_id)
@@ -7559,8 +7559,8 @@ class TestUnifiedRadixCacheActionRouting(CustomTestCase):
         _component_with_cache(ComponentType.FULL, cache).apply_component_action(
             FreeComponentDeviceSlot([indices], component_type=ComponentType.FULL)
         )
-        cache.token_to_kv_pool_allocator.full_attn_allocator.free.assert_called_once_with(
-            indices
+        cache.token_to_kv_pool_allocator.full_attn_allocator.free_segment.assert_called_once_with(
+            indices, start_pos=0
         )
 
     def test_apply_component_action_device_kv_swa_uses_free_swa(self):
@@ -7681,7 +7681,7 @@ 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.free_full.assert_called_once_with(incoming_full)
+        alloc.free_full_segment.assert_called_once_with(incoming_full, start_pos=0)
         # Never by indexing the tensor: the unified composite has no
         # `full_to_swa_index_mapping` to index into.
         alloc.full_to_swa_index_mapping.__setitem__.assert_not_called()