minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[NVIDIA] Enable automatic NUMA configuration (#19452)

Browse Source
This commit is contained in:
Trevor Morris
2026-03-27 18:44:13 -07:00
committed by GitHub
parent 83997080a6
commit 7160b6cb76
7 changed files with 479 additions and 191 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/scheduler.py
+2 -8
View File
@@ -207,10 +207,8 @@ from sglang.srt.utils import (
get_available_gpu_memory,
get_bool_env_var,
get_int_env_var,
get_numa_node,
is_mps,
kill_itself_when_parent_died,
numa_bind_to_node,
point_to_point_pyobj,
require_mlp_sync,
set_gpu_proc_affinity,
@@ -224,6 +222,7 @@ from sglang.srt.utils.hf_transformers_utils import (
get_tokenizer_from_processor,
)
from sglang.srt.utils.network import get_zmq_socket
from sglang.srt.utils.numa_utils import get_numa_node_if_available, numa_bind_to_node
from sglang.srt.utils.tensor_bridge import use_mlx
from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter
from sglang.utils import TypeBasedDispatcher, get_exception_traceback
@@ -3542,12 +3541,7 @@ def run_scheduler_process(
set_gpu_proc_affinity(
server_args.pp_size, server_args.tp_size, server_args.nnodes, 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}")
numa_node = get_numa_node_if_available(server_args, gpu_id)
if numa_node is not None and not envs.SGLANG_NUMA_BIND_V2.get():
numa_bind_to_node(numa_node)
python/sglang/srt/mem_cache/hi_mamba_radix_cache.py
-4
View File
@@ -47,7 +47,6 @@ from sglang.srt.mem_cache.radix_cache import (
split_node_hash_value,
)
from sglang.srt.observability.metrics_collector import StorageMetricsCollector
from sglang.srt.utils import bind_to_closest_numa_node_cuda
if TYPE_CHECKING:
from sglang.srt.mem_cache.cache_init_params import CacheInitParams
@@ -104,9 +103,6 @@ class HiMambaRadixCache(MambaRadixCache):
"switching to page first direct layout"
)
if not server_args.disable_hicache_numa_detect:
bind_to_closest_numa_node_cuda()
self.page_size = params.page_size
self.hybrid_kv_cache = params.token_to_kv_pool_allocator.get_kvcache()
if not isinstance(self.hybrid_kv_cache, HybridLinearKVPool):
python/sglang/srt/mem_cache/hiradix_cache.py
-4
View File
@@ -45,7 +45,6 @@ from sglang.srt.mem_cache.radix_cache import (
)
from sglang.srt.mem_cache.utils import convert_to_bigram_key
from sglang.srt.observability.metrics_collector import StorageMetricsCollector
from sglang.srt.utils import bind_to_closest_numa_node_cuda
if TYPE_CHECKING:
from sglang.srt.mem_cache.cache_init_params import CacheInitParams
@@ -59,9 +58,6 @@ class HiRadixCache(RadixCache):
def __init__(self, params: CacheInitParams, server_args: ServerArgs):
self._enable_metrics_flag = params.enable_metrics
if not server_args.disable_hicache_numa_detect:
bind_to_closest_numa_node_cuda()
self.page_size = params.page_size
self.kv_cache = params.token_to_kv_pool_allocator.get_kvcache()
python/sglang/srt/server_args.py
+1 -7
View File
@@ -555,7 +555,6 @@ class ServerArgs:
hicache_write_policy: str = "write_through"
hicache_io_backend: str = "kernel"
hicache_mem_layout: str = "layer_first"
disable_hicache_numa_detect: bool = False
hicache_storage_backend: Optional[str] = None
hicache_storage_prefetch_policy: str = "best_effort"
hicache_storage_backend_extra_config: Optional[str] = None
@@ -5069,11 +5068,6 @@ class ServerArgs:
default=ServerArgs.hicache_mem_layout,
help="The layout of host memory pool for hierarchical cache.",
)
parser.add_argument(
"--disable-hicache-numa-detect",
action="store_true",
help="Disable binding the process to the NUMA node closest to the active CUDA device when hierarchical cache is enabled.",
)
parser.add_argument(
"--hicache-storage-backend",
type=str,
@@ -5542,7 +5536,7 @@ class ServerArgs:
"--numa-node",
type=int,
nargs="+",
help="Sets the numa node for the subprocesses. i-th element corresponds to i-th subprocess.",
help="Sets the numa node for the subprocesses. i-th element corresponds to i-th subprocess. If unset, will be automatically detected on NUMA systems.",
)
parser.add_argument(
"--enable-deterministic-inference",
python/sglang/srt/utils/common.py
-161
View File
@@ -3408,30 +3408,6 @@ def get_device_sm_nvidia_smi():
return (0, 0) # Default/fallback value
def get_libnuma():
libnuma = None
for libnuma_so in ["libnuma.so", "libnuma.so.1"]:
try:
libnuma = ctypes.CDLL(libnuma_so)
except OSError as e:
logger.error(f"{e}")
libnuma = None
if libnuma is not None:
break
return libnuma
def numa_bind_to_node(node: int):
libnuma = get_libnuma()
if libnuma is None or libnuma.numa_available() < 0:
logger.error("numa not available on this system, skip bind action")
else:
libnuma.numa_run_on_node(ctypes.c_int(node))
libnuma.numa_set_preferred(ctypes.c_int(node))
def json_list_type(value):
try:
return orjson.loads(value)
@@ -3729,140 +3705,3 @@ def get_or_create_event_loop():
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
return loop
def get_numa_node_count() -> int:
"""
Get the number of NUMA nodes available on the system.
Must be called after is_numa_available() is True.
Returns:
int: The number of NUMA nodes.
"""
libnuma = get_libnuma()
return libnuma.numa_max_node() + 1
def is_numa_available() -> bool:
try:
libnuma = get_libnuma()
return libnuma.numa_available() >= 0
except Exception:
return False
def get_system_nvgpu_count() -> int:
"""
Get the total number of GPUs in the system (not affected by CUDA_VISIBLE_DEVICES).
Returns:
int: The total number of physical GPUs.
"""
result = subprocess.run(
["nvidia-smi", "--list-gpus"],
capture_output=True,
text=True,
check=True,
)
gpu_lines = [
line
for line in result.stdout.strip().split("\n")
if line.strip().startswith("GPU")
]
return len(gpu_lines)
@lru_cache(maxsize=1)
def get_device_numa_node_cuda(gpu_id: int = 0) -> int:
"""
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
evenly across NUMA nodes based on GPU ID intervals.
For example, with 8 GPUs and 2 NUMA nodes: GPUs 0-3 -> node 0, GPUs 4-7 -> node 1.
Returns:
int: The NUMA node ID (e.g., 0, 1).
Raises:
RuntimeError: If device information cannot be retrieved.
"""
physical_device_id = get_physical_device_id(gpu_id)
# Query NUMA topology from nvidia-smi
result = subprocess.run(
["nvidia-smi", "topo", "-C", "-i", str(physical_device_id)],
capture_output=True,
text=True,
check=True,
)
output_line = result.stdout.strip()
prefix = "NUMA IDs of closest CPU:"
if output_line.startswith(prefix):
numa_id_str = output_line[len(prefix) :].strip()
if numa_id_str.isdigit():
return int(numa_id_str)
# Fall back: distribute GPUs evenly across NUMA nodes
numa_count = get_numa_node_count()
gpu_count = get_system_nvgpu_count()
if gpu_count >= numa_count:
gpus_per_numa = gpu_count // numa_count # >= 1
numa_node = physical_device_id // gpus_per_numa # 0 ~ numa_count - 1
else:
logger.warning(
f"GPU count {gpu_count} is less than NUMA count {numa_count}. Using first NUMA node."
)
numa_node = 0
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
if torch.version.cuda is None:
return False
return True
def bind_to_closest_numa_node_cuda():
"""
Bind the current process to the NUMA node closest to the active CUDA device.
Uses `numa` library calls via ctypes to set the CPU affinity of the process.
"""
if is_numa_available() and nvgpu_available():
node_id = get_current_device_numa_node_cuda()
numa_bind_to_node(node_id)
python/sglang/srt/utils/numa_utils.py
+165 -7
View File
@@ -1,26 +1,30 @@
import ctypes
import glob
import logging
import math
import multiprocessing
import os
import random
import shutil
import time
from contextlib import contextmanager
from pathlib import Path
from typing import Optional
import psutil
from sglang.srt.environ import envs
from sglang.srt.server_args import ServerArgs
from sglang.srt.utils import get_numa_node
from sglang.srt.utils import is_cuda
_is_cuda = is_cuda()
logger = logging.getLogger(__name__)
@contextmanager
def configure_subprocess(server_args: ServerArgs, gpu_id: int):
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)
numa_node = get_numa_node_if_available(server_args, 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)
@@ -58,3 +62,157 @@ def _mp_set_executable(executable: str, debug_str: str):
), f"{multiprocessing.spawn.get_executable()=}"
multiprocessing.spawn.set_executable(old_executable)
logger.info(f"mp.set_executable revert to {old_executable}")
def get_numa_node_if_available(server_args: ServerArgs, gpu_id: int) -> Optional[int]:
"""
Returns the NUMA node for the given GPU id. If it is not set in the server_args, it will try to query the NUMA node for the GPU.
If the NUMA node is not available, has already been configured externally, or the user lacks permission to set NUMA affinity, it will return None.
Args:
server_args: The server arguments.
gpu_id: The GPU id.
Returns:
The NUMA node for the given GPU id or None if it is not available.
"""
if server_args.numa_node is not None:
return server_args.numa_node[gpu_id]
if _is_numa_available():
queried_numa_node = _query_numa_node_for_gpu(gpu_id)
if len(queried_numa_node) == 0:
return None
if len(queried_numa_node) > 1:
# get_numa_node_for_gpu could return multiple nodes, we use the first one for now.
# I don't think there any hardware configs that would have more than one.
logger.warning(
f"Multiple NUMA nodes found for GPU {gpu_id}: {queried_numa_node}. Using the first one."
)
return queried_numa_node[0]
return None
def get_libnuma():
libnuma = None
for libnuma_so in ["libnuma.so", "libnuma.so.1"]:
try:
libnuma = ctypes.CDLL(libnuma_so)
except OSError as e:
logger.debug(f"{e}")
libnuma = None
if libnuma is not None:
break
return libnuma
def numa_bind_to_node(node: int):
libnuma = get_libnuma()
if libnuma is None or libnuma.numa_available() < 0:
logger.warning("numa not available on this system, skip bind action")
else:
libnuma.numa_run_on_node(ctypes.c_int(node))
libnuma.numa_set_preferred(ctypes.c_int(node))
def _can_set_mempolicy() -> bool:
"""Check if the process has permission to use NUMA memory policy syscalls."""
try:
libnuma = get_libnuma()
if libnuma is None or libnuma.numa_available() < 0:
return False
mode = ctypes.c_int()
ret = libnuma.get_mempolicy(
ctypes.byref(mode), None, ctypes.c_ulong(0), None, ctypes.c_ulong(0)
)
return ret == 0
except Exception:
return False
def _is_numa_available() -> bool:
"""
Check if NUMA is available and not already configured externally.
"""
if not _is_cuda:
return False
# Check if this is a numa system.
if not os.path.isdir("/sys/devices/system/node/node1"):
return False
# Check if affinity is already constrained
pid = os.getpid()
process = psutil.Process(pid)
cpu_affinity = process.cpu_affinity()
all_cpus = list(range(psutil.cpu_count()))
constrained_affinity = cpu_affinity != all_cpus
if constrained_affinity:
logger.warning(
"NUMA affinity is already constrained for process, skipping NUMA node configuration for GPU. Remove your constraints to allow automatic configuration."
)
return False
if not shutil.which("numactl") and envs.SGLANG_NUMA_BIND_V2.get():
logger.warning(
"numactl command not found, skipping NUMA node configuration for GPU. Install numactl (e.g., apt-get install numactl) to enable automatic NUMA binding."
)
return False
if not _can_set_mempolicy():
logger.warning(
"User lacks permission to set NUMA affinity, skipping NUMA node configuration for GPU. If using docker, try adding --cap-add SYS_NICE to your docker run command."
)
return False
return True
def _query_numa_node_for_gpu(device_id: int):
"""
Get the NUMA node affinity list for a GPU device.
Args:
device_id: GPU device index.
Returns:
List of NUMA node IDs that have affinity with the device.
"""
try:
import pynvml
except ModuleNotFoundError:
logger.warning("pynvml not installed, skipping NUMA node configuration for GPU")
return []
try:
pynvml.nvmlInit()
handle = pynvml.nvmlDeviceGetHandleByIndex(device_id)
numa_node_count = len(glob.glob("/sys/devices/system/node/node[0-9]*"))
c_ulong_bits = ctypes.sizeof(ctypes.c_ulong) * 8
node_set_size = max(1, math.ceil(numa_node_count / c_ulong_bits))
node_set = pynvml.nvmlDeviceGetMemoryAffinity(
handle,
node_set_size,
pynvml.NVML_AFFINITY_SCOPE_NODE,
)
# Decode the bitmask into a list of NUMA node IDs
numa_nodes = []
for node_id in range(numa_node_count):
mask_array_index = node_id // c_ulong_bits
mask_bit_index = node_id % c_ulong_bits
if node_set[mask_array_index] & (1 << mask_bit_index):
numa_nodes.append(node_id)
return numa_nodes
except pynvml.NVMLError as e:
logger.warning(
f"NVML error querying memory affinity for GPU {device_id}: {e}, skipping NUMA node configuration for GPU"
)
return []
finally:
try:
pynvml.nvmlShutdown()
except Exception:
pass # Ignore shutdown errors