From d9ed80b9f195c98dacf9979ada35d24b56aa7bfc Mon Sep 17 00:00:00 2001 From: Zhiqiang Xie Date: Fri, 16 Jan 2026 00:21:04 -0800 Subject: [PATCH] fix AMD CI failure of NUMA binding (#17184) --- python/sglang/srt/mem_cache/hiradix_cache.py | 10 +++------- python/sglang/srt/utils/common.py | 21 ++++++++++++++------ 2 files changed, 18 insertions(+), 13 deletions(-) diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index 5a9e8715b..cb3c8b24b 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -24,7 +24,7 @@ from sglang.srt.mem_cache.radix_cache import ( split_node_hash_value, ) 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: 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" ) - if ( - not server_args.disable_hicache_numa_detect - and is_numa_available() - and torch.cuda.is_available() - ): - bind_to_closest_numa_node() + 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() diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py index 31e1074e1..d162d730f 100644 --- a/python/sglang/srt/utils/common.py +++ b/python/sglang/srt/utils/common.py @@ -3956,7 +3956,7 @@ def is_numa_available() -> bool: 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). @@ -3978,7 +3978,7 @@ def get_system_gpu_count() -> int: @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. @@ -4017,7 +4017,7 @@ def get_current_device_numa_node() -> int: # Fall back: distribute GPUs evenly across NUMA nodes numa_count = get_numa_node_count() - gpu_count = get_system_gpu_count() + gpu_count = get_system_nvgpu_count() if gpu_count >= numa_count: gpus_per_numa = gpu_count // numa_count # >= 1 @@ -4031,11 +4031,20 @@ def get_current_device_numa_node() -> int: 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. Uses `numa` library calls via ctypes to set the CPU affinity of the process. """ - node_id = get_current_device_numa_node() - numa_bind_to_node(node_id) + if is_numa_available() and nvgpu_available(): + node_id = get_current_device_numa_node_cuda() + numa_bind_to_node(node_id)