面向FPGA的Mamba高效加速:量化与硬件协设计
计算与语言
2025-10-13 v2
摘要
状态空间模型(SSM)如Mamba最近受到广泛关注。与基于Transformer的大型语言模型(LLM)相比,Mamba以序列长度为线性计算复杂度,并展现出卓越的性能。然而,Mamba因激活离群值分散且计算依赖复杂,导致现有LLM加速器效率不高。本文提出LightMamba,通过协同设计量化算法和FPGA加速器架构实现高效Mamba推理。我们首先提出一种面向FPGA的后训练量化算法,采用旋转辅助量化和幂OfTwo SSM量化,将大部分计算压缩至4位。进一步设计的FPGA加速器部分展开Mamba计算,以平衡效率与硬件成本。通过计算重排序、细粒度分块和融合,加速器的硬件利用率和内存效率得到显著提升。我们在Xilinx Versal VCK190 FPGA上实现了LightMamba,能耗效率比GPU基线提高4.65倍至6.06倍。在Alveo U280 FPGA上,LightMamba达到93个token/s,仅为GPU基线的1.43倍。我们的代码已公开于https://github.com/PKU-SEC-Lab/LightMamba。
引用
@article{arxiv.2502.15260,
title = {LightMamba: Efficient Mamba Acceleration on FPGA with Quantization and Hardware Co-design},
author = {Renjie Wei and Songqiang Xu and Linfeng Zhong and Zebin Yang and Qingyu Guo and Yuan Wang and Runsheng Wang and Meng Li},
journal= {arXiv preprint arXiv:2502.15260},
year = {2025}
}
备注
Accepted by DATE 2025