import types import unittest from unittest.mock import patch import torch from sglang.srt.mem_cache.hicache_storage import ( HiCacheStorageConfig, PoolName, PoolTransfer, ) from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase register_cpu_ci(est_time=1, suite="base-a-test-cpu") class ReplicateConfigWithGroupIds: def __init__(self): self.group_ids = None class ReplicateConfigWithoutGroupIds: __slots__ = () class ReplicateConfigWithClassGroupIdsAndRequiredInit: group_ids = None def __init__(self, required): self.group_ids = required def _fake_mooncake_modules(fake_store_cls, replicate_config_cls): mooncake = types.ModuleType("mooncake") mooncake_store = types.ModuleType("mooncake.store") mooncake_store.MooncakeDistributedStore = fake_store_cls mooncake_store.ReplicateConfig = replicate_config_cls return { "mooncake": mooncake, "mooncake.store": mooncake_store, } def _fake_pool_host_mla_module(): pool_host_mla = types.ModuleType("sglang.srt.mem_cache.pool_host.mla") class MLATokenToKVPoolHost: pass pool_host_mla.MLATokenToKVPoolHost = MLATokenToKVPoolHost return pool_host_mla def _fake_pool_host_module(): pool_host = types.ModuleType("sglang.srt.mem_cache.pool_host") class HostKVCache: pass class HostTensorAllocator: pass pool_host.HostKVCache = HostKVCache pool_host.HostTensorAllocator = HostTensorAllocator return pool_host def _fake_host_pool_modules(): return { "sglang.srt.mem_cache.pool_host": _fake_pool_host_module(), "sglang.srt.mem_cache.pool_host.mla": _fake_pool_host_mla_module(), } def _fake_store_class(): class FakeMooncakeDistributedStore: instances = [] def __init__(self): self.batch_put_calls = [] self.existing_keys = set() self.objects = {} type(self).instances.append(self) def setup(self, *args, **kwargs): return 0 def register_buffer(self, *args, **kwargs): return 0 def put(self, key, value, *args): self.objects[key] = value return 0 def is_exist(self, key): return 1 if key in self.objects or key in self.existing_keys else 0 def get(self, key): return self.objects.get(key) def batch_is_exist(self, keys): return [1 if key in self.existing_keys else 0 for key in keys] def batch_put_from(self, keys, ptrs, sizes, *args): self.batch_put_calls.append( { "method": "batch_put_from", "keys": list(keys), "ptrs": list(ptrs), "sizes": list(sizes), "args": args, } ) self.existing_keys.update(keys) return [0] * len(keys) def batch_put_from_multi_buffers(self, keys, ptrs, sizes, *args): self.batch_put_calls.append( { "method": "batch_put_from_multi_buffers", "keys": list(keys), "ptrs": list(ptrs), "sizes": list(sizes), "args": args, } ) self.existing_keys.update(keys) return [0] * len(keys) return FakeMooncakeDistributedStore class FakeHostKVCache: def __init__(self, objects_per_page): self.objects_per_page = objects_per_page self.kv_buffer = torch.empty((1024,), dtype=torch.uint8) self.layout = "page_first" self.page_size = 1 def get_ksize_per_token(self): return 1 def get_page_buffer_meta(self, indices): page_count = len(indices) // self.page_size ptrs = [] sizes = [] for page_idx in range(page_count): for object_idx in range(self.objects_per_page): ptrs.append(1000 + page_idx * 100 + object_idx) sizes.append(8) return ptrs, sizes def get_split_heads_page_buffer_meta(self, indices, split_factor): page_count = len(indices) // self.page_size ptrs = [] sizes = [] for page_idx in range(page_count): for object_idx in range(2 * split_factor): ptrs.append(2000 + page_idx * 100 + object_idx) sizes.append(8) return ptrs, sizes class FakeIndexerPool: page_size = 1 def __init__(self): self.buffer = torch.empty((128,), dtype=torch.uint8) def get_hybrid_pool_buffer(self): return [self.buffer] def get_page_buffer_meta(self, indices): return [3000 + i for i in range(len(indices))], [8] * len(indices) class FakeMultiBufferPool: page_size = 1 def __init__(self): self.buffers = [ torch.empty((128,), dtype=torch.uint8), torch.empty((128,), dtype=torch.uint8), ] def get_hybrid_pool_buffer(self): return self.buffers def get_page_buffer_meta(self, indices): ptrs = [] sizes = [] for i in range(len(indices)): ptrs.extend([4000 + i * 10, 4001 + i * 10]) sizes.extend([8, 16]) return ptrs, sizes def _make_config( *, enable_group_semantics=True, extra_backend_tag=None, model_name=None, is_mla_model=False, should_split_heads=False, tp_rank=0, tp_size=1, tp_lcm_size=None, ): extra_config = { "master_server_address": "127.0.0.1:50051", "check_server": False, "global_segment_size": 1024 * 1024, "enable_group_semantics": enable_group_semantics, } if extra_backend_tag is not None: extra_config["extra_backend_tag"] = extra_backend_tag return HiCacheStorageConfig( tp_rank=tp_rank, tp_size=tp_size, pp_rank=0, pp_size=1, attn_cp_rank=0, attn_cp_size=1, is_mla_model=is_mla_model, enable_storage_metrics=False, is_page_first_layout=True, model_name=model_name, tp_lcm_size=tp_lcm_size, should_split_heads=should_split_heads, extra_config=extra_config, ) def _make_store( *, enable_group_semantics=True, replicate_config_cls=ReplicateConfigWithGroupIds, extra_backend_tag=None, model_name=None, is_mla_model=False, should_split_heads=False, tp_rank=0, tp_size=1, tp_lcm_size=None, ): fake_store_cls = _fake_store_class() cfg = _make_config( enable_group_semantics=enable_group_semantics, extra_backend_tag=extra_backend_tag, model_name=model_name, is_mla_model=is_mla_model, should_split_heads=should_split_heads, tp_rank=tp_rank, tp_size=tp_size, tp_lcm_size=tp_lcm_size, ) with patch.dict( "sys.modules", { **_fake_mooncake_modules(fake_store_cls, replicate_config_cls), **_fake_host_pool_modules(), }, ): from sglang.srt.mem_cache.storage.mooncake_store.mooncake_store import ( MooncakeStore, ) store = MooncakeStore(cfg) return store, fake_store_cls.instances[-1] class TestMooncakeGroupSemantics(CustomTestCase): def test_group_id_detection_uses_class_attribute_without_instantiating(self): fake_store_cls = _fake_store_class() with patch.dict( "sys.modules", { **_fake_mooncake_modules( fake_store_cls, ReplicateConfigWithClassGroupIdsAndRequiredInit, ), **_fake_host_pool_modules(), }, ): from sglang.srt.mem_cache.storage.mooncake_store.mooncake_store import ( MooncakeBaseStore, ) replicate_config_cls, supports_group_ids = ( MooncakeBaseStore()._import_mooncake_group_semantics() ) self.assertIs( replicate_config_cls, ReplicateConfigWithClassGroupIdsAndRequiredInit ) self.assertTrue(supports_group_ids) def test_flag_off_uses_three_arg_batch_put(self): store, fake_store = _make_store(enable_group_semantics=False) store.register_mem_pool_host(FakeHostKVCache(objects_per_page=2)) result = store.batch_set_v1(["page0"], torch.tensor([0])) self.assertEqual(result, [True]) self.assertEqual(len(fake_store.batch_put_calls), 1) self.assertEqual( fake_store.batch_put_calls[0]["keys"], ["page0_0_k", "page0_0_v"] ) self.assertEqual(fake_store.batch_put_calls[0]["args"], ()) def test_mha_group_ids_use_tagged_logical_page_key(self): store, fake_store = _make_store(extra_backend_tag="tag") store.register_mem_pool_host(FakeHostKVCache(objects_per_page=2)) result = store.batch_set_v1(["page0", "page1"], torch.tensor([0, 1])) self.assertEqual(result, [True, True]) call = fake_store.batch_put_calls[0] self.assertEqual( call["keys"], [ "tag_page0_0_k", "tag_page0_0_v", "tag_page1_0_k", "tag_page1_0_v", ], ) self.assertEqual(len(call["args"]), 1) self.assertEqual( call["args"][0].group_ids, [ "sglang-hicache:tag_page0", "sglang-hicache:tag_page0", "sglang-hicache:tag_page1", "sglang-hicache:tag_page1", ], ) def test_old_mooncake_falls_back_to_three_arg_batch_put(self): store, fake_store = _make_store( enable_group_semantics=True, replicate_config_cls=ReplicateConfigWithoutGroupIds, ) store.register_mem_pool_host(FakeHostKVCache(objects_per_page=2)) result = store.batch_set_v1(["page0"], torch.tensor([0])) self.assertEqual(result, [True]) self.assertEqual(fake_store.batch_put_calls[0]["args"], ()) def test_mla_group_ids(self): store, fake_store = _make_store(is_mla_model=True) store.register_mem_pool_host(FakeHostKVCache(objects_per_page=1)) result = store.batch_set_v1(["page0"], torch.tensor([0])) self.assertEqual(result, [True]) call = fake_store.batch_put_calls[0] self.assertEqual(call["keys"], ["page0__k"]) self.assertEqual(call["args"][0].group_ids, ["sglang-hicache:page0"]) def test_split_heads_group_ids(self): store, fake_store = _make_store( should_split_heads=True, tp_rank=1, tp_size=2, tp_lcm_size=4, ) store.register_mem_pool_host(FakeHostKVCache(objects_per_page=4)) result = store.batch_set_v1(["page0"], torch.tensor([0])) self.assertEqual(result, [True]) call = fake_store.batch_put_calls[0] self.assertEqual( call["keys"], ["page0_2_k", "page0_2_v", "page0_3_k", "page0_3_v"], ) self.assertEqual( call["args"][0].group_ids, ["sglang-hicache:page0"] * 4, ) def test_existing_filter_keeps_group_ids_aligned_with_missing_keys(self): store, fake_store = _make_store() store.register_mem_pool_host(FakeHostKVCache(objects_per_page=2)) fake_store.existing_keys.update({"page0_0_k", "page1_0_v"}) result = store.batch_set_v1(["page0", "page1"], torch.tensor([0, 1])) self.assertEqual(result, [True, True]) call = fake_store.batch_put_calls[0] self.assertEqual(call["keys"], ["page0_0_v", "page1_0_k"]) self.assertEqual( call["args"][0].group_ids, ["sglang-hicache:page0", "sglang-hicache:page1"], ) def test_v2_indexer_group_ids_use_logical_page_key(self): store, fake_store = _make_store(extra_backend_tag="tag", is_mla_model=True) indexer_pool = FakeIndexerPool() store.register_mem_host_pool_v2(indexer_pool, PoolName.INDEXER) result = store.batch_set_v2( [ PoolTransfer( name=PoolName.INDEXER, keys=["page0", "page1"], host_indices=torch.tensor([0, 1]), ) ] ) self.assertEqual(result[PoolName.INDEXER], [True, True]) call = fake_store.batch_put_calls[0] self.assertEqual(call["keys"], ["tag_page0__indexer", "tag_page1__indexer"]) self.assertEqual( call["args"][0].group_ids, ["sglang-hicache:tag_page0", "sglang-hicache:tag_page1"], ) def test_v2_multi_buffer_put_passes_group_ids(self): store, fake_store = _make_store(extra_backend_tag="tag", is_mla_model=True) multi_buffer_pool = FakeMultiBufferPool() store.register_mem_host_pool_v2(multi_buffer_pool, PoolName.DEEPSEEK_V4_C4) result = store.batch_set_v2( [ PoolTransfer( name=PoolName.DEEPSEEK_V4_C4, keys=["page0", "page1"], host_indices=torch.tensor([0, 1]), ) ] ) self.assertEqual(result[PoolName.DEEPSEEK_V4_C4], [True, True]) call = fake_store.batch_put_calls[0] self.assertEqual(call["method"], "batch_put_from_multi_buffers") self.assertEqual( call["keys"], ["tag_page0__deepseek_v4_c4", "tag_page1__deepseek_v4_c4"], ) self.assertEqual(call["ptrs"], [[4000, 4001], [4010, 4011]]) self.assertEqual(call["sizes"], [[8, 16], [8, 16]]) self.assertEqual( call["args"][0].group_ids, ["sglang-hicache:tag_page0", "sglang-hicache:tag_page1"], ) def test_model_names_isolate_the_same_logical_key(self): store_a, fake_store_a = _make_store( enable_group_semantics=False, model_name="org/model-a" ) store_b, fake_store_b = _make_store( enable_group_semantics=False, model_name="org/model-b" ) store_a.register_mem_pool_host(FakeHostKVCache(objects_per_page=2)) store_b.register_mem_pool_host(FakeHostKVCache(objects_per_page=2)) self.assertEqual(store_a.batch_set_v1(["page0"], torch.tensor([0])), [True]) self.assertEqual(store_b.batch_set_v1(["page0"], torch.tensor([0])), [True]) keys_a = fake_store_a.batch_put_calls[0]["keys"] keys_b = fake_store_b.batch_put_calls[0]["keys"] self.assertEqual(keys_a, ["org-model-a_page0_0_k", "org-model-a_page0_0_v"]) self.assertEqual(keys_b, ["org-model-b_page0_0_k", "org-model-b_page0_0_v"]) self.assertTrue(set(keys_a).isdisjoint(keys_b)) if __name__ == "__main__": unittest.main(verbosity=3)