diff --git a/python/sglang/srt/mem_cache/pool_host/base.py b/python/sglang/srt/mem_cache/pool_host/base.py index 5ef35c25d..39b809af2 100644 --- a/python/sglang/srt/mem_cache/pool_host/base.py +++ b/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) diff --git a/test/registered/unit/mem_cache/test_mem_pool_host.py b/test/registered/unit/mem_cache/test_mem_pool_host.py new file mode 100644 index 000000000..2faea127f --- /dev/null +++ b/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()