Skip to content

feat: interleave CPU KV cache pages across NUMA nodes#1399

Open
sufubao wants to merge 1 commit into
ModelTC:mainfrom
sufubao:cpu_cache_numa
Open

feat: interleave CPU KV cache pages across NUMA nodes#1399
sufubao wants to merge 1 commit into
ModelTC:mainfrom
sufubao:cpu_cache_numa

Conversation

@sufubao

@sufubao sufubao commented Jul 21, 2026

Copy link
Copy Markdown
Collaborator

CPU KV cache 的共享内存由 prefault 线程首次触碰,物理页落点取决于当时的线程调度。多 Socket 上连续 cache 分块可能集中到同一内存控制器,GPU 并发 load/offload 的带宽和运行间波动都会受影响。

这个改动在 prefault 前对 SHM VMA 设置 MPOL_INTERLEAVE,并对每个 attacher 的 VMA 设置相同策略。策略只影响后续缺页,不迁移已分配页。单 NUMA、设置 LIGHTLLM_DISABLE_NUMA_INTERLEAVE=1mbind 失败时继续使用 first-touch。

同一台 8×RTX 5090、双 NUMA 机器上,用普通 4K 页分别测了跨 NUMA 6 GPU(6 GiB)和同 NUMA GPU0-3(4 GiB),每组 3 轮 ABBA:

GPU 组合 聚合带宽 first-touch 中位数 interleave 中位数 变化
0,1,2,3,6,7(跨 NUMA) GPU→CPU offload 144.40 GB/s 170.26 GB/s +17.9%
0,1,2,3,6,7(跨 NUMA) CPU→GPU load 259.33 GB/s 325.44 GB/s +25.5%
0,1,2,3(均在 NUMA0) GPU→CPU offload 114.65 GB/s 143.28 GB/s +25.0%
0,1,2,3(均在 NUMA0) CPU→GPU load 163.22 GB/s 224.43 GB/s +37.5%

两组的三轮配对结果均为正。跨 NUMA 组的 first-touch offload/load CV 为 8.79%/9.54%,interleave 为 0.07%/0.20%;同 NUMA 组从 9.17%/13.98% 降到 0.09%/0.05%。interleave 的 numa_maps 每次都是精确 50/50,说明同 NUMA TP4 的收益也来自稳定页落点并同时使用两个内存控制器。

这是 pinned SHM 的 GPU DMA 微基准,不等同于完整模型 E2E,收益仍取决于拓扑和 workload。容器 seccomp 需要允许 mbind;被阻止时会记录 warning 并安全回退。py_compile、black 和 flake8 检查通过。

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request introduces NUMA node interleaving for CPU KV cache pages in lightllm/utils/kv_cache_utils.py by implementing interleave_pages_across_numa_nodes and helper functions. The feedback highlights two critical issues with the mbind syscall implementation: first, the hardcoded SYS_MBIND syscall number (237) is specific to x86_64 and will cause failures on aarch64 architectures; second, passing 65 as maxnode when nodemask is a 64-bit c_ulong leads to an out-of-bounds read of stack memory, potentially causing mbind to fail with EINVAL.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.


def interleave_pages_across_numa_nodes(libc, addr: int, size: int) -> bool:
MPOL_INTERLEAVE = 3
SYS_MBIND = 237

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

Hardcoding SYS_MBIND = 237 is specific to the x86_64 architecture. On aarch64 (ARM64), the syscall number for mbind is 235. Since ARM64 is a very common platform for LLM inference (e.g., NVIDIA Grace Hopper GH200), calling syscall 237 on ARM64 will execute rt_tgsigqueueinfo instead of mbind, leading to undefined behavior or crashes. We should dynamically determine the syscall number based on the system architecture.

Suggested change
SYS_MBIND = 237
import platform
arch = platform.machine()
if arch == "x86_64":
SYS_MBIND = 237
elif arch == "aarch64":
SYS_MBIND = 235
else:
logger.warning(f"Unsupported architecture {arch} for NUMA interleave")
return False

ctypes.c_ulong(size),
ctypes.c_int(mode),
ctypes.byref(nodemask),
ctypes.c_ulong(65),

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

Passing 65 as maxnode to mbind causes the kernel to read 128 bits (16 bytes) from userspace because the bit count is rounded up to the nearest multiple of unsigned long (64 bits). Since nodemask is defined as ctypes.c_ulong(mask), it only allocates 8 bytes (64 bits) on 64-bit systems. This leads to an out-of-bounds read of 8 bytes of uninitialized stack memory/garbage. If any of the garbage bits correspond to an offline or invalid node on the system, mbind will fail with EINVAL (errno 22), causing the NUMA interleave to be skipped entirely. Since max(nodes) >= 64 is already checked and skipped, all active nodes fit within the first 64 bits. Thus, maxnode should be set to 64 to prevent reading out-of-bounds and potential random failures.

Suggested change
ctypes.c_ulong(65),
ctypes.c_ulong(64),

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant