作者：安謀科技 (Arm China) 主任軟件工程師 Bolt Liu

跨 NUMA 內存訪問可能會限制 llama.cpp 在 Arm Neoverse 平臺上的擴展能力。本文將為你詳細分析這一問題，并通過引入原型驗證補丁來加以解決。測試結果表明，在基于 Neoverse N2 平臺的系統(tǒng)上運行 llama3_Q4_0 模型時，該補丁可使文本生成性能提升多達 55%。

跨 NUMA 內存訪問問題

在 llama.cpp 中，當線程數超過 NUMA 節(jié)點中的核心數時，性能會下降。本示例基于每個 NUMA 節(jié)點包含 64 個核心的系統(tǒng)，并使用 llama3_Q4_0 模型進行測試。

根本原因分析

當線程跨越多個 NUMA 節(jié)點運行時，性能下降主要源于以下兩個原因：

ggml_barrier() 函數調用中的跨 NUMA 節(jié)點原子操作非常耗時。

MulMat 算子的張量緩沖區(qū)存在大量的跨 NUMA 內存訪問。

ggml_barrier 問題

在多線程環(huán)境下運行 llama.cpp 時，每個線程負責計算張量數據的一部分。在所有線程完成計算后，會使用 barrier 來確保數據同步。隨著線程數的增加，性能會下降；當線程數超過 NUMA 節(jié)點（每個節(jié)點 64 個核心）時，性能下降現(xiàn)象愈發(fā)明顯。

MulMat 算子問題

llama.cpp 中的 MulMat 算子是主要的性能瓶頸。理論上，增加線程數應該可以提高性能，但實際情況并非如此。原因在于張量緩沖區(qū)是通過 malloc() 分配的，而 malloc() 不具有 NUMA 感知能力，會導致大量的跨 NUMA 內存訪問。

優(yōu)化方法

為緩解上述的跨 NUMA 問題，本文采用了兩種優(yōu)化方法：ggml_barrier 優(yōu)化，以及 MulMat 算子優(yōu)化。

ggml_barrier 優(yōu)化

此優(yōu)化方案采用“分而治之”的思想對 ggml_barrier 進行優(yōu)化：

使用 NUMA 局部原子變量來實現(xiàn) barrier，速度很快。

僅由特定的 NUMA 節(jié)點中的最后一個完成張量計算的線程與其他 NUMA 節(jié)點的最后一個線程進行跨 NUMA 同步。

通過這種方式，執(zhí)行跨 NUMA 全局原子操作的線程數減少到與所涉及的 NUMA 節(jié)點數相當。

ggml_barrier 經優(yōu)化后，即使存在跨 NUMA 操作，性能也不會明顯下降。

MulMat 算子優(yōu)化

MulMat 算子在計算過程中，會使用三個張量緩沖區(qū)：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）。此優(yōu)化方案將緩沖區(qū)分割成多個片段，以便計算線程盡可能從本地 NUMA 節(jié)點訪問內存。

對于 dst 和 src0 緩沖區(qū)，每個線程根據線程 ID 訪問緩沖區(qū)的一部分。通過將緩沖區(qū)分割成 N 個片段（其中 N 為 NUMA 節(jié)點數），在滿足以下條件時，每個線程可以訪問本地 NUMA 節(jié)點中的緩沖區(qū)部分：

線程 ID 與物理核心 ID 有良好的對應關系，可通過線程親和性進行設置。

緩沖區(qū)片段被移動到所需的 NUMA 節(jié)點，這可通過 move_pages() 系統(tǒng)調用來實現(xiàn)。

MulMat 算子在計算期間，src1 被量化并存儲在另一個名為 wdata 的緩沖區(qū)中。MulMat 算子的計算公式如下：dst = src0 * wdata

當每個線程根據線程 ID 訪問 src0 和 dst 時，wdata 緩沖區(qū)需要被這些線程全部訪問。

鑒于量化過程并非 MulMat 算子的主要性能瓶頸，此優(yōu)化方案為每個 NUMA 節(jié)點分別創(chuàng)建一個局部的 wdata 緩沖區(qū)，各 NUMA 節(jié)點中的所有線程將 src1 量化至其自有的 wdata，因而總共有 N 個 wdata 副本（N 為 NUMA 節(jié)點數）。

優(yōu)化 MulMat 算子計算的張量數據布局后，可以看到性能有了明顯的提升。

總體性能比較

llama.cpp 的批處理基準測試結果如下：

NUMA 經優(yōu)化后，性能有了明顯的提升。

在 S_TG t/s 指標上，基準版本的最佳值為 26.52，線程數為 32；啟用 NUMA 優(yōu)化后，最佳值為 41.15，線程數為 40，性能提升 55%。

在 S t/s 指標上，基準版本的最佳值為 48.73，線程數為 36；啟用 NUMA 優(yōu)化后，最佳值為 74.67，線程數為 54，性能提升 53.2%。

本示例使用雙 NUMA 節(jié)點系統(tǒng)，NUMA 優(yōu)化前后的內存帶寬數據（如下表所示）表明：優(yōu)化后，各 NUMA 節(jié)點的帶寬保持均衡；而優(yōu)化前，瓶頸會出現(xiàn)在 NUMA 節(jié)點 0 中。

各位開發(fā)者，歡迎獲取 GitHub 上的 NUMA 優(yōu)化補丁：

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

* 本文為 Arm 原創(chuàng)文章，轉載請留言聯(lián)系獲得授權并注明出處。