fix AMD CI failure of NUMA binding (#17184)

This commit is contained in:
Zhiqiang Xie
2026-01-16 00:21:04 -08:00
committed by GitHub
parent 7c39ea68f3
commit d9ed80b9f1
2 changed files with 18 additions and 13 deletions
+3 -7
View File
@@ -24,7 +24,7 @@ from sglang.srt.mem_cache.radix_cache import (
split_node_hash_value, split_node_hash_value,
) )
from sglang.srt.metrics.collector import StorageMetricsCollector from sglang.srt.metrics.collector import StorageMetricsCollector
from sglang.srt.utils import bind_to_closest_numa_node, is_numa_available from sglang.srt.utils import bind_to_closest_numa_node_cuda
if TYPE_CHECKING: if TYPE_CHECKING:
from sglang.srt.mem_cache.cache_init_params import CacheInitParams from sglang.srt.mem_cache.cache_init_params import CacheInitParams
@@ -44,12 +44,8 @@ class HiRadixCache(RadixCache):
"Page first layout is not supported with direct IO backend, switching to page first direct layout" "Page first layout is not supported with direct IO backend, switching to page first direct layout"
) )
if ( if not server_args.disable_hicache_numa_detect:
not server_args.disable_hicache_numa_detect bind_to_closest_numa_node_cuda()
and is_numa_available()
and torch.cuda.is_available()
):
bind_to_closest_numa_node()
self.page_size = params.page_size self.page_size = params.page_size
self.kv_cache = params.token_to_kv_pool_allocator.get_kvcache() self.kv_cache = params.token_to_kv_pool_allocator.get_kvcache()
+15 -6
View File
@@ -3956,7 +3956,7 @@ def is_numa_available() -> bool:
return False return False
def get_system_gpu_count() -> int: def get_system_nvgpu_count() -> int:
""" """
Get the total number of GPUs in the system (not affected by CUDA_VISIBLE_DEVICES). Get the total number of GPUs in the system (not affected by CUDA_VISIBLE_DEVICES).
@@ -3978,7 +3978,7 @@ def get_system_gpu_count() -> int:
@lru_cache(maxsize=1) @lru_cache(maxsize=1)
def get_current_device_numa_node() -> int: def get_current_device_numa_node_cuda() -> 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 currently active CUDA device.
@@ -4017,7 +4017,7 @@ def get_current_device_numa_node() -> int:
# Fall back: distribute GPUs evenly across NUMA nodes # Fall back: distribute GPUs evenly across NUMA nodes
numa_count = get_numa_node_count() numa_count = get_numa_node_count()
gpu_count = get_system_gpu_count() gpu_count = get_system_nvgpu_count()
if gpu_count >= numa_count: if gpu_count >= numa_count:
gpus_per_numa = gpu_count // numa_count # >= 1 gpus_per_numa = gpu_count // numa_count # >= 1
@@ -4031,11 +4031,20 @@ def get_current_device_numa_node() -> int:
return numa_node return numa_node
def bind_to_closest_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. 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. Uses `numa` library calls via ctypes to set the CPU affinity of the process.
""" """
node_id = get_current_device_numa_node() if is_numa_available() and nvgpu_available():
numa_bind_to_node(node_id) node_id = get_current_device_numa_node_cuda()
numa_bind_to_node(node_id)