Hotdry.

Article

NUMA 架构下 AI 推理服务的内存分配策略工程实践

深入解析多插槽 NUMA 拓扑感知的 AI 推理内存分配机制,包括 GPU 带宽约束下的内存布局与动态预算调度实现路径。

2026-04-20systems

在当前大模型推理服务规模化部署的实际场景中,硬件层面的 NUMA(Non-Uniform Memory Access)架构对推理性能的影响已经成为不可忽视的工程变量。随着 AMD MI300X 等基于 chiplet 设计的 GPU 逐渐普及,传统的「均匀内存访问」假设正在被打破,推理服务需要重新审视内存分配的拓扑感知策略。

NUMA 拓扑感知的本质需求

现代双路甚至四路服务器上,每个 CPU 插槽拥有独立的内存控制器和本地内存空间,跨插槽访问将引入显著的额外延迟。在 AI 推理场景中,这一问题被进一步放大:Transformer 模型的多头注意力机制需要频繁访问 K、V 张量,而这些张量的物理位置直接决定了内存带宽利用率和推理尾延迟。

研究表明,NUMA 感知的内存分配策略可以带来接近 50% 的性能提升,同时维持 80% 至 97% 的 L2 缓存命中率。这一收益在长序列推理和大批量请求场景下尤为显著,因为跨 NUMA 节点的内存访问会形成严重的带宽争用,最终表现为推理延迟的尾部膨胀。

多插槽拓扑感知的实现路径

在工程实践中,实现 NUMA 感知的内存分配需要从以下几个维度入手:

拓扑建模是基础。部署团队应当首先通过 lstoponumactl --hardware 明确服务器的 NUMA 节点数量、每个节点的 CPU 核心分布、内存容量以及 PCIe 设备与节点的亲和性关系。对于配备多 GPU 的服务器,需要额外确认每块 GPU 卡属于哪个 NUMA 节点,避免出现 GPU 与其本地内存跨越不同节点的情况。

内存亲和性策略是核心环节。Linux 提供了 numactllibnuma 接口,可以显式控制进程或内存页面的分配节点。对于推理服务进程,建议使用 numactl --membind=0 绑定到特定 NUMA 节点,确保推理过程中分配的缓冲区始终位于本地内存。对于需要多节点协作的场景,可采用交叉内存策略(--interleave=all)但需评估其对延迟的影响。

CPU-GPU 亲和性绑定同样关键。GPU 计算结果回传 CPU 时,如果 PCIe 通道跨越了 NUMA 边界,将产生额外的传输开销。建议将推理进程同时绑定到与目标 GPU 处于同一 NUMA 节点的 CPU 核心上,确保数据传输路径最短。

GPU 带宽约束下的内存布局

Chiplet 架构的 GPU 引入了更细粒度的 NUMA 概念。以 AMD MI300X 为例,其多个 GPU 计算芯片(CCD)各自拥有独立的 HBM 内存控制器,芯片间的数据移动需要通过 Infinity Fabric 总线完成,带宽显著低于芯片内部访问。针对这一特性,业界提出了 Swizzled Head-first Mapping 等优化策略:把同一个注意力头的 K、V 张量映射到同一个 CCD 的本地内存中,最大化芯片内 locality,减少跨芯片的数据迁移。

对于推理服务开发者而言,这一优化通常需要在模型加载阶段显式指定张量的物理分配位置。PyTorch 生态中可以通过 torch.cuda.set_device() 结合自定义的内存分配器实现基本的拓扑感知。更细粒度的控制则需要借助供应商提供的底层 API(如 AMD 的 rocBLAS 或 NVIDIA 的 CUDA 内存池接口)进行定制化开发。

动态内存预算调度机制

静态的 NUMA 绑定无法应对推理服务中动态变化的负载模式。当并发请求数剧烈波动时,固定的内存分配策略可能导致某些节点内存耗尽而其他节点闲置。动态内存预算调度机制正是为了解决这一问题:

实时监控与阈值设定是第一步。建议在推理服务中集成 numa_membalance 指标监控,当单个 NUMA 节点的内存使用率超过 80% 时触发告警,超过 90% 时启动内存回收流程。监控粒度建议细化到每个推理批次(batch)的内存申请与释放周期。

基于权重的动态调度可以兼顾公平性与效率。为不同优先级或不同模型的推理任务分配不同的内存配额权重,高优先级任务在内存紧张时优先获得本地内存保障,低优先级任务可被动态迁移到其他 NUMA 节点或降级到磁盘交换空间。

冷热数据分层管理是另一个有效策略。将活跃的 KV 缓存保留在 NUMA 本地内存中,将不活跃的中间计算结果迁移到跨节点的统一存储或 CPU 内存。这一机制的实现需要推理框架支持张量的生命周期管理与显式迁移接口。

实践建议与参数阈值

综合业界经验,以下参数阈值可作为工程落地的参考起点:推理进程的 numactl --membind 绑定应作为默认配置;NUMA 节点内存使用率告警阈值设为 80%,熔断阈值设为 90%;对于批量大小动态可变的推理服务,建议预留 15% 的内存余量应对突发流量;GPU 与 CPU 的 PCIe 亲和性检查应纳入服务启动的健康检查流程。

需要强调的是,NUMA 感知优化的收益高度依赖于具体硬件拓扑和 workload 特征。部署团队应当基于实际压测结果调整策略,而非盲目套用固定配置。


参考资料

  • NUMA-Aware GPU Attention Optimization(arxiv.org)
  • Dynamic GPU Scheduling With Multi-Resource Awareness(CUHK)

systems