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

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

跨 NUMA 內(nèi)存訪問問題

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

根本原因分析

當(dāng)線程跨越多個 NUMA 節(jié)點(diǎn)運(yùn)行時，性能下降主要源于以下兩個原因：

ggml_barrier() 函數(shù)調(diào)用中的跨 NUMA 節(jié)點(diǎn)原子操作非常耗時。

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

ggml_barrier 問題

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

MulMat 算子問題

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

優(yōu)化方法

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

ggml_barrier 優(yōu)化

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

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

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

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

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

MulMat 算子優(yōu)化

MulMat 算子在計(jì)算過程中，會使用三個張量緩沖區(qū)：dst（例如 attn_out、ffn_gate、ffn_out、ffn_up、Kcur、Qcur、Vcur、FP32）、src0（權(quán)重）和 src1（例如 attn_norm、ffn_gate_par、ffn_norm、kqv_out FP32）。此優(yōu)化方案將緩沖區(qū)分割成多個片段，以便計(jì)算線程盡可能從本地 NUMA 節(jié)點(diǎn)訪問內(nèi)存。

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

線程 ID 與物理核心 ID 有良好的對應(yīng)關(guān)系，可通過線程親和性進(jìn)行設(shè)置。

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

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

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

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

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

總體性能比較

llama.cpp 的批處理基準(zhǔn)測試結(jié)果如下：

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

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

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

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

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

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

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