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

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)