diff --git a/lightllm/utils/envs_utils.py b/lightllm/utils/envs_utils.py index 4dca81ae6..33819d6f5 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 enable_cpu_cache_numa_interleave() -> bool: + """是否启用 CPU KV cache 共享内存的 NUMA 交错分配策略。""" + return enable_env_vars("LIGHTLLM_ENABLE_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 06c4211eb..e81caafe7 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, + enable_cpu_cache_numa_interleave, get_llm_data_type, get_added_mtp_kv_layer_num, ) @@ -181,20 +182,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 +221,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 +316,111 @@ 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 + + +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``，所以只影响后续缺页分配，不迁移已经分配的物理页。该功能默认 + 关闭，只有设置 ``LIGHTLLM_ENABLE_NUMA_INTERLEAVE`` 后才会启用；未启用、单 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 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 + + 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), + ) + + 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