Support RDMA bandwidth calculation for heterogeneous GPU server configurations

deep_ep/utils/envs.py

@@ -162,7 +162,9 @@ def check_nvlink_connections(group: dist.ProcessGroup) -> None:
         # noinspection PyTypeChecker
         devices = os.environ.get('CUDA_VISIBLE_DEVICES', '0,1,2,3,4,5,6,7').strip(',').split(',')
         physical_device_idx = int(devices[torch.cuda.current_device()])
-        physical_device_indices = [0, ] * group.size()
+        physical_device_indices = [
+            0,
+        ] * group.size()
         dist.all_gather_object(physical_device_indices, physical_device_idx, group)
 
         # Check whether they are all connected via NVLink
@@ -202,8 +204,7 @@ def get_nvlink_gbs(factor: float = 0.9) -> float:
     """
     # noinspection PyBroadException
     try:
-        result = subprocess.run(['nvidia-smi', 'nvlink', '-s'],
-                                capture_output=True, text=True, check=True)
+        result = subprocess.run(['nvidia-smi', 'nvlink', '-s'], capture_output=True, text=True, check=True)
         output = result.stdout
         pattern = r'GPU \d+:.*?(?=^GPU \d+:|^$)'
         match = re.search(pattern, output, re.MULTILINE | re.DOTALL)
@@ -242,16 +243,58 @@ def check_fast_rdma_atomic_support(nic_name: str = _DEFAULT_NIC_NAME) -> bool:
         return False
 
 
+def _get_local_gpu_count() -> int:
+    """
+    Get the number of GPUs/ranks sharing local RDMA NICs.
+    """
+
+    try:
+        result = subprocess.run(['nvidia-smi', '-L'], capture_output=True, text=True)
+        count = len(re.findall(r'^GPU \d+:', result.stdout, re.MULTILINE)) if result.returncode == 0 else 0
+        if count > 0:
+            return count
+    except Exception:
+        pass
+
+    if torch.cuda.is_available():
+        count = torch.cuda.device_count()
+        if count > 0:
+            return count
+
+    return 1
+
+
+def _get_rdma_nic_prefix(nic_name: str) -> str:
+    """
+    Get the RDMA NIC prefix used to count peer NIC devices.
+    """
+    match = re.match(r'(.+_)\d+$', nic_name)
+    return match.group(1) if match else nic_name
+
+
+def _get_active_rdma_nic_count(ibstat_output: str, nic_name: str) -> int:
+    """
+    Count active RDMA NICs from `ibstat` output using the selected NIC prefix.
+    """
+    nic_prefix = _get_rdma_nic_prefix(nic_name)
+    nic_pattern = rf'{re.escape(nic_prefix)}\d+' if nic_prefix != nic_name else re.escape(nic_name)
+    ca_blocks = re.findall(r"^CA '[^']+'.*?(?=^CA '|\Z)", ibstat_output, re.MULTILINE | re.DOTALL)
+    count = sum(
+        1 for block in ca_blocks
+        if re.match(rf"^CA '{nic_pattern}'", block) and re.search(r'\bState:\s*Active\b', block) and re.search(r'\bRate:\s*\d+', block))
+    return max(count, 1)
+
+
 @functools.lru_cache()
 def get_rdma_gbs(nic_name: str = _DEFAULT_NIC_NAME) -> float:
     """
-    Get the RDMA bandwidth in GB/s, cached.
+    Get the per-GPU RDMA bandwidth in GB/s, cached.
 
     Arguments:
         nic_name: the NIC device name.
 
     Returns:
