[HiCache] Optimize LogicalHostPool free-list release (#33998)

python/sglang/srt/mem_cache/memory_pool_host.py

@@ -670,13 +670,29 @@ class LogicalHostPool:
     @synchronized
     def clear(self):
         self.free_slots = torch.arange(self.size, dtype=torch.int64)
+        # Match HostKVCache's lazy release path: defer large free-list merges
+        # until an allocation needs the released slots.
+        self.release_slots = []
+        self.num_release_slots = 0
 
     def destroy(self) -> None:
         """Logical anchors own no backing buffers or registrations to release."""
         return None
 
     def available_size(self):
-        return len(self.free_slots)
+        return len(self.free_slots) + self.num_release_slots
+
+    def _merge_release_slots(self):
+        if self.num_release_slots == 0:
+            return
+
+        if len(self.free_slots) == 0 and len(self.release_slots) == 1:
+            self.free_slots = self.release_slots[0]
+        else:
+            self.free_slots = torch.cat([self.free_slots, *self.release_slots])
+
+        self.release_slots = []
+        self.num_release_slots = 0
 
     @synchronized
     def alloc(self, need_size: int) -> Optional[torch.Tensor]:
@@ -687,6 +703,10 @@ class LogicalHostPool:
             )
         if need_size > self.available_size():
             return None
+
+        if need_size > len(self.free_slots):
+            self._merge_release_slots()
+
         select_index = self.free_slots[:need_size]
         self.free_slots = self.free_slots[need_size:]
         return select_index
@@ -698,9 +718,12 @@ class LogicalHostPool:
                 "LogicalHostPool free must be page-aligned, "
                 f"got len(indices)={len(indices)}, page_size={self.page_size}"
             )
-        self.free_slots = torch.cat(
-            [self.free_slots, indices.to(dtype=torch.int64, device="cpu").flatten()]
-        )
+        indices_cpu = indices.to(dtype=torch.int64, device="cpu").flatten()
+        if indices_cpu.numel() == 0:
+            return 0
+
+        self.release_slots.append(indices_cpu)
+        self.num_release_slots += len(indices_cpu)
         return len(indices)
 
     def backup_from_device_all_layer(

test/registered/unit/mem_cache/test_mem_pool_host.py

@@ -8,6 +8,7 @@ import torch
 from sglang.srt.mem_cache.memory_pool import MHATokenToKVPool
 from sglang.srt.mem_cache.memory_pool_host import (
     DeepSeekV4PagedHostPool,
+    LogicalHostPool,
     MambaPoolHost,
 )
 from sglang.srt.mem_cache.pool_host.mha import MHATokenToKVPoolHost
@@ -127,6 +128,10 @@ class TestLazyHostPoolRelease(CustomTestCase):
         pool.clear()
         return pool
 
+    @staticmethod
+    def _make_logical_pool():
+        return LogicalHostPool(size=8, page_size=2)
+
     def _assert_lazy_release(self, pool):
         self.assertEqual(pool.free(torch.empty(0, dtype=torch.int64)), 0)
         self.assertEqual(pool.num_release_slots, 0)
@@ -181,6 +186,17 @@ class TestLazyHostPoolRelease(CustomTestCase):
         pool.clear()
         self.assertEqual(len(pool.alloc(1)), 2)
 
+    def test_logical_pool_lazy_release(self):
+        pool = self._make_logical_pool()
+        self._assert_lazy_release(pool)
+
+        # Preserve the logical pool's strict page-alignment checks.
+        pool.clear()
+        with self.assertRaises(ValueError):
+            pool.alloc(1)
+        with self.assertRaises(ValueError):
+            pool.free(torch.tensor([0]))
+
 
 if __name__ == "__main__":
     unittest.main()