如何在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() 系统调用来实现。

MulMat 算子在计算期间，src1 被量化并存储在另一个名为 wdata 的缓冲区中。MulMat 算子的计算公式如下：dst = src0 * wdata

当每个线程根据线程 ID 访问 src0 和 dst 时，wdata 缓冲区需要被这些线程全部访问。

鉴于量化过程并非 MulMat 算子的主要性能瓶颈，此优化方案为每个 NUMA 节点分别创建一个局部的 wdata 缓冲区，各 NUMA 节点中的所有线程将 src1 量化至其自有的 wdata，因而总共有 N 个 wdata 副本（N 为 NUMA 节点数）。

优化 MulMat 算子计算的张量数据布局后，可以看到性能有了明显的提升。

总体性能比较

llama.cpp 的批处理基准测试结果如下：

NUMA 经优化后，性能有了明显的提升。

在 S_TG t/s 指标上，基准版本的最佳值为 26.52，线程数为 32；启用 NUMA 优化后，最佳值为 41.15，线程数为 40，性能提升 55%。

在 S t/s 指标上，基准版本的最佳值为 48.73，线程数为 36；启用 NUMA 优化后，最佳值为 74.67，线程数为 54，性能提升 53.2%。

本示例使用双 NUMA 节点系统，NUMA 优化前后的内存带宽数据（如下表所示）表明：优化后，各 NUMA 节点的带宽保持均衡；而优化前，瓶颈会出现在 NUMA 节点 0 中。

各位开发者，欢迎获取 GitHub 上的 NUMA 优化补丁：

https://github.com/ggml-org/llama.cpp/pull/14232

* 本文为 Arm 原创文章，转载请留言联系获得授权并注明出处。