如何在Arm Neoverse N2平台上提升llama.cpp扩展性能-电子发烧友网

作者：安谋科技 (Arm China) 主任软件工程师 Bolt Liu

跨 NUMA 内存访问可能会限制 llama.cpp 在 Arm Neoverse 平台上的扩展能力。本文将为你详细分析这一问题，并通过引入原型验证补丁来加以解决。测试结果表明，在基于 Neoverse N2 平台的系统上运行 llama3_Q4_0 模型时，该补丁可使文本生成性能提升多达 55%。

跨 NUMA 内存访问问题

在 llama.cpp 中，当线程数超过 NUMA 节点中的核心数时，性能会下降。本示例基于每个 NUMA 节点包含 64 个核心的系统，并使用 llama3_Q4_0 模型进行测试。

根本原因分析

当线程跨越多个 NUMA 节点运行时，性能下降主要源于以下两个原因：

ggml_barrier() 函数调用中的跨 NUMA 节点原子操作非常耗时。

MulMat 算子的张量缓冲区存在大量的跨 NUMA 内存访问。

ggml_barrier 问题

在多线程环境下运行 llama.cpp 时，每个线程负责计算张量数据的一部分。在所有线程完成计算后，会使用 barrier 来确保数据同步。随着线程数的增加，性能会下降；当线程数超过 NUMA 节点（每个节点 64 个核心）时，性能下降现象愈发明显。

MulMat 算子问题

llama.cpp 中的 MulMat 算子是主要的性能瓶颈。理论上，增加线程数应该可以提高性能，但实际情况并非如此。原因在于张量缓冲区是通过 malloc() 分配的，而 malloc() 不具有 NUMA 感知能力，会导致大量的跨 NUMA 内存访问。

优化方法

为缓解上述的跨 NUMA 问题，本文采用了两种优化方法：ggml_barrier 优化，以及 MulMat 算子优化。

ggml_barrier 优化

此优化方案采用“分而治之”的思想对 ggml_barrier 进行优化：

使用 NUMA 局部原子变量来实现 barrier，速度很快。

仅由特定的 NUMA 节点中的最后一个完成张量计算的线程与其他 NUMA 节点的最后一个线程进行跨 NUMA 同步。

通过这种方式，执行跨 NUMA 全局原子操作的线程数减少到与所涉及的 NUMA 节点数相当。

ggml_barrier 经优化后，即使存在跨 NUMA 操作，性能也不会明显下降。

MulMat 算子优化

MulMat 算子在计算过程中，会使用三个张量缓冲区：dst（例如 attn_out、ffn_gate、ffn_out、ffn_up、Kcur、Qcur、Vcur、FP32）、src0（权重）和 src1（例如 attn_norm、ffn_gate_par、ffn_norm、kqv_out FP32）。此优化方案将缓冲区分割成多个片段，以便计算线程尽可能从本地 NUMA 节点访问内存。

对于 dst 和 src0 缓冲区，每个线程根据线程 ID 访问缓冲区的一部分。通过将缓冲区分割成 N 个片段（其中 N 为 NUMA 节点数），在满足以下条件时，每个线程可以访问本地 NUMA 节点中的缓冲区部分：

线程 ID 与物理核心 ID 有良好的对应关系，可通过线程亲和性进行设置。

缓冲区片段被移动到所需的 NUMA 节点，这可通过 move_pages() 系统调用来实现。