From 54c24b29a826c5422a31246fe54ff27d775029a9 Mon Sep 17 00:00:00 2001 From: sufubao Date: Wed, 15 Jul 2026 13:54:26 +0800 Subject: [PATCH 1/3] feat: interleave cpu kv cache pages across numa nodes via mbind --- lightllm/utils/kv_cache_utils.py | 99 +++++++++++++++++++++++++++----- 1 file changed, 85 insertions(+), 14 deletions(-) diff --git a/lightllm/utils/kv_cache_utils.py b/lightllm/utils/kv_cache_utils.py index 06c4211eb..eaad2fa18 100644 --- a/lightllm/utils/kv_cache_utils.py +++ b/lightllm/utils/kv_cache_utils.py @@ -170,6 +170,87 @@ def get_merged_head_dim(self): ) // self.data_type.itemsize +def _get_default_hugepage_size() -> int: + try: + with open("/proc/meminfo", "r") as f: + for line in f: + if line.startswith("Hugepagesize:"): + parts = line.split() + if len(parts) >= 2: + kb = int(parts[1]) + return kb * 1024 + except Exception: + pass + return 2 * 1024 * 1024 + + +def _get_online_numa_nodes() -> List[int]: + for path in ("/sys/devices/system/node/has_memory", "/sys/devices/system/node/online"): + try: + with open(path, "r") as f: + online = f.read().strip() + nodes: List[int] = [] + for part in online.split(","): + if "-" in part: + start, end = part.split("-") + nodes.extend(range(int(start), int(end) + 1)) + else: + nodes.append(int(part)) + return nodes + except Exception: + continue + return [0] + + +def interleave_pages_across_numa_nodes(libc, addr: int, size: int) -> bool: + MPOL_INTERLEAVE = 3 + SYS_MBIND = {"x86_64": 237, "aarch64": 235}.get(os.uname().machine) + + if SYS_MBIND is None: + logger.warning(f"unsupported architecture {os.uname().machine}, skip cpu cache numa interleave") + return False + + if enable_huge_page(): + huge_sz = _get_default_hugepage_size() + size = triton.cdiv(size, huge_sz) * huge_sz + + def _mbind(mode, mask): + nodemask = ctypes.c_ulong(mask) + libc.syscall.restype = ctypes.c_long + return libc.syscall( + ctypes.c_long(SYS_MBIND), + ctypes.c_void_p(addr), + ctypes.c_ulong(size), + ctypes.c_int(mode), + ctypes.byref(nodemask), + ctypes.c_ulong(ctypes.sizeof(nodemask) * 8 + 1), + ctypes.c_uint(0), + ) + + if os.getenv("LIGHTLLM_DISABLE_NUMA_INTERLEAVE", "0") == "1": + logger.info("cpu cache numa interleave disabled by LIGHTLLM_DISABLE_NUMA_INTERLEAVE") + return False + nodes = _get_online_numa_nodes() + if len(nodes) <= 1: + return False + if max(nodes) >= 64: + logger.warning(f"more than 64 numa nodes ({nodes}), skip cpu cache numa interleave") + return False + try: + ret = _mbind(MPOL_INTERLEAVE, sum(1 << n for n in nodes)) + if ret != 0: + logger.warning( + f"mbind MPOL_INTERLEAVE failed (errno={ctypes.get_errno()}), " + f"cpu kv cache pages will use default first-touch numa policy" + ) + return False + logger.info(f"cpu kv cache pages interleaved across numa nodes {nodes}") + return True + except Exception as e: + logger.warning(f"cpu cache numa interleave skipped: {e}") + return False + + @lru_cache(maxsize=None) def create_shm_kv_cache_ptr(key: int, size: int) -> int: libc = ctypes.CDLL("/usr/lib/x86_64-linux-gnu/libc.so.6", use_errno=True) @@ -181,20 +262,6 @@ def create_shm_kv_cache_ptr(key: int, size: int) -> int: requested_size = size use_hugetlb = enable_huge_page() - # 计算大页大小（默认从 /proc/meminfo 读取 Hugepagesize） - def _get_default_hugepage_size() -> int: - try: - with open("/proc/meminfo", "r") as f: - for line in f: - if line.startswith("Hugepagesize:"): - parts = line.split() - if len(parts) >= 2: - kb = int(parts[1]) - return kb * 1024 - except Exception: - pass - return 2 * 1024 * 1024 # fallback 2MB - shmflg = 0o666 | 0o1000 # 权限和 IPC_CREAT 标志 if use_hugetlb: # 向上对齐到大页大小 @@ -234,6 +301,8 @@ def _get_default_hugepage_size() -> int: raise Exception("Error attaching shared memory") logger.info(f"Shared cpu kv cache tensor memory at address: {shm_addr}") + interleave_pages_across_numa_nodes(libc, shm_addr, size_to_alloc) + # Best-effort memory prefaulting in background to speed up subsequent cudaHostRegister def _pre_warm_memory(): page_size = _get_default_hugepage_size() if use_hugetlb else 4096 @@ -327,4 +396,6 @@ def attach_shm_kv_cache_ptr(key: int, size: int) -> int: raise Exception(f"Error attaching shared memory (errno={err})") logger.info(f"Attached to SHM key={key}, shmid={shmid}, addr={shm_addr}") + + interleave_pages_across_numa_nodes(libc, shm_addr, size) return shm_addr From 1390cf6986cf1e22c0e0640bca2a0fa97623c371 Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Mon, 3 Aug 2026 02:18:04 +0000 Subject: [PATCH 2/3] refactor: organize cpu cache NUMA interleave helpers --- lightllm/utils/envs_utils.py | 6 + lightllm/utils/kv_cache_utils.py | 186 +++++++++++++++++-------------- 2 files changed, 111 insertions(+), 81 deletions(-) diff --git a/lightllm/utils/envs_utils.py b/lightllm/utils/envs_utils.py index 4dca81ae6..8ee0fd982 100644 --- a/lightllm/utils/envs_utils.py +++ b/lightllm/utils/envs_utils.py @@ -238,6 +238,12 @@ def enable_huge_page(): return enable_env_vars("LIGHTLLM_HUGE_PAGE_ENABLE") +@lru_cache(maxsize=None) +def disable_cpu_cache_numa_interleave() -> bool: + """是否禁用 CPU KV cache 共享内存的 NUMA 交错分配策略。""" + return enable_env_vars("LIGHTLLM_DISABLE_NUMA_INTERLEAVE") + + @lru_cache(maxsize=None) def get_added_mtp_kv_layer_num() -> int: # mtp 模式下需要在mem manger上扩展draft model使用的layer diff --git a/lightllm/utils/kv_cache_utils.py b/lightllm/utils/kv_cache_utils.py index eaad2fa18..e30e093e4 100644 --- a/lightllm/utils/kv_cache_utils.py +++ b/lightllm/utils/kv_cache_utils.py @@ -11,6 +11,7 @@ from lightllm.utils.envs_utils import ( get_env_start_args, enable_huge_page, + disable_cpu_cache_numa_interleave, get_llm_data_type, get_added_mtp_kv_layer_num, ) @@ -170,87 +171,6 @@ def get_merged_head_dim(self): ) // self.data_type.itemsize -def _get_default_hugepage_size() -> int: - try: - with open("/proc/meminfo", "r") as f: - for line in f: - if line.startswith("Hugepagesize:"): - parts = line.split() - if len(parts) >= 2: - kb = int(parts[1]) - return kb * 1024 - except Exception: - pass - return 2 * 1024 * 1024 - - -def _get_online_numa_nodes() -> List[int]: - for path in ("/sys/devices/system/node/has_memory", "/sys/devices/system/node/online"): - try: - with open(path, "r") as f: - online = f.read().strip() - nodes: List[int] = [] - for part in online.split(","): - if "-" in part: - start, end = part.split("-") - nodes.extend(range(int(start), int(end) + 1)) - else: - nodes.append(int(part)) - return nodes - except Exception: - continue - return [0] - - -def interleave_pages_across_numa_nodes(libc, addr: int, size: int) -> bool: - MPOL_INTERLEAVE = 3 - SYS_MBIND = {"x86_64": 237, "aarch64": 235}.get(os.uname().machine) - - if SYS_MBIND is None: - logger.warning(f"unsupported architecture {os.uname().machine}, skip cpu cache numa interleave") - return False - - if enable_huge_page(): - huge_sz = _get_default_hugepage_size() - size = triton.cdiv(size, huge_sz) * huge_sz - - def _mbind(mode, mask): - nodemask = ctypes.c_ulong(mask) - libc.syscall.restype = ctypes.c_long - return libc.syscall( - ctypes.c_long(SYS_MBIND), - ctypes.c_void_p(addr), - ctypes.c_ulong(size), - ctypes.c_int(mode), - ctypes.byref(nodemask), - ctypes.c_ulong(ctypes.sizeof(nodemask) * 8 + 1), - ctypes.c_uint(0), - ) - - if os.getenv("LIGHTLLM_DISABLE_NUMA_INTERLEAVE", "0") == "1": - logger.info("cpu cache numa interleave disabled by LIGHTLLM_DISABLE_NUMA_INTERLEAVE") - return False - nodes = _get_online_numa_nodes() - if len(nodes) <= 1: - return False - if max(nodes) >= 64: - logger.warning(f"more than 64 numa nodes ({nodes}), skip cpu cache numa interleave") - return False - try: - ret = _mbind(MPOL_INTERLEAVE, sum(1 << n for n in nodes)) - if ret != 0: - logger.warning( - f"mbind MPOL_INTERLEAVE failed (errno={ctypes.get_errno()}), " - f"cpu kv cache pages will use default first-touch numa policy" - ) - return False - logger.info(f"cpu kv cache pages interleaved across numa nodes {nodes}") - return True - except Exception as e: - logger.warning(f"cpu cache numa interleave skipped: {e}") - return False - - @lru_cache(maxsize=None) def create_shm_kv_cache_ptr(key: int, size: int) -> int: libc = ctypes.CDLL("/usr/lib/x86_64-linux-gnu/libc.so.6", use_errno=True) @@ -399,3 +319,107 @@ def attach_shm_kv_cache_ptr(key: int, size: int) -> int: interleave_pages_across_numa_nodes(libc, shm_addr, size) return shm_addr + + +def _get_default_hugepage_size() -> int: + try: + with open("/proc/meminfo", "r") as f: + for line in f: + if line.startswith("Hugepagesize:"): + parts = line.split() + if len(parts) >= 2: + kb = int(parts[1]) + return kb * 1024 + except Exception: + pass + return 2 * 1024 * 1024 + + +def _get_online_numa_nodes() -> List[int]: + for path in ("/sys/devices/system/node/has_memory", "/sys/devices/system/node/online"): + try: + with open(path, "r") as f: + online = f.read().strip() + nodes: List[int] = [] + for part in online.split(","): + if "-" in part: + start, end = part.split("-") + nodes.extend(range(int(start), int(end) + 1)) + else: + nodes.append(int(part)) + return nodes + except Exception: + continue + return [0] + + +def interleave_pages_across_numa_nodes(libc, addr: int, size: int) -> bool: + """为 CPU KV cache 的共享内存映射设置 NUMA 交错分配策略。 + + CPU KV cache 使用 SysV SHM 在多个进程间共享。默认的 first-touch 策略会把物理页分配到 + 首次触页线程所在的 NUMA 节点；在多 Socket 机器上，后台 prefault 线程的调度位置可能导致 + 大量 cache 页集中到单个内存控制器，限制多个 GPU 并发 load/offload 的主机内存带宽。 + + 本函数通过 ``mbind(MPOL_INTERLEAVE)`` 将映射范围内尚未分配的物理页按页偏移交错放置到 + 可用 NUMA 节点。调用方应在首次触页前设置策略：creator 在启动 prefault 线程前调用； + HugeTLB 的共享策略不会可靠地传播到其他进程的 VMA，因此 attacher 也需要在访问映射前调用。 + + 调用未设置 ``MPOL_MF_MOVE``，所以只影响后续缺页分配，不迁移已经分配的物理页。禁用开关、 + 单 NUMA、不支持的架构或 syscall 失败都会安全回退到原有 first-touch 行为。 + + Args: + libc: 使用 ``use_errno=True`` 加载的 libc 对象，用于发起 raw ``mbind`` syscall。 + addr: ``shmat`` 返回的、按页对齐的映射起始虚拟地址。 + size: 需要设置策略的映射长度；HugeTLB 模式下会向上对齐到默认大页大小。 + + Returns: + 策略成功安装时返回 ``True``；跳过或安装失败时返回 ``False``。 + """ + MPOL_INTERLEAVE = 3 + SYS_MBIND = {"x86_64": 237, "aarch64": 235}.get(os.uname().machine) + + if SYS_MBIND is None: + logger.warning(f"unsupported architecture {os.uname().machine}, skip cpu cache numa interleave") + return False + + if enable_huge_page(): + huge_sz = _get_default_hugepage_size() + size = triton.cdiv(size, huge_sz) * huge_sz + + def _mbind(mode, mask): + nodemask = ctypes.c_ulong(mask) + libc.syscall.restype = ctypes.c_long + return libc.syscall( + ctypes.c_long(SYS_MBIND), + ctypes.c_void_p(addr), + ctypes.c_ulong(size), + ctypes.c_int(mode), + ctypes.byref(nodemask), + # Raw syscall ABI decrements maxnode before copying the bitmap. + # Passing mask width + 1 preserves every bit while copying exactly one c_ulong. + ctypes.c_ulong(ctypes.sizeof(nodemask) * 8 + 1), + ctypes.c_uint(0), + ) + + if disable_cpu_cache_numa_interleave(): + logger.info("cpu cache numa interleave disabled by LIGHTLLM_DISABLE_NUMA_INTERLEAVE") + return False + nodes = _get_online_numa_nodes() + if len(nodes) <= 1: + return False + if max(nodes) >= 64: + logger.warning(f"more than 64 numa nodes ({nodes}), skip cpu cache numa interleave") + return False + try: + ret = _mbind(MPOL_INTERLEAVE, sum(1 << n for n in nodes)) + if ret != 0: + logger.warning( + f"mbind MPOL_INTERLEAVE failed (errno={ctypes.get_errno()}), " + f"cpu kv cache pages will use default first-touch numa policy" + ) + return False + logger.info(f"cpu kv cache pages interleaved across numa nodes {nodes}") + return True + except Exception as e: + logger.warning(f"cpu cache numa interleave skipped: {e}") + return False From 227e76defa947ed337261feb036096e9ee880998 Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Mon, 3 Aug 2026 03:17:01 +0000 Subject: [PATCH 3/3] fix: make CPU cache NUMA interleave opt-in --- lightllm/utils/envs_utils.py | 6 +++--- lightllm/utils/kv_cache_utils.py | 13 +++++++------ 2 files changed, 10 insertions(+), 9 deletions(-) diff --git a/lightllm/utils/envs_utils.py b/lightllm/utils/envs_utils.py index 8ee0fd982..33819d6f5 100644 --- a/lightllm/utils/envs_utils.py +++ b/lightllm/utils/envs_utils.py @@ -239,9 +239,9 @@ def enable_huge_page(): @lru_cache(maxsize=None) -def disable_cpu_cache_numa_interleave() -> bool: - """是否禁用 CPU KV cache 共享内存的 NUMA 交错分配策略。""" - return enable_env_vars("LIGHTLLM_DISABLE_NUMA_INTERLEAVE") +def enable_cpu_cache_numa_interleave() -> bool: + """是否启用 CPU KV cache 共享内存的 NUMA 交错分配策略。""" + return enable_env_vars("LIGHTLLM_ENABLE_NUMA_INTERLEAVE") @lru_cache(maxsize=None) diff --git a/lightllm/utils/kv_cache_utils.py b/lightllm/utils/kv_cache_utils.py index e30e093e4..e81caafe7 100644 --- a/lightllm/utils/kv_cache_utils.py +++ b/lightllm/utils/kv_cache_utils.py @@ -11,7 +11,7 @@ from lightllm.utils.envs_utils import ( get_env_start_args, enable_huge_page, - disable_cpu_cache_numa_interleave, + enable_cpu_cache_numa_interleave, get_llm_data_type, get_added_mtp_kv_layer_num, ) @@ -364,8 +364,9 @@ def interleave_pages_across_numa_nodes(libc, addr: int, size: int) -> bool: 可用 NUMA 节点。调用方应在首次触页前设置策略：creator 在启动 prefault 线程前调用； HugeTLB 的共享策略不会可靠地传播到其他进程的 VMA，因此 attacher 也需要在访问映射前调用。 - 调用未设置 ``MPOL_MF_MOVE``，所以只影响后续缺页分配，不迁移已经分配的物理页。禁用开关、 - 单 NUMA、不支持的架构或 syscall 失败都会安全回退到原有 first-touch 行为。 + 调用未设置 ``MPOL_MF_MOVE``，所以只影响后续缺页分配，不迁移已经分配的物理页。该功能默认 + 关闭，只有设置 ``LIGHTLLM_ENABLE_NUMA_INTERLEAVE`` 后才会启用；未启用、单 NUMA、不支持的 + 架构或 syscall 失败都会安全回退到原有 first-touch 行为。 Args: libc: 使用 ``use_errno=True`` 加载的 libc 对象，用于发起 raw ``mbind`` syscall。 @@ -378,6 +379,9 @@ def interleave_pages_across_numa_nodes(libc, addr: int, size: int) -> bool: MPOL_INTERLEAVE = 3 SYS_MBIND = {"x86_64": 237, "aarch64": 235}.get(os.uname().machine) + if not enable_cpu_cache_numa_interleave(): + return False + if SYS_MBIND is None: logger.warning(f"unsupported architecture {os.uname().machine}, skip cpu cache numa interleave") return False @@ -401,9 +405,6 @@ def _mbind(mode, mask): ctypes.c_uint(0), ) - if disable_cpu_cache_numa_interleave(): - logger.info("cpu cache numa interleave disabled by LIGHTLLM_DISABLE_NUMA_INTERLEAVE") - return False nodes = _get_online_numa_nodes() if len(nodes) <= 1: return False