T-MAC 是一種創(chuàng)新的基于查找表（LUT）的方法，專為在 CPU 上高效執(zhí)行低比特大型語言模型（LLMs）推理而設(shè)計，無需權(quán)重反量化，支持混合精度矩陣乘法（mpGEMM），顯著降低了推理開銷并提升了計算速度。
為增強設(shè)備上的智能性，在邊緣設(shè)備部署大型語言模型（LLMs）成為了一個趨勢，比如微軟的 Windows 11 AI + PC。
目前部署的大語言模型多會量化到低比特。然而，低比特 LLMs 在推理過程中需要進行低精度權(quán)重和高精度激活向量的混合精度矩陣乘法（mpGEMM）?，F(xiàn)有的系統(tǒng)由于硬件缺乏對 mpGEMM 的原生支持，不得不將權(quán)重反量化以進行高精度計算。這種間接的方式導致了顯著的推理開銷，并且無法隨著比特數(shù)進一步降低而獲得加速。
為此，微軟亞洲研究院、中國科學技術(shù)大學、中國科學院大學的研究人員聯(lián)合開發(fā)了 T-MAC。T-MAC 采用基于查找表（LUT）的計算范式，無需反量化，直接支持混合精度矩陣乘，其高效的推理性能以及其統(tǒng)一且可擴展的特性為在資源受限的邊緣設(shè)備上實際部署低比特 LLMs 鋪平了道路。
此外，當前大模型的部署普遍依賴于專用加速器，如 NPU 和 GPU 等，而 T-MAC 可以擺脫專用加速器的依賴，僅利用 CPU 部署 LLMs，推理速度甚至能夠超過同一片上的專用加速器，使 LLMs 可以部署在各類包括 PC、手機、樹莓派等邊緣端設(shè)備。T-MAC 現(xiàn)已開源。
在 CPU 上高效部署低比特大語言模型
T-MAC 的關(guān)鍵創(chuàng)新在于采用基于查找表（LUT）的計算范式，而非傳統(tǒng)的乘累加（MAC）計算范式。T-MAC 利用查找表直接支持低比特計算，從而消除了其他系統(tǒng)中必須的反量化 (dequantization) 操作，并且顯著減少了乘法和加法操作的數(shù)量。
經(jīng)過實驗，T-MAC 展現(xiàn)出了卓越的性能：在配備了最新高通 Snapdragon X Elite 芯片組的 Surface AI PC 上，3B BitNet-b1.58 模型的生成速率可達每秒 48 個 token，2bit 7B llama 模型的生成速率可達每秒 30 個 token，4bit 7B llama 模型的生成速率可達每秒 20 個 token。
這甚至超越了 NPU 的性能！
當部署 llama-2-7b-4bit 模型時，盡管使用 NPU 可以生成每秒 10.4 個 token，但 CPU 在 T-MAC 的助力下，僅使用兩核便能達到每秒 12.6 個 token，最高甚至可以飆升至每秒 22 個 token。
這些都遠超人類的平均閱讀速度，相比于原始的 llama.cpp 框架提升了 4 至 5 倍。
即使在較低端的設(shè)備如 Raspberry Pi 5 上，T-MAC 針對 3B BitNet-b1.58 也能達到每秒 11 個 token 的生成速率。T-MAC 也具有顯著的功耗優(yōu)勢：達到相同的生成速率，T-MAC 所需的核心數(shù)僅為原始 llama.cpp 的 1/4 至 1/6，降低能耗的同時也為其它應用留下計算資源。
值得注意的是，T-MAC 的計算性能會隨著比特數(shù)的降低而線性提高，這一現(xiàn)象在基于反量化去實現(xiàn)的 GPU 和 NPU 中是難以觀察到的。但 T-MAC 能夠在 2 比特下實現(xiàn)單核每秒 10 個 token，四核每秒 28 個 token，大大超越了 NPU 的性能。
矩陣乘不需乘，只需查表 (LUT)
對于低比特參數(shù) (weights)，T-MAC 將每一個比特單獨進行分組（例如，一組 4 個比特），這些比特與激活向量相乘，預先計算所有可能的部分和，然后使用 LUT 進行存儲。
之后，T-MAC 采用移位和累加操作來支持從 1 到 4 的可擴展位數(shù)。通過這種方法，T-MAC 拋棄了 CPU 上效率不高的 FMA（乘加）指令，轉(zhuǎn)而使用功耗更低效率也更高的 TBL / PSHUF（查表）指令。
以比特為核心的計算，取代以數(shù)據(jù)類型為核心的計算
傳統(tǒng)的基于反量化的計算，實際上是以數(shù)據(jù)類型為核心的計算，這種方式需要對每一種不同的數(shù)據(jù)類型單獨定制。
每種激活和權(quán)重的位寬組合，如 W4A16（權(quán)重 int4 激活 float16）和 W2A8，都需要特定的權(quán)重布局和計算內(nèi)核。
舉個例子，W3 的布局需要將 2 位和另外 1 位分開打包，并利用不同的交錯或混洗方法進行內(nèi)存對齊或快速解碼。然后，相應的計算內(nèi)核需要將這種特定布局解包到硬件支持的數(shù)據(jù)類型進行執(zhí)行。
而 T-MAC 通過從比特的視角觀察低比特矩陣乘計算，只需為單獨的一個比特設(shè)計最優(yōu)的數(shù)據(jù)結(jié)構(gòu)，然后通過堆疊的方式擴展到更高的 2/3/4 比特。
同時，對于不同精度的激活向量（float16 / float32 / int8），僅有構(gòu)建表的過程需要發(fā)生變化，在查表的時候不再需要考慮不同的數(shù)據(jù)結(jié)構(gòu)。
同時，傳統(tǒng)基于反量化的方法，從 4-比特降低到 3/2/1-比特時，盡管內(nèi)存占用更少，但是計算量并未減小，而且由于反量化的開銷不減反增，性能反而可能會更差。
但 T-MAC 的計算量隨著比特數(shù)降低能夠線性減少，從而在更低比特帶來更好加速，為最新的工作 BitNet，EfficientQAT 等發(fā)布的 1-比特 / 2-比特模型提供了高效率的部署方案。
高度優(yōu)化的算子實現(xiàn)
基于比特為核心的計算具有許多優(yōu)勢，但將其實現(xiàn)在 CPU 上仍具有不小的挑戰(zhàn)：
（1）與激活和權(quán)重的連續(xù)數(shù)據(jù)訪問相比，表的訪問是隨機的。表在快速片上內(nèi)存中的駐留對于最終的推理性能尤為重要；
（2）然而，片上內(nèi)存是有限的，查找表（LUT）方法相比傳統(tǒng)的 mpGEMV 增大了片上內(nèi)存的使用。這是因為查找表需要保存激活向量與所有可能的位模式相乘的結(jié)果。這比激活本身要多得多。
為此，微軟亞洲研究院的研究員們深入探究了基于查表的計算數(shù)據(jù)流，為這種計算范式設(shè)計了高效的數(shù)據(jù)結(jié)構(gòu)和計算流程，其中包括：
1. 將 LUT 存入片上內(nèi)存，以利用 CPU 上的查表向量指令（TBL / PSHUF）提升隨機訪存性能。
2. 改變矩陣 axis 計算順序，以盡可能提升放入片上內(nèi)存的有限 LUT 的數(shù)據(jù)重用率。
3. 為查表單獨設(shè)計最優(yōu)矩陣分塊（Tiling）方式，結(jié)合 autotvm 搜索最優(yōu)分塊參數(shù)
4. 參數(shù) weights 的布局優(yōu)化
a）weights 重排，以盡可能連續(xù)訪問并提升緩存命中率
b）weights 交錯，以提升解碼效率
5. 對 Intel / ARM CPU 做針對性優(yōu)化，包括
a）寄存器重排以快速建立查找表
b）通過取平均數(shù)指令做快速 8-比特累加
研究員們在一個基礎(chǔ)實現(xiàn)上，一步步應用各種優(yōu)化，最終相對于 SOTA 低比特算子獲得顯著加速：
廣告聲明：文內(nèi)含有的對外跳轉(zhuǎn)鏈接（包括不限于超鏈接、二維碼、口令等形式），用于傳遞更多信息，節(jié)省甄選時間，結(jié)果僅供參考，IT之家所有文章均包含本聲明。