本文將為你介紹如何利用 Arm i8mm 指令，具體來說，是通過帶符號 8 位整數(shù)矩陣乘加指令 smmla，來優(yōu)化 llama.cpp 中 Q6_K 和 Q4_K 量化模型推理。

llama.cpp 量化

llama.cpp 是一個開源的 C++ 庫，用于運(yùn)行大語言模型 (LLM)，針對加速 CPU 推理進(jìn)行了優(yōu)化。通過量化等技術(shù)（例如 8 位或 4 位整數(shù)格式）來減少內(nèi)存占用并加快計(jì)算速度，從而實(shí)現(xiàn)在消費(fèi)級和服務(wù)器級硬件上高效部署模型。

llama.cpp 支持多種量化方式。量化可在模型精度和性能之間取得平衡。數(shù)據(jù)量越小，推理速度越快，但可能會因困惑度升高而致使精度降低。例如，Q8_0 采用 8 位整數(shù)表示一個數(shù)據(jù)點(diǎn)，而 Q6_K 則將數(shù)據(jù)量縮減至 6 位。

量化以塊為單位進(jìn)行，同一個塊中的數(shù)據(jù)點(diǎn)共享一個縮放因子。例如，Q8_0 的處理以 32 個數(shù)據(jù)點(diǎn)為一個塊，具體過程如下：

從原始數(shù)據(jù)中提取 32 個浮點(diǎn)值，記為 f[0:32]

計(jì)算絕對值的最大值，即 mf = max(abs(f[0:32]))

計(jì)算縮放因子：scale_factor = mf / (max(int8)) = mf / 127

量化：q[i] = round(f[i] / scale_factor)

反量化：v[i] = q[i] * scale_factor

Q6_K 則更為復(fù)雜。如下圖所示，數(shù)據(jù)點(diǎn)分為兩個層級：

一個超級塊包含 256 個數(shù)據(jù)點(diǎn)，并對應(yīng)一個浮點(diǎn)格式的超級塊縮放因子

每個超級塊由 16 個子塊組成。每個子塊包含 16 個數(shù)據(jù)點(diǎn)，這些數(shù)據(jù)點(diǎn)共享一個整數(shù)格式的子塊級縮放因子。

圖 1：Llama.cpp Q6_K 量化

利用 Arm i8mm 指令

優(yōu)化 llama.cpp

與大多數(shù)人工智能 (AI) 工作負(fù)載相同，在 LLM 推理過程中，大部分 CPU 周期都耗費(fèi)在矩陣乘法運(yùn)算上。Arm i8mm（具體是指 smmla 指令）能夠有效加速 8 位整數(shù)矩陣乘法運(yùn)算。

為了說明 smmla 指令的作用及其高效性，假設(shè)我們要對下圖中的兩個矩陣進(jìn)行乘法運(yùn)算。

圖 2：矩陣乘法

按照教科書上的方法，我們可以逐一計(jì)算輸出矩陣中的四個標(biāo)量，即第一個輸出標(biāo)量是矩陣 x 的第一行與矩陣 y 的第一列的內(nèi)積。依此類推，需要進(jìn)行四次內(nèi)積運(yùn)算。

還有一種更高效的方法，即外積法。如下圖所示，我們可以用矩陣 x 的第一列乘以矩陣 y 的第一行，一次性得出四個部分輸出標(biāo)量。將這兩個部分輸出相加就能得到結(jié)果，這樣只需要兩次外積運(yùn)算即可。

圖 3：外積

smmla 指令實(shí)現(xiàn)了向量級別的外積運(yùn)算，如下圖所示。請注意，vmmlaq_s32 是實(shí)現(xiàn) smmla 指令的編譯器內(nèi)建函數(shù)。

每個輸入向量 (int8x16) 被拆分為兩個 int8x8 向量

計(jì)算四對 int8x8 向量的內(nèi)積

將結(jié)果存儲到輸出向量 (int32x4) 的四個通道中

圖 4：smmla 指令

借助 smmla 指令，我們可以通過同時處理兩行和兩列來加速矩陣乘法。如下圖所示，計(jì)算步驟如下：

從矩陣 x 中加載兩行數(shù)據(jù) (int8x16) 到 vx0 和 vx1，從矩陣 y 中加載兩列數(shù)據(jù)到 vy0 和 vy1

對 vx0 和 vx1 進(jìn)行“壓縮”操作，將這兩個向量的下半部分合并為一個向量，上半部分合并為另一個向量。這是確保 smmla 指令正確工作的必要步驟。對 vy0 和 vy1 執(zhí)行相同操作

使用兩條 smmla 指令計(jì)算四個臨時標(biāo)量結(jié)果

處理下一個數(shù)據(jù)塊并累積臨時結(jié)果，直到處理完所有數(shù)據(jù)

圖 5：使用 smmla 指令進(jìn)行矩陣乘法

我們利用 smmla 指令對 llama.cpp 的 Q6_K 和 Q4_K 矩陣乘法內(nèi)核進(jìn)行了優(yōu)化，并在 Arm Neoverse N2 平臺上進(jìn)行了測試，觀察到性能有顯著提升。下圖展示了 Q6_K 優(yōu)化前后 llama.cpp 的性能對比，其中：

S_TG 代表詞元生成速度，數(shù)值越高代表性能越好

S_PP 代表提示詞預(yù)填充速度，數(shù)值越高代表性能越好

圖 6：Arm i8mm 提升 llama.cpp Q6_K 模型性能

上游補(bǔ)丁

[1]利用 Arm i8mm 優(yōu)化 llama.cpp Q6_K 內(nèi)核：

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

[2]利用 Arm i8mm 優(yōu)化 llama.cpp Q4_K 內(nèi)核：

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