feat(hicache): Add shared memory allocator for host KV cache (#29326)

Co-authored-by: Jimmy Shong <69131491+Jiminator@users.noreply.github.com>
Co-authored-by: Zhangheng <hzh0425@apache.org>

test/registered/unit/mem_cache/test_mem_pool_host.py

@@ -83,6 +83,24 @@ class TestHostKVCache(CustomTestCase):
         self.assertIn("Double-free", msg)
         self.assertIn(str(indices.tolist()), msg)
 
+    def test_shm_allocator(self):
+        shm_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="shm",
+        )
+        self.assertIsNotNone(shm_host_pool.fd)
+        self.assertGreaterEqual(shm_host_pool.fd, 0)
+
+        indices = shm_host_pool.alloc(4)
+        self.assertEqual(len(indices), 4)
+        shm_host_pool.free(indices)
+
     def test_empty_free_keeps_release_list_empty(self):
         self.assertEqual(self.host_pool.free(torch.empty(0, dtype=torch.int64)), 0)
         self.assertEqual(self.host_pool.num_release_slots, 0)

test/registered/unit/mem_cache/test_mmap_allocator.py

@@ -0,0 +1,119 @@
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=10, suite="base-a-test-cpu")
+
+import sys
+
+sys.modules["libtpu"] = None
+import mmap
+import os
+import unittest
+
+import torch
+
+from sglang.srt.mem_cache.pool_host.common import ShmHostTensorAllocator
+from sglang.srt.mem_cache.storage.mmap import alloc_mmap, alloc_shm
+
+
+class TestMmapAllocator(unittest.TestCase):
+    def test_alloc_mmap(self):
+        dims = (10, 1024)
+        dtype = torch.float32
+        tensor = alloc_mmap(dims, dtype)
+        self.assertEqual(tensor.shape, dims)
+        self.assertEqual(tensor.dtype, dtype)
+        # Verify it has mapped memory address
+        self.assertGreater(tensor.data_ptr(), 0)
+
+    def test_alloc_shm(self):
+        dims = (10, 1024)
+        dtype = torch.float32
+        tensor, fd, mm = alloc_shm(dims, dtype)
+
+        self.assertEqual(tensor.shape, dims)
+        self.assertEqual(tensor.dtype, dtype)
+        self.assertGreater(tensor.data_ptr(), 0)
+        self.assertGreaterEqual(fd, 0)
+        self.assertIsInstance(mm, mmap.mmap)
+
+        # Check that we can write to the tensor
+        tensor[0, 0] = 42.0
+        self.assertEqual(tensor[0, 0].item(), 42.0)
+
+        # Check that the FD is open and valid
+        try:
+            os.lseek(fd, 0, os.SEEK_SET)
+        except OSError:
+            self.fail("FD is not valid or closed")
+
+        # Cleanup
+        mm.close()
+        os.close(fd)
+
+    def test_shm_host_tensor_allocator(self):
+        allocator = ShmHostTensorAllocator()
+        dims = (2, 512)
+        dtype = torch.int32
+
+        tensor = allocator.allocate(dims, dtype, "cpu")
+        self.assertEqual(tensor.shape, dims)
+        self.assertEqual(tensor.dtype, dtype)
+        self.assertIsNotNone(allocator.fd)
+        self.assertGreaterEqual(allocator.fd, 0)
+
+        # Write data and check
+        tensor[1, 1] = 99
+        self.assertEqual(tensor[1, 1].item(), 99)
+
+        # Test destructor cleans up fd
+        fd = allocator.fd
+        # Trigger GC / deletion
+        del allocator
+
+        # Verify fd is closed
+        with self.assertRaises(OSError):
+            os.fstat(fd)
+
+    def test_alloc_shm_unlinked(self):
+        dims = (4, 256)
+        dtype = torch.float32
+        tensor, fd, mm = alloc_shm(dims, dtype)
+
+        # On Linux, the path of an unlinked fd shows up in /proc/self/fd/
+        # with a ' (deleted)' suffix.
+        fd_path = f"/proc/self/fd/{fd}"
+        try:
+            resolved_path = os.readlink(fd_path)
+            self.assertIn("sglang_host_pool_", resolved_path)
+            self.assertTrue(resolved_path.endswith(" (deleted)"))
+        except OSError:
+            # If procfs is not available or readlink fails, fallback to direct path existence check
+            self.assertFalse(os.path.exists(f"/dev/shm/sglang_host_pool_"))
+
+        # Cleanup
+        mm.close()
+        os.close(fd)
+
+    def test_alloc_shm_hugepage_warning(self):
+        from sglang.srt.environ import envs
+
+        envs.SGLANG_HUGEPAGE_SIZE.override("2MB")
+        try:
+            # Should succeed by falling back to plain page size mapping
+            dims = (2, 2)
+            tensor, fd, mm = alloc_shm(dims, torch.float32)
+            self.assertEqual(tensor.shape, dims)
+            mm.close()
+            os.close(fd)
+        finally:
+            envs.SGLANG_HUGEPAGE_SIZE.override(None)
+
+    def test_shm_host_tensor_allocator_invalid_device(self):
+        allocator = ShmHostTensorAllocator()
+        with self.assertRaises(AssertionError) as ctx:
+            allocator.allocate((2, 2), torch.float32, device="cuda")
+        self.assertIn("only supports CPU allocations", str(ctx.exception))
+
+
+if __name__ == "__main__":
+    unittest.main()