[HiCache] add GPU id to IB dev topo for mooncake storage backend (#13112)
This commit is contained in:
@@ -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.
|
||||
|
||||
|
||||
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user