[HiCache] Detect for double-free in HostKVCache (#29310)

python/sglang/srt/mem_cache/pool_host/base.py

@@ -230,6 +230,9 @@ class HostKVCache(abc.ABC):
             (self.size,), dtype=torch.uint8, device=self.device
         )
         self.free_slots = torch.arange(self.size, dtype=torch.int64)
+        # Per-slot flag used to detect double-free.
+        # slot_used[k] is true if slot k is allocated.
+        self.slot_used = torch.zeros(self.size, dtype=torch.bool)
 
     def available_size(self):
         return len(self.free_slots)
@@ -245,9 +248,21 @@ class HostKVCache(abc.ABC):
         select_index = self.free_slots[:need_size]
         self.free_slots = self.free_slots[need_size:]
 
+        assert not self.slot_used[select_index].any(), (
+            f"Double-alloc detected: slots already allocated: "
+            f"{select_index[self.slot_used[select_index]].tolist()}."
+        )
+        self.slot_used[select_index] = True
+
         return select_index
 
     @synchronized
     def free(self, indices: torch.Tensor) -> int:
-        self.free_slots = torch.cat([self.free_slots, indices.cpu()])
+        indices_cpu = indices.cpu()
+        assert self.slot_used[indices_cpu].all(), (
+            f"Double-free detected: slots not currently allocated: "
+            f"{indices_cpu[~self.slot_used[indices_cpu]].tolist()}."
+        )
+        self.slot_used[indices_cpu] = False
+        self.free_slots = torch.cat([self.free_slots, indices_cpu])
         return len(indices)

test/registered/unit/mem_cache/test_mem_pool_host.py

@@ -0,0 +1,83 @@
+"""Unit tests for HostKVCache alloc/free bookkeeping (double-alloc / double-free detection)."""
+
+import unittest
+
+import torch
+
+from sglang.srt.mem_cache.memory_pool import MHATokenToKVPool
+from sglang.srt.mem_cache.memory_pool_host import MHATokenToKVPoolHost
+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")
+
+
+class TestHostKVCache(CustomTestCase):
+    def setUp(self):
+        self.page_size = 2
+        # Small device pool is enough to construct the host pool.
+        self.device_pool = MHATokenToKVPool(
+            size=self.page_size * 2,
+            page_size=self.page_size,
+            dtype=torch.float16,
+            head_num=2,
+            head_dim=4,
+            layer_num=2,
+            device="cpu",
+            enable_memory_saver=False,
+        )
+        self.host_pool = MHATokenToKVPoolHost(
+            device_pool=self.device_pool,
+            host_to_device_ratio=2.0,
+            host_size=0,
+            page_size=self.page_size,
+            layout="layer_first",
+            pin_memory=False,
+            device="cpu",
+            allocator_type="default",
+        )
+
+    def test_double_alloc(self):
+        indices = self.host_pool.alloc(4)
+        self.assertEqual(len(indices), 4)
+        # Mimic bookkeeping corruption: push an already-used slot back to the
+        # head of free_slots so the next alloc would hand out an in-use slot.
+        leak = torch.tensor([int(indices[0])])
+        self.host_pool.free_slots = torch.cat([leak, self.host_pool.free_slots])
+        with self.assertRaises(AssertionError) as ctx:
+            self.host_pool.alloc(4)
+        msg = str(ctx.exception)
+        self.assertIn("Double-alloc", msg)
+        self.assertIn(f"[{int(leak[0])}]", msg)
+
+    def test_double_free(self):
+        indices = self.host_pool.alloc(4)
+        self.assertEqual(len(indices), 4)
+        self.host_pool.free(indices[:2])
+        # indices[1] is double freed.
+        with self.assertRaises(AssertionError) as ctx:
+            self.host_pool.free(indices[1:])
+        msg = str(ctx.exception)
+        self.assertIn("Double-free", msg)
+        self.assertIn(f"[{int(indices[1])}]", msg)
+
+    def test_free_unallocated(self):
+        indices = torch.tensor([1])
+        with self.assertRaises(AssertionError) as ctx:
+            self.host_pool.free(indices)
+        msg = str(ctx.exception)
+        self.assertIn("Double-free", msg)
+        self.assertIn(f"[{int(indices[0])}]", msg)
+
+    def test_free_after_clear(self):
+        indices = self.host_pool.alloc(4)
+        self.host_pool.clear()
+        with self.assertRaises(AssertionError) as ctx:
+            self.host_pool.free(indices)
+        msg = str(ctx.exception)
+        self.assertIn("Double-free", msg)
+        self.assertIn(str(indices.tolist()), msg)
+
+
+if __name__ == "__main__":
+    unittest.main()