From 44e391b6e9c8f6474c07f293c0fc761779033d53 Mon Sep 17 00:00:00 2001 From: Teng Ma Date: Wed, 12 Nov 2025 15:00:45 +0800 Subject: [PATCH] [PD] Add custom gpu id to device topo support (#12817) Co-authored-by: Shangming Cai --- .../mooncake/transfer_engine.py | 59 ++++++++++++++++++- 1 file changed, 58 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/disaggregation/mooncake/transfer_engine.py b/python/sglang/srt/disaggregation/mooncake/transfer_engine.py index 54657bb46..c938f7e06 100644 --- a/python/sglang/srt/disaggregation/mooncake/transfer_engine.py +++ b/python/sglang/srt/disaggregation/mooncake/transfer_engine.py @@ -1,3 +1,4 @@ +import json import logging from typing import List, Optional @@ -6,6 +7,62 @@ from sglang.srt.utils import get_bool_env_var, get_free_port, maybe_wrap_ipv6_ad logger = logging.getLogger(__name__) +def get_ib_devices_for_gpu(ib_device_str: Optional[str], gpu_id: int) -> Optional[str]: + """ + Parse IB device string and get IB devices for a specific GPU ID. + + Supports both formats: + 1. Old format: "ib0, ib1, ib2" + 2. New format: {0: "ib0, ib1", 1: "ib2, ib3", 2: "ib4"} + + Args: + ib_device_str: The original IB device string + gpu_id: The GPU ID to get devices for + + Returns: + IB devices string for the GPU, or None if not available + """ + if ib_device_str is None or not ib_device_str.strip(): + return None + + ib_device_str = ib_device_str.strip() + + # Check if it's JSON format (new format) + try: + parsed_json = json.loads(ib_device_str) + if isinstance(parsed_json, dict): + # Validate format - keys should be integers (or string rep), values should be strings + gpu_mapping = {} + for gpu_key, ib_devices in parsed_json.items(): + if ( + isinstance(gpu_key, str) + and gpu_key.isdigit() + and isinstance(ib_devices, str) + ): + gpu_mapping[int(gpu_key)] = ib_devices.strip() + elif isinstance(gpu_key, int) and isinstance(ib_devices, str): + gpu_mapping[gpu_key] = ib_devices.strip() + else: + raise ValueError( + f"Invalid format: keys must be integers (or string representations of integers) and values must be strings" + ) + + if not gpu_mapping: + raise ValueError("No valid GPU mappings found in JSON") + + # Return devices for specific GPU + if gpu_id in gpu_mapping: + return gpu_mapping[gpu_id] + else: + raise ValueError( + f"No IB devices configured for GPU {gpu_id}. Available GPUs: {list(gpu_mapping.keys())}" + ) + + except json.JSONDecodeError: + # Not JSON format, treat as old format - return same devices for all GPUs + return ib_device_str + + class MooncakeTransferEngine: def __init__(self, hostname: str, gpu_id: int, ib_device: Optional[str] = None): @@ -21,7 +78,7 @@ class MooncakeTransferEngine: self.engine = TransferEngine() self.hostname = hostname self.gpu_id = gpu_id - self.ib_device = ib_device + self.ib_device = get_ib_devices_for_gpu(ib_device, gpu_id) self.initialize( hostname=self.hostname,