Add auto bind numa node (#15678)

Signed-off-by: Michael Qiu <qiudayu.qdy@antgroup.com>

python/sglang/srt/environ.py

@@ -485,6 +485,7 @@ class Envs:
 
     # Numa
     SGLANG_NUMA_BIND_V2 = EnvBool(True)
+    SGLANG_AUTO_NUMA_BIND = EnvBool(False)
 
     # Metrics
     SGLANG_ENABLE_METRICS_DEVICE_TIMER = EnvBool(False)

python/sglang/srt/managers/scheduler.py

@@ -199,6 +199,7 @@ from sglang.srt.utils import (
     get_available_gpu_memory,
     get_bool_env_var,
     get_int_env_var,
+    get_numa_node,
     get_zmq_socket,
     kill_itself_when_parent_died,
     numa_bind_to_node,
@@ -3229,10 +3230,14 @@ def run_scheduler_process(
         set_gpu_proc_affinity(
             server_args.pp_size, server_args.tp_size, server_args.nnodes, gpu_id
         )
-    if (
-        numa_node := server_args.numa_node
-    ) is not None and not envs.SGLANG_NUMA_BIND_V2.get():
-        numa_bind_to_node(numa_node[gpu_id])
+    numa_node = None
+    if (numa_nodes := server_args.numa_node) is not None:
+        numa_node = numa_nodes[gpu_id]
+    elif envs.SGLANG_AUTO_NUMA_BIND.get():
+        numa_node = get_numa_node(gpu_id)
+        logger.info(f"auto get NUMA node {numa_node} for GPU {gpu_id}")
+    if numa_node is not None and not envs.SGLANG_NUMA_BIND_V2.get():
+        numa_bind_to_node(numa_node)
 
     # Set up tracing
     if server_args.enable_trace:

python/sglang/srt/utils/common.py

@@ -4180,9 +4180,9 @@ def get_system_nvgpu_count() -> int:
 
 
 @lru_cache(maxsize=1)
-def get_current_device_numa_node_cuda() -> int:
+def get_device_numa_node_cuda(gpu_id: int = 0) -> int:
     """
-    Retrieve the NUMA node ID of the CPU socket closest to the currently active CUDA device.
+    Retrieve the NUMA node ID of the CPU socket closest to the gpu_id.
 
     First tries to query nvidia-smi topology. If it returns a single NUMA ID, uses that directly.
     If it returns multiple NUMA IDs (comma/dash separated), falls back to distributing GPUs
@@ -4196,10 +4196,8 @@ def get_current_device_numa_node_cuda() -> int:
     Raises:
         RuntimeError: If device information cannot be retrieved.
     """
-    import torch
 
-    logical_device_id = torch.cuda.current_device()
-    physical_device_id = get_physical_device_id(logical_device_id)
+    physical_device_id = get_physical_device_id(gpu_id)
 
     # Query NUMA topology from nvidia-smi
     result = subprocess.run(
@@ -4233,6 +4231,32 @@ def get_current_device_numa_node_cuda() -> int:
     return numa_node
 
 
+def get_numa_node(gpu_id):
+    numa_node = None
+    try:
+        device = get_device()
+        if device == "cuda":
+            numa_node = get_device_numa_node_cuda(gpu_id)
+        else:
+            logger.info(f"Now only supports NVIDIA devices")
+    except Exception as e:
+        logger.error(f"Error: {e}")
+
+    return numa_node
+
+
+@lru_cache(maxsize=1)
+def get_current_device_numa_node_cuda() -> int:
+    """
+    Retrieve the NUMA node ID of the CPU socket closest to the currently active CUDA device.
+    """
+
+    logical_device_id = torch.cuda.current_device()
+    numa_node = get_device_numa_node_cuda(logical_device_id)
+
+    return numa_node
+
+
 def nvgpu_available() -> bool:
     if not torch.cuda.is_available():
         return False

python/sglang/srt/utils/numa_utils.py

@@ -8,16 +8,20 @@ from pathlib import Path
 
 from sglang.srt.environ import envs
 from sglang.srt.server_args import ServerArgs
+from sglang.srt.utils import get_numa_node
 
 logger = logging.getLogger(__name__)
 
 
 @contextmanager
 def configure_subprocess(server_args: ServerArgs, gpu_id: int):
-    if (
-        numa_nodes := server_args.numa_node
-    ) is not None and envs.SGLANG_NUMA_BIND_V2.get():
+    numa_node = None
+    if (numa_nodes := server_args.numa_node) is not None:
         numa_node = numa_nodes[gpu_id]
+    elif envs.SGLANG_AUTO_NUMA_BIND.get():
+        numa_node = get_numa_node(gpu_id)
+
+    if numa_node is not None and envs.SGLANG_NUMA_BIND_V2.get():
         numactl_args = f"--cpunodebind={numa_node} --membind={numa_node}"
         executable, debug_str = _create_numactl_executable(numactl_args=numactl_args)
         with _mp_set_executable(executable=executable, debug_str=debug_str):