From 329ffc89b9129be3ef9135108cd74f7a699e819b Mon Sep 17 00:00:00 2001 From: Michele Palazzi Date: Thu, 17 Sep 2026 04:51:03 +0200 Subject: [PATCH] [Mooncake] Fix silent SSD offload corruption when TP/PP ranks share ssd_offload_path (#31926) Signed-off-by: Michele Palazzi Co-authored-by: Teng Ma --- .../sglang/srt/managers/cache_controller.py | 1 + .../sglang/srt/mem_cache/hicache_storage.py | 2 + .../storage/mooncake_store/mooncake_store.py | 17 +- .../test_mooncake_ssd_offload_path.py | 202 ++++++++++++++++++ 4 files changed, 221 insertions(+), 1 deletion(-) create mode 100644 test/registered/unit/mem_cache/test_mooncake_ssd_offload_path.py diff --git a/python/sglang/srt/managers/cache_controller.py b/python/sglang/srt/managers/cache_controller.py index d38055c22..b246d8412 100644 --- a/python/sglang/srt/managers/cache_controller.py +++ b/python/sglang/srt/managers/cache_controller.py @@ -745,6 +745,7 @@ class HiCacheController: model_name=model_name, tp_lcm_size=tp_lcm_size, should_split_heads=should_split_heads, + dp_rank=self.dp_rank, extra_config=storage_backend_extra_config, ) diff --git a/python/sglang/srt/mem_cache/hicache_storage.py b/python/sglang/srt/mem_cache/hicache_storage.py index f469daac8..4603cf219 100644 --- a/python/sglang/srt/mem_cache/hicache_storage.py +++ b/python/sglang/srt/mem_cache/hicache_storage.py @@ -37,6 +37,8 @@ class HiCacheStorageConfig: model_name: Optional[str] tp_lcm_size: Optional[int] = None should_split_heads: bool = False + # with dp-attention, tp_rank is attention-group-local; dp_rank disambiguates + dp_rank: int = 0 extra_config: Optional[dict] = None diff --git a/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py b/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py index 097532d9f..45c1e9023 100644 --- a/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py +++ b/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py @@ -504,7 +504,22 @@ class MooncakeStore(HiCacheStorage, MooncakeBaseStore): if self.config.enable_ssd_offload: setup_kwargs["enable_ssd_offload"] = True if self.config.ssd_offload_path is not None: - setup_kwargs["ssd_offload_path"] = self.config.ssd_offload_path + # Each rank embeds its own Mooncake client. Sharing one + # offload directory corrupts silently: bucket ids are + # generated per process and resumed from the same startup + # scan after a restart, and bucket files are opened with + # O_CREAT|O_TRUNC, so a filename collision truncates + # another rank's bucket. Give every rank a private subdir. + ssd_offload_path = self.config.ssd_offload_path + if storage_config is not None: + ssd_offload_path = os.path.join( + ssd_offload_path, + f"rank_{storage_config.dp_rank}" + f"_{storage_config.tp_rank}_{storage_config.pp_rank}" + f"_{storage_config.attn_cp_rank}", + ) + os.makedirs(ssd_offload_path, exist_ok=True) + setup_kwargs["ssd_offload_path"] = ssd_offload_path if self.config.tenant_id != DEFAULT_TENANT_ID: setup_kwargs["tenant_id"] = self.config.tenant_id diff --git a/test/registered/unit/mem_cache/test_mooncake_ssd_offload_path.py b/test/registered/unit/mem_cache/test_mooncake_ssd_offload_path.py new file mode 100644 index 000000000..f7c19fdeb --- /dev/null +++ b/test/registered/unit/mem_cache/test_mooncake_ssd_offload_path.py @@ -0,0 +1,202 @@ +import os +import tempfile +import types +import unittest +from unittest.mock import patch + +from sglang.srt.mem_cache.hicache_storage import HiCacheStorageConfig +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") + + +def _fake_store_class(): + class FakeMooncakeDistributedStore: + instances = [] + + def __init__(self): + self.setup_kwargs = None + self.objects = {} + type(self).instances.append(self) + + def setup(self, *args, **kwargs): + self.setup_kwargs = kwargs + return 0 + + def register_buffer(self, *args, **kwargs): + return 0 + + def put(self, key, value, *args): + self.objects[key] = value + return 0 + + def get(self, key): + return self.objects.get(key) + + def is_exist(self, key): + return 1 if key in self.objects else 0 + + def batch_is_exist(self, keys): + return [1 if key in self.objects else 0 for key in keys] + + def remove(self, key): + self.objects.pop(key, None) + return 0 + + return FakeMooncakeDistributedStore + + +def _fake_mooncake_modules(fake_store_cls): + mooncake = types.ModuleType("mooncake") + mooncake_store = types.ModuleType("mooncake.store") + mooncake_store.MooncakeDistributedStore = fake_store_cls + return { + "mooncake": mooncake, + "mooncake.store": mooncake_store, + } + + +def _fake_host_pool_modules(): + pool_host = types.ModuleType("sglang.srt.mem_cache.pool_host") + + class HostKVCache: + pass + + class HostTensorAllocator: + pass + + pool_host.HostKVCache = HostKVCache + pool_host.HostTensorAllocator = HostTensorAllocator + + pool_host_mla = types.ModuleType("sglang.srt.mem_cache.pool_host.mla") + + class MLATokenToKVPoolHost: + pass + + pool_host_mla.MLATokenToKVPoolHost = MLATokenToKVPoolHost + return { + "sglang.srt.mem_cache.pool_host": pool_host, + "sglang.srt.mem_cache.pool_host.mla": pool_host_mla, + } + + +def _make_config( + *, + tp_rank, + pp_rank, + ssd_offload_path, + dp_rank=0, + tp_size=8, + attn_cp_rank=0, + attn_cp_size=1, +): + return HiCacheStorageConfig( + tp_rank=tp_rank, + tp_size=tp_size, + pp_rank=pp_rank, + pp_size=1, + dp_rank=dp_rank, + attn_cp_rank=attn_cp_rank, + attn_cp_size=attn_cp_size, + is_mla_model=False, + enable_storage_metrics=False, + is_page_first_layout=True, + model_name="test", + extra_config={ + "master_server_address": "127.0.0.1:50051", + "check_server": False, + "global_segment_size": 1024 * 1024, + "enable_ssd_offload": True, + "ssd_offload_path": ssd_offload_path, + }, + ) + + +def _make_store( + *, + tp_rank, + pp_rank, + ssd_offload_path, + dp_rank=0, + tp_size=8, + attn_cp_rank=0, + attn_cp_size=1, +): + fake_store_cls = _fake_store_class() + cfg = _make_config( + tp_rank=tp_rank, + pp_rank=pp_rank, + ssd_offload_path=ssd_offload_path, + dp_rank=dp_rank, + tp_size=tp_size, + attn_cp_rank=attn_cp_rank, + attn_cp_size=attn_cp_size, + ) + with patch.dict( + "sys.modules", + { + **_fake_mooncake_modules(fake_store_cls), + **_fake_host_pool_modules(), + }, + ): + from sglang.srt.mem_cache.storage.mooncake_store.mooncake_store import ( + MooncakeStore, + ) + + MooncakeStore(cfg) + return fake_store_cls.instances[-1] + + +class TestMooncakeSsdOffloadPath(CustomTestCase): + def test_each_rank_gets_private_subdirectory(self): + with tempfile.TemporaryDirectory() as base: + for tp_rank in (0, 3): + client = _make_store(tp_rank=tp_rank, pp_rank=0, ssd_offload_path=base) + expected = os.path.join(base, f"rank_0_{tp_rank}_0_0") + self.assertEqual(client.setup_kwargs.get("ssd_offload_path"), expected) + self.assertTrue(os.path.isdir(expected)) + + def test_rank_directories_are_distinct(self): + with tempfile.TemporaryDirectory() as base: + paths = { + _make_store( + tp_rank=tp_rank, pp_rank=0, ssd_offload_path=base + ).setup_kwargs["ssd_offload_path"] + for tp_rank in range(4) + } + self.assertEqual(len(paths), 4) + + def test_dp_ranks_are_distinct_when_attn_tp_rank_is_zero(self): + # dp-attention with attn_tp_size 1: every DP rank reports tp_rank 0 + with tempfile.TemporaryDirectory() as base: + paths = { + _make_store( + tp_rank=0, pp_rank=0, ssd_offload_path=base, dp_rank=dp_rank + ).setup_kwargs["ssd_offload_path"] + for dp_rank in range(8) + } + self.assertEqual(len(paths), 8) + self.assertIn(os.path.join(base, "rank_5_0_0_0"), paths) + + def test_cp_ranks_are_distinct_with_dp_attention(self): + """CP peers with identical DP/TP/PP ranks must not share SSD files.""" + with tempfile.TemporaryDirectory() as base: + paths = { + _make_store( + tp_rank=0, + pp_rank=0, + ssd_offload_path=base, + dp_rank=dp_rank, + tp_size=1, + attn_cp_rank=cp_rank, + attn_cp_size=4, + ).setup_kwargs["ssd_offload_path"] + for dp_rank in range(2) + for cp_rank in range(4) + } + self.assertEqual(len(paths), 8) + + +if __name__ == "__main__": + unittest.main()