From be156d68043171d2d9b9f6fbd7b9f7d90dde7b21 Mon Sep 17 00:00:00 2001 From: Teng Ma Date: Thu, 14 May 2026 14:07:02 +0800 Subject: [PATCH] [HiCache] enable ssd offload support for mooncake store (#24277) Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: stmatengss <11641725+stmatengss@users.noreply.github.com> Co-authored-by: Zhangheng --- python/sglang/srt/environ.py | 2 + .../storage/mooncake_store/README.md | 39 ++++++++++++ .../storage/mooncake_store/mooncake_store.py | 60 +++++++++++++++---- 3 files changed, 91 insertions(+), 10 deletions(-) diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 34560e6a5..7c1752a72 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -319,6 +319,8 @@ class Envs: MOONCAKE_MASTER_METRICS_PORT = EnvInt(9003) MOONCAKE_CHECK_SERVER = EnvBool(False) MOONCAKE_STANDALONE_STORAGE = EnvBool(False) + MOONCAKE_ENABLE_SSD_OFFLOAD = EnvBool(False) + MOONCAKE_OFFLOAD_FILE_STORAGE_PATH = EnvStr(None) # AMD & ROCm SGLANG_USE_AITER = EnvBool(False) diff --git a/python/sglang/srt/mem_cache/storage/mooncake_store/README.md b/python/sglang/srt/mem_cache/storage/mooncake_store/README.md index e9eadb0ee..a88d4ad74 100644 --- a/python/sglang/srt/mem_cache/storage/mooncake_store/README.md +++ b/python/sglang/srt/mem_cache/storage/mooncake_store/README.md @@ -250,6 +250,45 @@ In particular, for the `global segment size`, if at least one `store service` in **Important:** when `tp > 1`, each Tensor Parallel (TP) rank launches its own Mooncake backend instance and contributes `1/global_segment_size` memory. Therefore, the total memory consumption equals `global segment size`. +**SSD Offload (`enable_ssd_offload`):** + +When `enable_ssd_offload` is set to `true`, SGLang will request that Mooncake enable SSD offloading for the KV cache. This allows Mooncake to spill overflow data from DRAM to local SSDs, effectively expanding the available L3 cache capacity. + +If you need to explicitly control the SSD spill directory, set `ssd_offload_path` or the `MOONCAKE_OFFLOAD_FILE_STORAGE_PATH` environment variable. SGLang forwards this value to `MooncakeDistributedStore.setup(..., ssd_offload_path=...)`, while other SSD offload tuning parameters continue to be read directly by the Mooncake C++ library. + +You can enable it in any of the three supported configuration methods: + +- **Via `--hicache-storage-backend-extra-config`:** + ```bash + python -m sglang.launch_server \ + --enable-hierarchical-cache \ + --hicache-storage-backend mooncake \ + --model-path [model_path] \ + --hicache-storage-backend-extra-config '{"master_server_address": "127.0.0.1:50051", "enable_ssd_offload": true, "ssd_offload_path": "/mnt/mooncake-ssd"}' + ``` + +- **Via JSON config file (`SGLANG_HICACHE_MOONCAKE_CONFIG_PATH`):** + ```json + { + "master_server_address": "127.0.0.1:50051", + "enable_ssd_offload": true, + "ssd_offload_path": "/mnt/mooncake-ssd" + } + ``` + +- **Via environment variable:** + ```bash + MOONCAKE_MASTER="127.0.0.1:50051" \ + MOONCAKE_ENABLE_SSD_OFFLOAD=1 \ + MOONCAKE_OFFLOAD_FILE_STORAGE_PATH="/mnt/mooncake-ssd" \ + python -m sglang.launch_server \ + --enable-hierarchical-cache \ + --hicache-storage-backend mooncake \ + --model-path [model_path] + ``` + +> **Note:** `enable_ssd_offload` requires a Mooncake version that supports the `enable_ssd_offload` parameter in `MooncakeDistributedStore.setup()`. If the installed version does not support it, SGLang will automatically fall back to the old behavior and print a warning. + **HiCache Related Parameters for SGLang Server** For a comprehensive overview of HiCache-related parameters, please refer to [this document](https://docs.sglang.io/advanced_features/hicache_design.html#related-parameters). 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 b9be8eb0b..d25b7811a 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 @@ -92,6 +92,8 @@ class MooncakeStoreConfig: check_server: bool standalone_storage: bool client_server_address: str + enable_ssd_offload: bool = False + ssd_offload_path: Optional[str] = None @staticmethod def from_file() -> "MooncakeStoreConfig": @@ -142,6 +144,12 @@ class MooncakeStoreConfig: client_server_address=config.get( "client_server_address", envs.MOONCAKE_CLIENT.default ), + enable_ssd_offload=config.get( + "enable_ssd_offload", envs.MOONCAKE_ENABLE_SSD_OFFLOAD.default + ), + ssd_offload_path=config.get( + "ssd_offload_path", envs.MOONCAKE_OFFLOAD_FILE_STORAGE_PATH.default + ), ) @staticmethod @@ -181,6 +189,8 @@ class MooncakeStoreConfig: check_server=envs.MOONCAKE_CHECK_SERVER.get(), standalone_storage=envs.MOONCAKE_STANDALONE_STORAGE.get(), client_server_address=envs.MOONCAKE_CLIENT.get(), + enable_ssd_offload=envs.MOONCAKE_ENABLE_SSD_OFFLOAD.get(), + ssd_offload_path=envs.MOONCAKE_OFFLOAD_FILE_STORAGE_PATH.get(), ) @staticmethod @@ -223,6 +233,12 @@ class MooncakeStoreConfig: client_server_address=extra_config.get( "client_server_address", envs.MOONCAKE_CLIENT.default ), + enable_ssd_offload=extra_config.get( + "enable_ssd_offload", envs.MOONCAKE_ENABLE_SSD_OFFLOAD.default + ), + ssd_offload_path=extra_config.get( + "ssd_offload_path", envs.MOONCAKE_OFFLOAD_FILE_STORAGE_PATH.default + ), ) @@ -373,16 +389,40 @@ class MooncakeStore(HiCacheStorage, MooncakeBaseStore): client_hostname = self.config.local_hostname transfer_engine = None - ret_code = self.store.setup( - client_hostname, - self.config.metadata_server, - per_tp_global_segment_size, - DEFAULT_LOCAL_BUFFER_SIZE, # Zero copy interface does not need local buffer - self.config.protocol, - device_name, - self.config.master_server_address, - transfer_engine, - ) + setup_kwargs = {} + 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 + + while True: + try: + ret_code = self.store.setup( + client_hostname, + self.config.metadata_server, + per_tp_global_segment_size, + DEFAULT_LOCAL_BUFFER_SIZE, # Zero copy interface does not need local buffer + self.config.protocol, + device_name, + self.config.master_server_address, + transfer_engine, + **setup_kwargs, + ) + break + except TypeError as e: + unsupported_kwargs = [ + key for key in list(setup_kwargs) if key in str(e) + ] + if not unsupported_kwargs: + raise + logger.warning( + "The installed Mooncake version does not support the " + f"{', '.join(unsupported_kwargs)} parameter(s) in setup(). " + f"Retrying without {', '.join(unsupported_kwargs)}. " + "Please upgrade Mooncake to enable SSD offload support." + ) + for key in unsupported_kwargs: + setup_kwargs.pop(key, None) if ret_code: raise RuntimeError( f"Failed to setup Mooncake store, error code: {ret_code}"