Skip to content

Commit d37d89c

Browse files
Functionhxclaude
andcommitted
Add NUMA balancing detection and warning
NUMA balancing auto-migrates pages between NUMA nodes, which causes severe performance degradation for bandwidth-intensive internode dispatch/combine kernels that use RDMA and NVLink. Detect when /proc/sys/kernel/numa_balancing is non-zero at import time and emit a clear warning telling the user how to disable it. Test Plan: - Verified that `import deep_ep` emits a warning when /proc/sys/kernel/numa_balancing contains "1" - Verified that `import deep_ep` is silent when /proc/sys/kernel/numa_balancing contains "0" or is absent Co-Authored-By: Claude <noreply@anthropic.com> Signed-off-by: Yuchen Fan <functionhx@gmail.com>
1 parent 60d4403 commit d37d89c

1 file changed

Lines changed: 25 additions & 1 deletion

File tree

deep_ep/__init__.py

Lines changed: 25 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,11 @@
11
import filecmp
22
import functools
33
import glob
4+
import os
45
import subprocess
6+
import warnings
7+
58
import torch
6-
import os
79

810
from .utils.find_pkgs import find_nccl_root
911

@@ -43,6 +45,27 @@ def find_cuda_home() -> str:
4345
return cuda_home
4446

4547

48+
def check_numa_balancing():
49+
"""
50+
Check whether Linux NUMA balancing is enabled. When enabled, the kernel
51+
auto-migrates pages across NUMA nodes, which can significantly hurt
52+
performance of bandwidth-intensive RDMA/NVLink communication patterns.
53+
"""
54+
if int(os.environ.get('EP_SUPPRESS_NUMA_CHECK', 0)):
55+
return
56+
try:
57+
with open('/proc/sys/kernel/numa_balancing', 'r') as f:
58+
if f.read().strip() != '0':
59+
warnings.warn(
60+
'NUMA balancing is enabled (/proc/sys/kernel/numa_balancing != 0). '
61+
'This can cause severe performance degradation for internode dispatch/combine kernels. '
62+
'Disable it with: echo 0 | sudo tee /proc/sys/kernel/numa_balancing',
63+
stacklevel=2,
64+
)
65+
except OSError:
66+
pass
67+
68+
4669
def check_nccl_so():
4770
"""
4871
Verify that the NCCL library loaded at runtime matches the linked version.
@@ -80,6 +103,7 @@ def init_jit():
80103
find_nccl_root()) # NCCL root
81104

82105
# Run initialization
106+
check_numa_balancing()
83107
check_nccl_so()
84108
init_jit()
85109

0 commit comments

Comments
 (0)