From 290fcd897163c2699209614ed76916c17cf2b310 Mon Sep 17 00:00:00 2001 From: Teng Ma Date: Mon, 17 Nov 2025 15:48:29 +0800 Subject: [PATCH] [HiCache] add GPU id to IB dev topo for mooncake storage backend (#13112) --- .../storage/mooncake_store/README.md | 10 +++++++++ .../storage/mooncake_store/mooncake_store.py | 21 ++++++++++++++++++- 2 files changed, 30 insertions(+), 1 deletion(-) 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 40f8c8655..4ccc8ca7e 100644 --- a/python/sglang/srt/mem_cache/storage/mooncake_store/README.md +++ b/python/sglang/srt/mem_cache/storage/mooncake_store/README.md @@ -104,6 +104,16 @@ Parameter Explanation: * `master_server_address`: The network address of the `master service`. The default port is 50051. * `protocol`: The protocol used by Mooncake. Supported values are `"rdma"` or `"tcp"`. For optimal performance, `"rdma"` is recommended. * `device_name`: For `"rdma"`, you can leave this empty in most cases. Mooncake auto-discovers RDMA NICs by default. If you want to pin specific NICs (e.g., `mlx5_0,mlx5_1`), just set `device_name` accordingly. To list available devices, use `ibv_devices`. + - For tensor parallel deployments where different ranks should use different devices, you can specify device configurations using JSON format: + ```json + { + "device_name": "{0: \"ib0,ib1\", 1: \"ib2,ib3\", 2: \"ib4,ib5\"}" + } + ``` + - Or in environment variables: + ```bash + MOONCAKE_DEVICE="{\"0\": \"ib0,ib1\", \"1\": \"ib2,ib3\", \"2\": \"ib4,ib5\"}" + ``` * `global_segment_size`: The amount of memory contributed to the global memory pool. Accepts either bytes (integer) or a string with the `gb` suffix, e.g., `"16gb"`. A larger value allows Mooncake to cache more KV tensors. * `local_buffer_size`: Local buffer is used to do request operations such as `Get` or `Put`. In this case, it is set to 0 because the instance functions solely as a storage server, contributing memory to the global pool without issuing any request operations. 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 040afee32..73d1ff013 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 @@ -190,13 +190,32 @@ class MooncakeStore(HiCacheStorage): if self.config.check_server: self.check_server() + # Handle JSON device_name configuration + device_name = self.config.device_name + if device_name and device_name.strip().startswith("{"): + try: + device_config = json.loads(device_name) + if storage_config and hasattr(storage_config, "tp_rank"): + tp_rank = storage_config.tp_rank + # Try both integer and string keys since JSON parsing may convert keys + device_name = device_config.get(tp_rank, "") + if not device_name: + device_name = device_config.get(str(tp_rank), "") + else: + device_name = "" + except (json.JSONDecodeError, AttributeError): + logger.warning( + f"Failed to parse device_name as JSON: {device_name}" + ) + device_name = "" + ret_code = self.store.setup( self.config.local_hostname, self.config.metadata_server, per_tp_global_segment_size, per_tp_local_buffer_size, self.config.protocol, - self.config.device_name, + device_name, self.config.master_server_address, ) if ret_code: