[PD] Optimize paged allocator free-list release (#37146)

Co-authored-by: wangwenming.41 <wangwenming.41@jd.com>
Co-authored-by: hnyls2002 <lsyincs@gmail.com>
Co-authored-by: Liangsheng Yin <hnyls2002@gmail.com>

test/registered/unit/mem_cache/test_paged_allocator_lazy_release.py

@@ -0,0 +1,70 @@
+"""need_sort allocators stage released pages and merge only under allocation pressure."""
+
+import unittest
+from unittest.mock import patch
+
+import torch
+
+from sglang.srt.mem_cache.allocator.paged import PagedTokenToKVPoolAllocator
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=2, suite="base-a-test-cpu")
+
+PAGE_SIZE = 2
+NUM_PAGES = 16
+
+
+def _make_allocator(*, need_sort: bool) -> PagedTokenToKVPoolAllocator:
+    return PagedTokenToKVPoolAllocator(
+        size=NUM_PAGES * PAGE_SIZE,
+        page_size=PAGE_SIZE,
+        dtype=torch.float16,
+        device="cpu",
+        kvcache=None,
+        need_sort=need_sort,
+    )
+
+
+def _page_ids(indices: torch.Tensor) -> torch.Tensor:
+    return indices[::PAGE_SIZE] // PAGE_SIZE
+
+
+class TestPagedAllocatorLazyRelease(CustomTestCase):
+    def test_release_stages_until_free_pages_run_dry(self):
+        allocator = _make_allocator(need_sort=True)
+        allocated = allocator.alloc(24)
+        free_pages_before = allocator.free_pages
+
+        with patch.object(torch, "cat", wraps=torch.cat) as cat_mock:
+            allocator.free(allocated[:2])
+            allocator.free(allocated[4:6])
+            self.assertIs(allocator.free_pages, free_pages_before)
+            self.assertEqual(allocator.num_staged_pages, 2)
+            self.assertEqual(allocator.available_size(), 12)
+            cat_mock.assert_not_called()
+
+            # free_pages still covers this: no merge.
+            primary = allocator.alloc(8)
+            self.assertTrue(torch.equal(_page_ids(primary), torch.arange(13, 17)))
+            cat_mock.assert_not_called()
+
+            # Pressure: one merge, staged pages come back sorted.
+            reused = allocator.alloc(4)
+            self.assertTrue(torch.equal(_page_ids(reused), torch.tensor([1, 3])))
+            self.assertEqual(cat_mock.call_count, 1)
+
+        self.assertEqual(allocator.num_staged_pages, 0)
+        self.assertEqual(allocator.staged_pages, [])
+
+    def test_clear_drops_staged_pages(self):
+        allocator = _make_allocator(need_sort=True)
+        allocator.free(allocator.alloc(4))
+
+        allocator.clear()
+        self.assertEqual(allocator.available_size(), allocator.size)
+        self.assertEqual(allocator.staged_pages, [])
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/mem_cache/test_paged_free_segment.py

@@ -67,11 +67,12 @@ class TestFreeSegment(unittest.TestCase):
         alloc.free_segment(row[:0], start_pos=0)
         self.assertEqual(len(alloc.free_pages), before)
 
-    def test_need_sort_routes_to_release_pages(self):
+    def test_need_sort_defers_released_pages(self):
         alloc = _make_allocator(need_sort=True)
         row = _make_kv_row(alloc, 2 * PAGE_SIZE)
         alloc.free_segment(row, start_pos=0)
-        self.assertEqual(len(alloc.release_pages), 2)
+        self.assertEqual(len(alloc.staged_pages), 1)
+        self.assertEqual(alloc.num_staged_pages, 2)
 
     def test_group_defers_until_group_end(self):
         alloc = _make_allocator()
@@ -108,8 +109,8 @@ class TestFreeSegment(unittest.TestCase):
         with self.assertRaises(AssertionError):
             alloc.free_group_end()
 
-    def test_group_end_debug_assert_covers_release_pages(self):
-        # need_sort routes frees into release_pages; the duplicate check must
+    def test_group_end_debug_assert_covers_staged_releases(self):
+        # need_sort stages frees in chunks; the duplicate check must
         # not go vacuous there (PD disaggregation runs with need_sort=True).
         alloc = _make_allocator(need_sort=True)
         alloc.debug_mode = True