-        gbs: the RDMA bandwidth in GB/s (0 if detection fails).
+        gbs: the RDMA bandwidth in GB/s per local GPU/rank (0 if detection fails).
     """
     # noinspection PyBroadException
     try:
@@ -262,7 +305,18 @@ def get_rdma_gbs(nic_name: str = _DEFAULT_NIC_NAME) -> float:
         match = re.search(pattern, output, re.DOTALL)
         assert match
         rate = int(match.group(1))
-        return rate / 8
+        nic_count = _get_active_rdma_nic_count(output, nic_name)
+        # Try to detect 802.3ad LACP bonding slaves and multiply by slave count
+        result = subprocess.run(f'cat /sys/class/infiniband/{nic_name}/device/net/*/upper_*/bonding/slaves',
+                                shell=True,
+                                capture_output=True,
+                                text=True)
+        slave_count = len(result.stdout.strip().split())
+        if slave_count >= 2:
+            nic_count *= slave_count
+        gpu_count = _get_local_gpu_count()
+
+        return rate * nic_count / gpu_count / 8
     except Exception as e:
         print(f'Failed to get RDMA connection speed: {e}')
         return 0

setup.py

@@ -58,16 +58,17 @@ def get_package_version():
         status_output = subprocess.check_output(status_cmd).decode('ascii').strip()
         if status_output:
             print(f'Warning: Git working directory is not clean. Uncommitted changes:\n{status_output}')
-            assert False, 'Git working directory is not clean'
+            raise AssertionError('Git working directory is not clean')
 
         cmd = ['git', 'rev-parse', '--short', 'HEAD']
         revision = '+' + subprocess.check_output(cmd).decode('ascii').rstrip()
-    except:
+    except Exception:
         revision = '+local'
     return f'{public_version}{revision}'
 
 
 class CustomBuildPy(build_py):
+
     def run(self):
         # Make clusters' cache setting default into `envs.py`
         self.generate_default_envs()
@@ -98,9 +99,7 @@ def generate_default_envs(self):
     cxx_flags = ['-O3', '-Wno-deprecated-declarations', '-Wno-unused-variable', '-Wno-sign-compare', '-Wno-reorder', '-Wno-attributes']
     nvcc_flags = ['-O3', '-Xcompiler', '-O3', '--extended-lambda', '--diag-suppress=128,2417']
     sources = ['csrc/python_api.cpp', 'csrc/kernels/legacy/layout.cu', 'csrc/kernels/legacy/intranode.cu']
-    include_dirs = [f'{current_dir}/deep_ep/include',
-                    f'{current_dir}/third-party/fmt/include',
-                    '/usr/local/cuda/include/cccl']
+    include_dirs = [f'{current_dir}/deep_ep/include', f'{current_dir}/third-party/fmt/include', '/usr/local/cuda/include/cccl']
     library_dirs = []
     nvcc_dlink = []
     extra_link_args = ['-lcuda']
@@ -136,7 +135,7 @@ def generate_default_envs(self):
         nvcc_flags.append('-DDISABLE_SM90_FEATURES')
 
         # Disable internode and low-latency kernels
-        assert False, 'Not implemented'
+        raise AssertionError('Not implemented')
     else:
         # Prefer H800 series
         os.environ['TORCH_CUDA_ARCH_LIST'] = os.getenv('TORCH_CUDA_ARCH_LIST', '9.0')
@@ -189,30 +188,29 @@ def generate_default_envs(self):
         if name in os.environ:
             persistent_envs.append((name, os.environ[name]))
     if len(persistent_envs) > 0:
-        print(f' > Persistent envs:')
+        print(' > Persistent envs:')
         for k, v in persistent_envs:
             print(f'   > {k}: {v}')
     print()
 
-    setuptools.setup(
-        name='deep_ep',
-        version=get_package_version(),
-        packages=setuptools.find_packages(include=['deep_ep', 'deep_ep.*']),
-        package_data={
-            'deep_ep': [
-                'include/deep_ep/**/*',
-            ]
-        },
-        ext_modules=[
-            CUDAExtension(name='deep_ep._C',
-                          include_dirs=include_dirs,
-                          library_dirs=library_dirs,
-                          sources=sources,
-                          extra_compile_args=extra_compile_args,
-                          extra_link_args=extra_link_args)
-        ],
-        cmdclass={
-            'build_ext': BuildExtension,
-            'build_py': CustomBuildPy
-        }
-    )
+    setuptools.setup(name='deep_ep',
+                     version=get_package_version(),
+                     packages=setuptools.find_packages(include=['deep_ep', 'deep_ep.*']),
+                     package_data={'deep_ep': [
+                         'include/deep_ep/**/*',
+                     ]},
+                     install_requires=[
+                         "deepxtrace>=0.1.0",
+                     ],
+                     ext_modules=[
+                         CUDAExtension(name='deep_ep._C',
+                                       include_dirs=include_dirs,
+                                       library_dirs=library_dirs,
+                                       sources=sources,
+                                       extra_compile_args=extra_compile_args,
+                                       extra_link_args=extra_link_args)
+                     ],
+                     cmdclass={
+                         'build_ext': BuildExtension,
+                         'build_py': CustomBuildPy
+                     })