為增強(qiáng)設(shè)備上的智能性，在邊緣設(shè)備部署大型語(yǔ)言模型（LLMs）成為了一個(gè)趨勢(shì)，比如微軟的Windows 11 AI + PC。

目前部署的大語(yǔ)言模型多會(huì)量化到低比特。然而，低比特LLMs在推理過(guò)程中需要進(jìn)行低精度權(quán)重和高精度激活向量的混合精度矩陣乘法（mpGEMM）?，F(xiàn)有的系統(tǒng)由于硬件缺乏對(duì)mpGEMM的原生支持，不得不將權(quán)重反量化以進(jìn)行高精度計(jì)算。這種間接的方式導(dǎo)致了顯著的推理開(kāi)銷(xiāo)，并且無(wú)法隨著比特?cái)?shù)進(jìn)一步降低而獲得加速。

為此，微軟亞洲研究院、中國(guó)科學(xué)技術(shù)大學(xué)、中國(guó)科學(xué)院大學(xué)的研究人員聯(lián)合開(kāi)發(fā)了T-MAC。T-MAC采用基于查找表（LUT）的計(jì)算范式，無(wú)需反量化，直接支持混合精度矩陣乘，其高效的推理性能以及其統(tǒng)一且可擴(kuò)展的特性為在資源受限的邊緣設(shè)備上實(shí)際部署低比特LLMs鋪平了道路。

代碼：https://github.com/microsoft/T-MAC

論文：https://www.arxiv.org/pdf/2407.00088

此外，當(dāng)前大模型的部署普遍依賴(lài)于專(zhuān)用加速器，如NPU和GPU等，而T-MAC可以擺脫專(zhuān)用加速器的依賴(lài)，僅利用CPU部署LLMs，推理速度甚至能夠超過(guò)同一片上的專(zhuān)用加速器，使LLMs可以部署在各類(lèi)包括PC、手機(jī)、樹(shù)莓派等邊緣端設(shè)備。T-MAC現(xiàn)已開(kāi)源。

在CPU上高效部署低比特大語(yǔ)言模型

T-MAC的關(guān)鍵創(chuàng)新在于采用基于查找表（LUT）的計(jì)算范式，而非傳統(tǒng)的乘累加（MAC）計(jì)算范式。T-MAC 利用查找表直接支持低比特計(jì)算，從而消除了其他系統(tǒng)中必須的反量化(dequantization)操作，并且顯著減少了乘法和加法操作的數(shù)量。經(jīng)過(guò)實(shí)驗(yàn)，T-MAC展現(xiàn)出了卓越的性能：在配備了最新高通Snapdragon X Elite芯片組的Surface AI PC 上，3B BitNet-b1.58模型的生成速率可達(dá)每秒48個(gè)token，2bit 7B llama模型的生成速率可達(dá)每秒30個(gè)token，4bit 7B llama模型的生成速率可達(dá)每秒20個(gè)token。這甚至超越了NPU的性能！當(dāng)部署llama-2-7b-4bit模型時(shí)，盡管使用NPU可以生成每秒10.4個(gè)token，但CPU在T-MAC的助力下，僅使用兩核便能達(dá)到每秒12.6個(gè)token，最高甚至可以飆升至每秒22個(gè)token。這些都遠(yuǎn)超人類(lèi)的平均閱讀速度，相比于原始的llama.cpp框架提升了4至5倍。

即使在較低端的設(shè)備如Raspberry Pi 5上，T-MAC針對(duì)3B BitNet-b1.58也能達(dá)到每秒11個(gè)token的生成速率。T-MAC也具有顯著的功耗優(yōu)勢(shì)：達(dá)到相同的生成速率，T-MAC所需的核心數(shù)僅為原始llama.cpp的1/4至1/6，降低能耗的同時(shí)也為其它應(yīng)用留下計(jì)算資源。

值得注意的是，T-MAC的計(jì)算性能會(huì)隨著比特?cái)?shù)的降低而線性提高，這一現(xiàn)象在基于反量化去實(shí)現(xiàn)的GPU和NPU中是難以觀察到的。但T-MAC能夠在2比特下實(shí)現(xiàn)單核每秒10個(gè)token，四核每秒28個(gè)token，大大超越了NPU的性能。

圖1 BitNet on T-MAC vs llama.cpp on Apple M2

圖2 在不同端側(cè)設(shè)備CPU（Surface Laptop 7, NVIDIA AGX Orin, Apple M2-Ultra）的各核數(shù)下T-MAC和llama.cpp的token生成速度可達(dá)llama.cpp的4-5倍。達(dá)到相同的生成速率，T-MAC所需的核心數(shù)僅為原始llama.cpp的1/4至1/6

矩陣乘不需乘，只需查表 (LUT)

對(duì)于低比特參數(shù) (weights)，T-MAC將每一個(gè)比特單獨(dú)進(jìn)行分組（例如，一組4個(gè)比特），這些比特與激活向量相乘，預(yù)先計(jì)算所有可能的部分和，然后使用LUT進(jìn)行存儲(chǔ)。

之后，T-MAC采用移位和累加操作來(lái)支持從1到4的可擴(kuò)展位數(shù)。通過(guò)這種方法，T-MAC拋棄了CPU上效率不高的FMA（乘加）指令，轉(zhuǎn)而使用功耗更低效率也更高的TBL/PSHUF（查表）指令。

圖3 混合精度GEMV基于現(xiàn)有反量化的實(shí)現(xiàn)范式vs T-MAC基于查找表的新范式

以比特為核心的計(jì)算，取代以數(shù)據(jù)類(lèi)型為核心的計(jì)算

傳統(tǒng)的基于反量化的計(jì)算，實(shí)際上是以數(shù)據(jù)類(lèi)型為核心的計(jì)算，這種方式需要對(duì)每一種不同的數(shù)據(jù)類(lèi)型單獨(dú)定制。

每種激活和權(quán)重的位寬組合，如W4A16（權(quán)重int4激活float16） 和W2A8，都需要特定的權(quán)重布局和計(jì)算內(nèi)核。

舉個(gè)例子，W3的布局需要將2位和另外1位分開(kāi)打包，并利用不同的交錯(cuò)或混洗方法進(jìn)行內(nèi)存對(duì)齊或快速解碼。然后，相應(yīng)的計(jì)算內(nèi)核需要將這種特定布局解包到硬件支持的數(shù)據(jù)類(lèi)型進(jìn)行執(zhí)行。

而T-MAC通過(guò)從比特的視角觀察低比特矩陣乘計(jì)算，只需為單獨(dú)的一個(gè)比特設(shè)計(jì)最優(yōu)的數(shù)據(jù)結(jié)構(gòu)，然后通過(guò)堆疊的方式擴(kuò)展到更高的2/3/4比特。

同時(shí)，對(duì)于不同精度的激活向量（float16/float32/int8），僅有構(gòu)建表的過(guò)程需要發(fā)生變化，在查表的時(shí)候不再需要考慮不同的數(shù)據(jù)結(jié)構(gòu)。

圖4 以比特為核心的查表計(jì)算混合精度GEMV

同時(shí)，傳統(tǒng)基于反量化的方法，從4-比特降低到3/2/1-比特時(shí)，盡管內(nèi)存占用更少，但是計(jì)算量并未減小，而且由于反量化的開(kāi)銷(xiāo)不減反增，性能反而可能會(huì)更差。

但T-MAC的計(jì)算量隨著比特?cái)?shù)降低能夠線性減少，從而在更低比特帶來(lái)更好加速，為最新的工作BitNet，EfficientQAT等發(fā)布的1-比特/2-比特模型提供了高效率的部署方案。

圖5 使用不同端側(cè)設(shè)備CPU的單核，T-MAC在4到1比特的混合精度GEMV算子相較llama.cpp加速3-11倍。T-MAC的GEMM耗時(shí)能隨著比特?cái)?shù)減少線性減少，而基于反量化的llama.cpp無(wú)法做到（1比特llama.cpp的算子性能由其2比特實(shí)現(xiàn)推算得到）

高度優(yōu)化的算子實(shí)現(xiàn)

基于比特為核心的計(jì)算具有許多優(yōu)勢(shì)，但將其實(shí)現(xiàn)在CPU上仍具有不小的挑戰(zhàn)：

（1）與激活和權(quán)重的連續(xù)數(shù)據(jù)訪問(wèn)相比，表的訪問(wèn)是隨機(jī)的。表在快速片上內(nèi)存中的駐留對(duì)于最終的推理性能尤為重要；

（2）然而，片上內(nèi)存是有限的，查找表（LUT）方法相比傳統(tǒng)的mpGEMV增大了片上內(nèi)存的使用。這是因?yàn)椴檎冶硇枰４婕せ钕蛄颗c所有可能的位模式相乘的結(jié)果。這比激活本身要多得多。

圖6 T-MAC與llama.cpp在計(jì)算數(shù)據(jù)流上的不同

為此，微軟亞洲研究院的研究員們深入探究了基于查表的計(jì)算數(shù)據(jù)流，為這種計(jì)算范式設(shè)計(jì)了高效的數(shù)據(jù)結(jié)構(gòu)和計(jì)算流程，其中包括：

1. 將LUT存入片上內(nèi)存，以利用CPU上的查表向量指令（TBL/PSHUF）提升隨機(jī)訪存性能。

2. 改變矩陣axis計(jì)算順序，以盡可能提升放入片上內(nèi)存的有限LUT的數(shù)據(jù)重用率。

3. 為查表單獨(dú)設(shè)計(jì)最優(yōu)矩陣分塊（Tiling）方式，結(jié)合autotvm搜索最優(yōu)分塊參數(shù)

4. 參數(shù)weights的布局優(yōu)化

a）weights重排，以盡可能連續(xù)訪問(wèn)并提升緩存命中率

b）weights交錯(cuò)，以提升解碼效率

5. 對(duì)Intel/ARM CPU 做針對(duì)性優(yōu)化，包括

a）寄存器重排以快速建立查找表

b）通過(guò)取平均數(shù)指令做快速8-比特累加

研究員們?cè)谝粋€(gè)基礎(chǔ)實(shí)現(xiàn)上，一步步應(yīng)用各種優(yōu)化，最終相對(duì)于SOTA低比特算子獲得顯著加速：

圖7：在實(shí)現(xiàn)各種優(yōu)化后，T-MAC 4-比特算子最終相對(duì)于llama.cpp 獲得顯著加速