From 96724f490c6d8217f043ab6f60b6e68869c621c1 Mon Sep 17 00:00:00 2001 From: Mike Qiu Date: Mon, 9 Mar 2026 14:46:09 +0800 Subject: [PATCH] Add auto bind numa node (#15678) Signed-off-by: Michael Qiu --- python/sglang/srt/environ.py | 1 + python/sglang/srt/managers/scheduler.py | 13 +++++++--- python/sglang/srt/utils/common.py | 34 +++++++++++++++++++++---- python/sglang/srt/utils/numa_utils.py | 10 +++++--- 4 files changed, 46 insertions(+), 12 deletions(-) diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 5399fd2a4..f7f8257a4 100644 --- a/python/sglang/srt/environ.py +++ b/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) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index ec3a254ad..77b65379e 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/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: diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py index 2af599c02..fbbea936e 100644 --- a/python/sglang/srt/utils/common.py +++ b/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 diff --git a/python/sglang/srt/utils/numa_utils.py b/python/sglang/srt/utils/numa_utils.py index 1dac8c5fd..2c934af0d 100644 --- a/python/sglang/srt/utils/numa_utils.py +++ b/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):