"""Unit tests for mem_cache/utils.py — no server, no model loading.""" import hashlib import sys import types import unittest from array import array from types import SimpleNamespace from unittest.mock import MagicMock, patch from sglang.srt.mem_cache.evict_policy import ( FIFOStrategy, FILOStrategy, LFUStrategy, LRUStrategy, MRUStrategy, PriorityStrategy, SLRUStrategy, ) from sglang.srt.mem_cache.utils import ( compute_node_event_hash_values, compute_node_hash_values, get_eviction_strategy, get_hash_str, hash_str_to_int64, maybe_init_custom_mem_pool, split_node_hash_value, ) from sglang.test.ci.ci_register import register_cpu_ci register_cpu_ci(est_time=15, suite="base-a-test-cpu") def _legacy_get_hash_str(token_ids, prior_hash=None): hasher = hashlib.sha256() if prior_hash: hasher.update(bytes.fromhex(prior_hash)) for t in token_ids: if isinstance(t, tuple): for elem in t: hasher.update(elem.to_bytes(4, byteorder="little", signed=False)) else: hasher.update(t.to_bytes(4, byteorder="little", signed=False)) return hasher.hexdigest() def _legacy_page_hashes(key, page_size, prior_hash=None): hashes = [] running_hash = prior_hash for start in range(0, len(key), page_size): running_hash = _legacy_get_hash_str( key[start : start + page_size], running_hash ) hashes.append(running_hash) return hashes class _HashKey: def __init__(self, token_ids, is_bigram=False, cache_salt=None): self.token_ids = token_ids self.is_bigram = is_bigram self.cache_salt = cache_salt def __len__(self): if self.is_bigram: return max(0, len(self.token_ids) - 1) return len(self.token_ids) def __getitem__(self, index): if isinstance(index, slice): start = index.start or 0 stop = index.stop if index.stop is not None else len(self) if self.is_bigram: return _HashKey( self.token_ids[start : stop + 1], is_bigram=True, cache_salt=self.cache_salt, ) return _HashKey(self.token_ids[start:stop], cache_salt=self.cache_salt) if self.is_bigram: return (self.token_ids[index], self.token_ids[index + 1]) return self.token_ids[index] def raw_token_ids(self): return self.token_ids def hash_page(self, start, end, prior_hash=None): return _legacy_get_hash_str(self[start:end], prior_hash) def _single_hash_compatibility_cases(): prior_hash = _legacy_get_hash_str([7, 8, 9]) return [ ("empty_list", [], None), ("plain_list", [1, 2, 3, 4, 5], None), ("array_q", _HashKey(array("q", range(1, 258))), None), ("array_i_with_prior", _HashKey(array("I", range(1, 258))), prior_hash), ( "tuple_bigram_with_prior", [(10, 20), (20, 30), (30, 40), (40, 50)], prior_hash, ), ( "eagle_bigram_with_prior", _HashKey( array("q", ((i * 2654435761) & 0x00FFFFFF for i in range(258))), is_bigram=True, ), prior_hash, ), ("empty_bigram", _HashKey(array("q"), is_bigram=True), None), ] def _page_hash_compatibility_cases(): prior_hash = _legacy_get_hash_str([7, 8, 9]) return [ ("empty_list", [], 8, None), ("empty_bigram", _HashKey(array("q"), is_bigram=True), 8, None), ("array_q_page_64", _HashKey(array("q", range(1, 258))), 64, None), ( "array_i_page_64_with_prior", _HashKey(array("I", range(1, 258))), 64, prior_hash, ), ( "eagle_bigram_page_64_with_prior", _HashKey( array("q", ((i * 2654435761) & 0x00FFFFFF for i in range(258))), is_bigram=True, ), 64, prior_hash, ), ( "eagle_bigram_page_1_with_prior", _HashKey(array("q", [11, 22, 33, 44, 55]), is_bigram=True), 1, prior_hash, ), ] class TestGetEvictionStrategy(unittest.TestCase): def test_lru(self): self.assertIsInstance(get_eviction_strategy("lru"), LRUStrategy) def test_lfu(self): self.assertIsInstance(get_eviction_strategy("lfu"), LFUStrategy) def test_fifo(self): self.assertIsInstance(get_eviction_strategy("fifo"), FIFOStrategy) def test_mru(self): self.assertIsInstance(get_eviction_strategy("mru"), MRUStrategy) def test_filo(self): self.assertIsInstance(get_eviction_strategy("filo"), FILOStrategy) def test_priority(self): self.assertIsInstance(get_eviction_strategy("priority"), PriorityStrategy) def test_slru(self): self.assertIsInstance(get_eviction_strategy("slru"), SLRUStrategy) def test_case_insensitive(self): self.assertIsInstance(get_eviction_strategy("LRU"), LRUStrategy) self.assertIsInstance(get_eviction_strategy("Lru"), LRUStrategy) self.assertIsInstance(get_eviction_strategy("FIFO"), FIFOStrategy) def test_unknown_policy_raises_valueerror(self): with self.assertRaises(ValueError) as ctx: get_eviction_strategy("nonexistent") msg = str(ctx.exception) self.assertIn("Unknown eviction policy", msg) self.assertIn("lru", msg) self.assertIn("lfu", msg) self.assertIn("fifo", msg) self.assertIn("mru", msg) self.assertIn("filo", msg) self.assertIn("priority", msg) self.assertIn("slru", msg) def test_each_call_creates_new_instance(self): s1 = get_eviction_strategy("lru") s2 = get_eviction_strategy("lru") self.assertIsNot(s1, s2) class TestMaybeInitCustomMemPool(unittest.TestCase): @patch("sglang.srt.mem_cache.utils.envs.SGLANG_MOONCAKE_CUSTOM_MEM_POOL.get") def test_disabled_by_default(self, mock_env_get): mock_env_get.return_value = None enabled, pool, pool_type = maybe_init_custom_mem_pool("cuda:0") self.assertFalse(enabled) self.assertIsNone(pool) self.assertIsNone(pool_type) @patch("sglang.srt.mem_cache.utils.envs.SGLANG_MOONCAKE_CUSTOM_MEM_POOL.get") def test_enabled_via_env(self, mock_env_get): mock_env_get.return_value = "enabled" mock_init = MagicMock() mock_init.return_value = (True, "mock_pool_instance", "mooncake") mooncake_pkg = types.ModuleType("sglang.srt.disaggregation.mooncake") mooncake_utils = types.ModuleType("sglang.srt.disaggregation.mooncake.utils") mooncake_utils.init_mooncake_custom_mem_pool = mock_init with patch.dict( sys.modules, { "sglang.srt.disaggregation.mooncake": mooncake_pkg, "sglang.srt.disaggregation.mooncake.utils": mooncake_utils, }, ): enabled, pool, pool_type = maybe_init_custom_mem_pool("cuda:0") self.assertTrue(enabled) self.assertEqual(pool, "mock_pool_instance") self.assertEqual(pool_type, "mooncake") mock_init.assert_called_once_with("cuda:0") class TestGetHashStr(unittest.TestCase): def test_hash_str_matches_pre_optimization_per_token_loop(self): for name, tokens, prior_hash in _single_hash_compatibility_cases(): with self.subTest(name=name): self.assertEqual( get_hash_str(tokens, prior_hash), _legacy_get_hash_str(tokens, prior_hash), ) def test_page_hashes_match_pre_optimization_per_token_loop(self): for name, tokens, page_size, prior_hash in _page_hash_compatibility_cases(): with self.subTest(name=name): self.assertEqual( get_hash_str(tokens, prior_hash, page_size=page_size), _legacy_page_hashes(tokens, page_size, prior_hash), ) def test_hash_properties(self): self.assertEqual(get_hash_str([]), hashlib.sha256().hexdigest()) h1 = get_hash_str([1, 2, 3]) h2 = get_hash_str([3, 2, 1]) self.assertNotEqual(h1, h2) self.assertNotEqual(get_hash_str([100]), get_hash_str([200])) self.assertEqual(get_hash_str([1, 2]), get_hash_str([(1, 2)])) self.assertNotEqual(get_hash_str([(1, 2)]), get_hash_str([(2, 1)])) chained = get_hash_str([3, 4], prior_hash=get_hash_str([1, 2])) self.assertNotEqual(chained, get_hash_str([3, 4])) self.assertNotEqual( chained, get_hash_str([3, 4], prior_hash=get_hash_str([9, 9])) ) for tokens in [[], [1], [1, 2, 3], [(1, 2)], [1, 2, 3, 4, 5]]: with self.subTest(tokens=tokens): self.assertRegex(get_hash_str(tokens), r"^[0-9a-f]{64}$") def test_hash_key_hash_page_matches_get_hash_str(self): key = _HashKey(array("q", [1, 2, 3, 4, 5, 6]), is_bigram=True) prior_hash = get_hash_str([(9, 10)]) self.assertEqual( key.hash_page(1, 4, prior_hash), get_hash_str(key[1:4], prior_hash), ) class TestHashStrToInt64(unittest.TestCase): def test_zero_hash(self): result = hash_str_to_int64("0" * 64) self.assertEqual(result, 0) def test_small_positive_value(self): result = hash_str_to_int64("0000000000000001" + "0" * 48) self.assertEqual(result, 1) def test_large_positive_value(self): result = hash_str_to_int64("7fffffffffffffff" + "0" * 48) self.assertEqual(result, 2**63 - 1) def test_negative_overflow(self): result = hash_str_to_int64("8000000000000000" + "0" * 48) self.assertEqual(result, -(2**63)) def test_max_unsigned_maps_to_minus_one(self): result = hash_str_to_int64("f" * 16 + "0" * 48) self.assertEqual(result, -1) def test_only_first_16_chars_matter(self): h1 = hash_str_to_int64("a" * 16 + "0" * 48) h2 = hash_str_to_int64("a" * 16 + "f" * 48) self.assertEqual(h1, h2) def test_roundtrip_with_get_hash_str(self): hash_hex = get_hash_str([42, 99]) int64_val = hash_str_to_int64(hash_hex) self.assertIsInstance(int64_val, int) self.assertTrue(-(2**63) <= int64_val < 2**63) class TestComputeNodeHashValues(unittest.TestCase): def _make_node(self, key, parent=None, parent_hash_values=None): node = MagicMock() node.key = key node.parent = parent node.event_hash_value = None if parent is not None: parent.hash_value = parent_hash_values return node def test_root_node_hashes_match_legacy_page_hashes(self): cases = [ ("single_page", _HashKey(array("q", [1, 2, 3])), 16), ("multiple_pages", _HashKey(array("q", range(1, 31))), 16), ("page_aligned_boundary", _HashKey(array("q", range(1, 33))), 8), ("shorter_than_page", _HashKey(array("q", [1, 2, 3, 4, 5])), 16), ] for name, key, page_size in cases: with self.subTest(name=name): node = self._make_node(key) self.assertEqual( compute_node_hash_values(node, page_size=page_size), _legacy_page_hashes(key, page_size=page_size), ) def test_cache_salt_seeds_root_hash_chain(self): key = _HashKey(array("q", range(1, 17)), cache_salt="tenant-a") seed = hashlib.sha256(b"sglang-cache-salt-v1\0tenant-a").hexdigest() self.assertEqual( compute_node_event_hash_values(self._make_node(key), page_size=8), _legacy_page_hashes(key, page_size=8, prior_hash=seed), ) self.assertEqual( compute_node_hash_values(self._make_node(key), page_size=8), _legacy_page_hashes(key, page_size=8), ) other = _HashKey(array("q", range(1, 17)), cache_salt="tenant-b") self.assertNotEqual( compute_node_event_hash_values(self._make_node(key), page_size=8), compute_node_event_hash_values(self._make_node(other), page_size=8), ) def test_cache_salt_event_hashes_are_memoized(self): node = self._make_node( _HashKey(array("q", range(1, 17)), cache_salt="tenant-a") ) with patch( "sglang.srt.mem_cache.utils.get_hash_str", wraps=get_hash_str ) as mock_get_hash_str: first = compute_node_event_hash_values(node, page_size=8) second = compute_node_event_hash_values(node, page_size=8) self.assertIs(first, second) mock_get_hash_str.assert_called_once() def test_cache_salt_event_hash_walk_is_iterative(self): root = SimpleNamespace( key=_HashKey(array("q")), parent=None, hash_value=[], event_hash_value=None, ) node = root path = [] for token_id in range(1, 1102): node = SimpleNamespace( key=_HashKey(array("q", [token_id]), cache_salt="tenant-a"), parent=node, hash_value=None, event_hash_value=None, ) path.append(node) result = compute_node_event_hash_values(node, page_size=1) self.assertEqual(len(result), 1) self.assertTrue(all(item.event_hash_value is not None for item in path)) def test_parent_hash_is_used_only_when_parent_has_nonempty_key_and_hash(self): parent = MagicMock() parent.key = _HashKey(array("q", range(1, 17))) parent.hash_value = _legacy_page_hashes(parent.key, page_size=8) child_key = _HashKey(array("q", range(101, 109))) cases = [ ("valid_parent_hash", parent, parent.hash_value, parent.hash_value[-1]), ( "empty_parent_key", self._make_node(_HashKey(array("q"))), [get_hash_str([1, 2, 3])], None, ), ( "empty_parent_hash", self._make_node(_HashKey(array("q", range(1, 9)))), [], None, ), ( "none_parent_hash", self._make_node(_HashKey(array("q", range(1, 9)))), None, None, ), ] for name, parent, parent_hash_values, expected_prior in cases: with self.subTest(name=name): child = self._make_node( child_key, parent=parent, parent_hash_values=parent_hash_values ) self.assertEqual( compute_node_hash_values(child, page_size=8), _legacy_page_hashes( child_key, page_size=8, prior_hash=expected_prior ), ) class TestSplitNodeHashValue(unittest.TestCase): def test_none_input_returns_none_tuple(self): result = split_node_hash_value(None, 10, 4) self.assertEqual(result, (None, None)) def test_page_size_one_split(self): hash_values = [f"hash_{i}" for i in range(4)] new_hash, child_hash = split_node_hash_value(hash_values, 2, 1) self.assertEqual(new_hash, ["hash_0", "hash_1"]) self.assertEqual(child_hash, ["hash_2", "hash_3"]) def test_page_size_one_split_at_zero(self): hash_values = ["a", "b", "c"] new_hash, child_hash = split_node_hash_value(hash_values, 0, 1) self.assertEqual(new_hash, []) self.assertEqual(child_hash, ["a", "b", "c"]) def test_page_size_one_split_at_len(self): hash_values = ["a", "b", "c"] new_hash, child_hash = split_node_hash_value(hash_values, 3, 1) self.assertEqual(new_hash, ["a", "b", "c"]) self.assertEqual(child_hash, []) def test_page_size_larger_than_one(self): hash_values = ["p0", "p1", "p2", "p3", "p4", "p5"] new_hash, child_hash = split_node_hash_value(hash_values, 4, 2) self.assertEqual(new_hash, ["p0", "p1"]) self.assertEqual(child_hash, ["p2", "p3", "p4", "p5"]) def test_page_size_split_exact_page_boundary(self): hash_values = ["a", "b", "c", "d"] new_hash, child_hash = split_node_hash_value(hash_values, 4, 4) self.assertEqual(new_hash, ["a"]) self.assertEqual(child_hash, ["b", "c", "d"]) def test_page_size_split_within_page(self): hash_values = ["a", "b", "c"] new_hash, child_hash = split_node_hash_value(hash_values, 6, 4) self.assertEqual(new_hash, ["a"]) self.assertEqual(child_hash, ["b", "c"]) def test_total_length_preserved(self): hash_values = [f"h_{i}" for i in range(10)] for split_len in [0, 1, 3, 5, 7, 9, 10]: for page_size in [1, 2, 4]: new_hash, child_hash = split_node_hash_value( hash_values, split_len, page_size ) self.assertEqual( len(new_hash) + len(child_hash), len(hash_values), f"split_len={split_len}, page_size={page_size}", ) if __name__ == "__main__": unittest.main()