[HiCache] add GPU id to IB dev topo for mooncake storage backend (#13112)

This commit is contained in:
Teng Ma
2025-11-17 15:48:29 +08:00
committed by GitHub
parent 1dcde53928
commit 290fcd8971
2 changed files with 30 additions and 1 deletions

View File

@@ -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.

View File

@@ -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: