LLM端侧新范式:T-MAC让CPU生成速度超NPU五倍

高效部署低比特大语言模型于CPU:优化策略与实践

在CPU上高效部署低比特大语言模型

T-MAC的关键创新在于采用基于查找表(LUT)的计算范式,而非传统的乘累加(MAC)计算范式。T-MAC利用查找表直接支持低比特计算,从而消除了其他系统中必须的反量化操作,并且显著减少了乘法和加法操作的数量。

经过实验,T-MAC展现出了卓越的性能:在配备了最新高通Snapdragon X E lite芯片组的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倍。

即使在较低端的设备如Raspberry Pi 5上,T-MAC针对3B BitNet-b1.58也能达到每秒11个token的生成速率。T-MAC也具有显著的功耗优势:达到相同的生成速率,T-MAC所需的核心数仅为原始llama.cpp的1/4至1/6,降低能耗的同时也为其它应用留下计算资源。

值得注意的是,T-MAC的计算性能会随着比特数的降低而线性提高,这一现象在基于反量化去实现的GPU和NPU中是难以观察到的。但T-MAC能够在2比特下实现单核每秒10个token,四核每秒28个token,大大超越了NPU的性能。

矩阵乘不需乘,只需查表 (LUT)

对于低比特参数(weights),T-MAC将每一个比特单独进行分组(例如,一组4个比特),这些比特与激活向量相乘,预先计算所有可能的部分和,然后使用LUT进行存储。

之后,T-MAC采用移位和累加操作来支持从1到4的可扩展位数。通过这种方法,T-MAC抛弃了CPU上效率不高的FMA(乘加)指令,转而使用功耗更低效率也更高的TBL/PSHUF(查表)指令。

以比特为核心的计算,取代以数据类型为核心的计算

传统的基于反量化的计算,实际上是以数据类型为核心的计算,这种方式需要对每一种不同的数据类型单独定制。

每种激活和权重的位宽组合,如W4A16(权重int4激活float16)和W2A8,都需要特定的权重布局和计算内核。

举个例子,W3的布局需要将2位和另外1位分开打包,并利用不同的交错或混洗方法进行内存对齐或快速解码。然后,相应的计算内核需要将这种特定布局解包到硬件支持的数据类型进行执行。

而T-MAC通过从比特的视角观察低比特矩阵乘计算,只需为单独的一个比特设计最优的数据结构,然后通过堆叠的方式扩展到更高的2/3/4比特。

同时,对于不同精度的激活向量(float16/float32/int8),仅有构建表的过程需要发生变化,在查表的时候不再需要考虑不同的数据结构。

同时,传统基于反量化的方法,从4-比特降低到3/2/1-比特时,尽管内存占用更少,但是计算量并未减小,而且由于反量化的开销不减反增,性能反而可能会更差。

但T-MAC的计算量随着比特数降低能够线性减少,从而在更低比特带来更好加速,为最新的工作BitNet,EfficientQAT等发布的1-比特/2-比特模型提供了高效率的部署方案。

高度优化的算子实现

基于比特为核心的计算具有许多优势,但将其实现在CPU上仍具有不小的挑战:

(1)与激活和权重的连续数据访问相比,表的访问是随机的。表在快速片上内存中的驻留对于最终的推理性能尤为重要;

(2)然而,片上内存是有限的,查找表(LUT)方法相比传统的mpGEMV增大了片上内存的使用。这是因为查找表需要保存激活向量与所有可能的位模式相乘的结果。这比激活本身要多得多。

为此,微软亚洲研究院的研究员们深入探究了基于查表的计算数据流,为这种计算范式设计了高效的数据结构和计算流程,其中包括:

  1. 将LUT存入片上内存,以利用CPU上的查表向量指令(TBL/PSHUF)提升随机访存性能。

  2. 改变矩阵axis计算顺序,以尽可能提升放入片上内存的有限LUT的数据重用率。

  3. 为查表单独设计最优矩阵分块(Tiling)方式,结合autotvm搜索最优分块参数

  4. 参数weights的布局优化

a)weights重排,以尽可能连续访问并提升缓存命中率

b)weights交错,以提升解码效率

  1. 对Intel/ARM CPU 做针对性优化,包括

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

b)通过取平均数指令做快速8-比特累加

研究员们在一个基础实现上,一步步应用各种优化,最终相对于SOTA低比特算子获得显著加速: