"""Unit tests for EmbeddingCacheController paged host pool behavior.""" import threading import unittest from queue import Queue from unittest.mock import MagicMock import torch from sglang.srt.managers.schedule_batch import Modality from sglang.srt.mem_cache.storage.mooncake_store.embedding_cache_controller import ( EmbeddingCacheController, EmbeddingCacheEntry, EmbeddingPool, EntryState, EvictableLRU, PageRun, RangePageAllocator, build_transfer_buffers, ) from sglang.test.ci.ci_register import register_cpu_ci register_cpu_ci(est_time=5, suite="base-a-test-cpu") def _make_pool(num_pages=16, dim=4, page_size=2, modality="vision"): total_tokens = num_pages * page_size tensor = torch.empty((total_tokens, dim), dtype=torch.float32) return EmbeddingPool( modality=modality, dim=dim, dtype=torch.float32, page_size=page_size, tensor=tensor, num_pages=num_pages, allocator=RangePageAllocator(num_pages), page_bytes=page_size * dim * torch.float32.itemsize, pool_size_bytes=total_tokens * dim * torch.float32.itemsize, pin_memory=False, ) def _make_controller(num_pages=16, dim=4, page_size=2, enable_eviction=True): ctrl = EmbeddingCacheController.__new__(EmbeddingCacheController) ctrl.tp_world_size = 1 ctrl.tp_group = None ctrl.tp_rank = 0 ctrl.all_rank_get = False ctrl.hidden_dims = { Modality.IMAGE: dim, Modality.VIDEO: dim, Modality.AUDIO: dim, } ctrl.dtype = torch.float32 ctrl.element_size = torch.float32.itemsize ctrl.enable_eviction = enable_eviction ctrl.max_eviction_batch = 10 ctrl.mooncake_store = MagicMock() ctrl.total_pool_size_bytes = num_pages * page_size * dim * torch.float32.itemsize ctrl.vision_pool = _make_pool(num_pages, dim, page_size) ctrl.audio_pool = _make_pool(num_pages, dim, page_size, modality="audio") ctrl.pools = {"vision": ctrl.vision_pool, "audio": ctrl.audio_pool} ctrl.entries = {} ctrl.vision_pool.evictable = EvictableLRU() ctrl.audio_pool.evictable = EvictableLRU() ctrl.stats = { "total_allocated": 0, "total_evicted": 0, "eviction_count": 0, "allocation_failures": 0, } ctrl.ongoing_prefetch = {} ctrl.prefetch_queue = Queue() ctrl.insert_queue = Queue() ctrl.lock = threading.Lock() ctrl.stop_event = threading.Event() ctrl.io_thread = MagicMock() ctrl.prefetch_tp_group = None ctrl._copy_streams = {} return ctrl class TestRangePageAllocator(unittest.TestCase): def test_prefers_single_contiguous_run(self): allocator = RangePageAllocator(num_pages=8) runs = allocator.allocate(num_tokens=6, page_size=2) self.assertEqual(runs, [PageRun(start=0, length=3)]) self.assertEqual(allocator.free_ranges, [(3, 5)]) def test_free_merges_adjacent_ranges(self): allocator = RangePageAllocator(num_pages=8) first = allocator.allocate(num_tokens=4, page_size=2) second = allocator.allocate(num_tokens=4, page_size=2) allocator.free(first) allocator.free(second) self.assertEqual(allocator.free_ranges, [(0, 8)]) def test_scatter_fallback_returns_physical_order(self): allocator = RangePageAllocator(num_pages=4) a = allocator.allocate(num_tokens=2, page_size=2) b = allocator.allocate(num_tokens=2, page_size=2) c = allocator.allocate(num_tokens=2, page_size=2) d = allocator.allocate(num_tokens=2, page_size=2) allocator.free(a) allocator.free(c) runs = allocator.allocate(num_tokens=4, page_size=2) self.assertEqual(runs, [PageRun(start=0, length=1), PageRun(start=2, length=1)]) self.assertEqual([run.start for run in runs], sorted(run.start for run in runs)) self.assertEqual(b, [PageRun(start=1, length=1)]) self.assertEqual(d, [PageRun(start=3, length=1)]) def test_allocate_fails_when_total_free_pages_are_insufficient(self): allocator = RangePageAllocator(num_pages=2) allocator.allocate(num_tokens=4, page_size=2) self.assertIsNone(allocator.allocate(num_tokens=2, page_size=2)) self.assertEqual(allocator.free_pages, 0) class TestEntryStateAndPins(unittest.TestCase): def test_ready_entry_with_no_pins_is_evictable(self): entry = EmbeddingCacheEntry( hash="h", modality=Modality.IMAGE, num_tokens=2, dim=4, page_runs=[PageRun(0, 1)], state=EntryState.READY, ) self.assertTrue(entry.is_evictable()) def test_ready_entry_with_pin_is_not_evictable(self): entry = EmbeddingCacheEntry( hash="h", modality=Modality.IMAGE, num_tokens=2, dim=4, page_runs=[PageRun(0, 1)], state=EntryState.READY, ) entry.pin() self.assertFalse(entry.is_evictable()) entry.unpin() self.assertTrue(entry.is_evictable()) def test_multiple_pins_require_all_unpins(self): entry = EmbeddingCacheEntry( hash="h", modality=Modality.IMAGE, num_tokens=2, dim=4, page_runs=[PageRun(0, 1)], state=EntryState.READY, ) entry.pin() entry.pin() self.assertEqual(entry.ref_count, 2) self.assertFalse(entry.is_evictable()) class TestEvictableLruInvariant(unittest.TestCase): def _insert_entry( self, ctrl, mm_hash, modality=Modality.IMAGE, state=EntryState.READY, ): pool = ctrl._get_pool(modality) page_runs = pool.allocator.allocate(2, pool.page_size) entry = EmbeddingCacheEntry( hash=mm_hash, modality=modality, num_tokens=2, dim=pool.dim, page_runs=page_runs, state=state, ) ctrl.entries[mm_hash] = entry return entry def test_filling_entry_is_not_in_evictable_lru_until_ready(self): ctrl = _make_controller(num_pages=4, dim=4, page_size=2) entry = self._insert_entry( ctrl, "h", state=EntryState.FILLING, ) self.assertNotIn("h", ctrl.vision_pool.evictable) with ctrl.lock: ctrl._mark_ready(entry) self.assertEqual(list(ctrl.vision_pool.evictable.keys()), ["h"]) def test_first_read_pin_removes_candidate_and_last_release_reinserts(self): ctrl = _make_controller(num_pages=4, dim=4, page_size=2) entry = self._insert_entry(ctrl, "h") with ctrl.lock: ctrl._lru_touch("h") with ctrl.lock: ctrl._pin_read(entry) ctrl._pin_read(entry) self.assertNotIn("h", ctrl.vision_pool.evictable) with ctrl.lock: ctrl._unpin_read(entry) self.assertNotIn("h", ctrl.vision_pool.evictable) with ctrl.lock: ctrl._unpin_read(entry) self.assertEqual(list(ctrl.vision_pool.evictable.keys()), ["h"]) def test_evict_for_pool_pops_only_that_pool_candidates(self): ctrl = _make_controller(num_pages=4, dim=4, page_size=2) self._insert_entry(ctrl, "vision_h", modality=Modality.IMAGE) self._insert_entry(ctrl, "audio_h", modality=Modality.AUDIO) with ctrl.lock: ctrl._lru_touch("vision_h") ctrl._lru_touch("audio_h") required_pages = ctrl.vision_pool.allocator.free_pages + 1 with ctrl.lock: ctrl._evict_for_pool(ctrl.vision_pool, required_pages) self.assertNotIn("vision_h", ctrl.entries) self.assertIn("audio_h", ctrl.entries) self.assertEqual(list(ctrl.audio_pool.evictable.keys()), ["audio_h"]) class TestStoreToPool(unittest.TestCase): def test_store_to_pool_async_raises_on_cpu_tensor(self): ctrl = _make_controller(num_pages=8, dim=4, page_size=2) tensor = torch.empty((2, 4), dtype=torch.float32) with self.assertRaises(ValueError): ctrl.store_to_pool_async(["h"], [tensor], Modality.IMAGE) def _insert_ready_entry(ctrl, mm_hash, tensor, modality=Modality.IMAGE): """Manually write tensor into pool pages and create a READY entry.""" pool = ctrl._get_pool(modality) if tensor.ndim != 2: tensor = tensor.reshape(-1, tensor.shape[-1]) num_tokens = int(tensor.shape[0]) page_runs = pool.allocator.allocate(num_tokens, pool.page_size) copied = 0 for run in page_runs: valid = min(pool.page_size * run.length, num_tokens - copied) start = run.start * pool.page_size pool.tensor[start : start + valid].copy_(tensor[copied : copied + valid]) copied += valid entry = EmbeddingCacheEntry( hash=mm_hash, modality=modality, num_tokens=num_tokens, dim=int(tensor.shape[1]), page_runs=page_runs, state=EntryState.READY, ) ctrl.entries[mm_hash] = entry pool.evictable.touch(mm_hash) return entry class TestMooncakeLifecycle(unittest.TestCase): def test_prefetch_creates_filling_entry_and_get_success_marks_ready(self): ctrl = _make_controller(num_pages=4, dim=4, page_size=2) ctrl.prefetch("req", ["h"], [2], Modality.IMAGE) op = ctrl.ongoing_prefetch["req"] ctrl._finish_get(op, [True]) entry = ctrl.entries["h"] self.assertEqual(entry.state, EntryState.READY) def test_prefetch_get_failure_frees_entry(self): ctrl = _make_controller(num_pages=4, dim=4, page_size=2) ctrl.prefetch("req", ["h"], [2], Modality.IMAGE) op = ctrl.ongoing_prefetch["req"] ctrl._finish_get(op, [False]) self.assertNotIn("h", ctrl.entries) self.assertEqual(ctrl.vision_pool.allocator.free_pages, 4) def test_insert_batch_pins_and_releases_on_put(self): ctrl = _make_controller(num_pages=4, dim=4, page_size=2) tensor = torch.arange(8, dtype=torch.float32).reshape(2, 4) _insert_ready_entry(ctrl, "h", tensor) ctrl.insert_batch(["h"], Modality.IMAGE) op = ctrl.insert_queue.get_nowait() entry = ctrl.entries["h"] self.assertEqual(entry.ref_count, 1) self.assertNotIn("h", ctrl.vision_pool.evictable) ctrl._finish_put(op, [True]) self.assertEqual(entry.ref_count, 0) self.assertEqual(entry.state, EntryState.READY) self.assertIn("h", ctrl.vision_pool.evictable) class TestGetPoolViews(unittest.TestCase): def test_get_pool_views_returns_none_for_filling_entry(self): ctrl = _make_controller(num_pages=4, dim=4, page_size=2) ctrl.entries["h"] = EmbeddingCacheEntry( hash="h", modality=Modality.IMAGE, num_tokens=2, dim=4, page_runs=[PageRun(0, 1)], state=EntryState.FILLING, ) views = ctrl.get_pool_views(["h"]) self.assertIsNone(views[0]) def test_get_pool_views_returns_slices_and_release_unpins(self): ctrl = _make_controller(num_pages=4, dim=4, page_size=2) tensor = torch.arange(8, dtype=torch.float32).reshape(2, 4) _insert_ready_entry(ctrl, "h", tensor) views = ctrl.get_pool_views(["h"]) self.assertIsNotNone(views[0]) entry = ctrl.entries["h"] self.assertEqual(entry.ref_count, 1) flat = torch.cat(views[0], dim=0) self.assertTrue(torch.equal(flat, tensor)) ctrl.release_pool_views(["h"]) self.assertEqual(entry.ref_count, 0) self.assertIn("h", ctrl.vision_pool.evictable) class TestTransferBuffers(unittest.TestCase): def test_build_transfer_buffers_for_single_run(self): pool = _make_pool(num_pages=8, dim=4, page_size=2) entry = EmbeddingCacheEntry( hash="h", modality=Modality.IMAGE, num_tokens=5, dim=4, page_runs=[PageRun(2, 3)], state=EntryState.READY, ) ptrs, sizes = build_transfer_buffers(entry, pool) self.assertEqual(ptrs, [pool.tensor[4].data_ptr()]) self.assertEqual(sizes, [5 * 4 * torch.float32.itemsize]) def test_build_transfer_buffers_for_multiple_runs(self): pool = _make_pool(num_pages=8, dim=4, page_size=2) entry = EmbeddingCacheEntry( hash="h", modality=Modality.IMAGE, num_tokens=5, dim=4, page_runs=[PageRun(0, 1), PageRun(3, 2)], state=EntryState.READY, ) ptrs, sizes = build_transfer_buffers(entry, pool) self.assertEqual(ptrs, [pool.tensor[0].data_ptr(), pool.tensor[6].data_ptr()]) self.assertEqual( sizes, [ 2 * 4 * torch.float32.itemsize, 3 * 4 * torch.float32.itemsize, ], ) class TestMooncakeEmbeddingStoreWrappers(unittest.TestCase): def test_batch_put_multi_buffers_deduplicates_existing_keys(self): from sglang.srt.mem_cache.storage.mooncake_store.mooncake_embedding_store import ( MooncakeEmbeddingStore, ) store = MooncakeEmbeddingStore.__new__(MooncakeEmbeddingStore) store.store = MagicMock() store.store.batch_is_exist.return_value = [1, 0] store.store.batch_put_from_multi_buffers.return_value = [0] results = store.batch_put_from_multi_buffers( ["a", "b"], [[11], [22]], [[4], [4]], ) self.assertEqual(results, [True, True]) store.store.batch_put_from_multi_buffers.assert_called_once_with( ["emb_b"], [[22]], [[4]] ) def test_batch_get_multi_buffers_maps_positive_result_to_true(self): from sglang.srt.mem_cache.storage.mooncake_store.mooncake_embedding_store import ( MooncakeEmbeddingStore, ) store = MooncakeEmbeddingStore.__new__(MooncakeEmbeddingStore) store.store = MagicMock() store.store.batch_get_into_multi_buffers.return_value = [8, -1] results = store.batch_get_into_multi_buffers( ["a", "b"], [[11], [22]], [[4], [4]], ) self.assertEqual(results, [True, False]) if __name__ == "__main__": unittest.main()