轻量化Mamba在FPGA上实现高效推理,能效比提升6倍。
LightMamba: Efficient Mamba Acceleration on FPGA with Quantization and Hardware Co-design
- 提出旋转辅助的4比特量化算法,简化计算
- 在VCK190上实现4.65-6.06倍能效提升
- 适合边缘端低功耗部署,如FPGA硬件加速
状态空间模型(SSMs)如Mamba近年来备受关注。与基于Transformer的大语言模型相比,Mamba具有与序列长度呈线性关系的计算复杂度,并表现出更优性能。然而,由于激活值分布零散且计算依赖复杂,现有大语言模型加速器难以有效加速Mamba。本文提出LightMamba,通过量化算法与FPGA加速架构的协同设计,实现高效的Mamba推理。我们首先提出一种面向FPGA的后训练量化算法,包含旋转辅助量化和2的幂次量化,将大部分计算降至4比特。进一步设计了部分展开的FPGA加速器,结合计算重排序、细粒度分块与融合,显著提升硬件利用率与内存效率。我们在Xilinx Versal VCK190 FPGA上实现LightMamba,能效比相较GPU基线提升4.65至6.06倍。在Alveo U280 FPGA上,达到93 tokens/s,是GPU基线的1.43倍。代码已公开于https://github.com/PKU-SEC-Lab/LightMamba。
原文摘要 · Abstract (English)
State space models (SSMs) like Mamba have recently attracted much attention. Compared to Transformer-based large language models (LLMs), Mamba achieves linear computation complexity with the sequence length and demonstrates superior performance. However, Mamba is hard to accelerate due to the scattered activation outliers and the complex computation dependency, rendering existing LLM accelerators inefficient. In this paper, we propose LightMamba that co-designs the quantization algorithm and FPGA accelerator architecture for efficient Mamba inference. We first propose an FPGA-friendly post-training quantization algorithm that features rotation-assisted quantization and power-of-two SSM quantization to reduce the majority of computation to 4-bit. We further design an FPGA accelerator that partially unrolls the Mamba computation to balance the efficiency and hardware costs. Through computation reordering as well as fine-grained tiling and fusion, the hardware utilization and memory efficiency of the accelerator get drastically improved. We implement LightMamba on Xilinx Versal VCK190 FPGA and achieve 4.65x to 6.06x higher energy efficiency over the GPU baseline. When evaluated on Alveo U280 FPGA, LightMamba reaches 93 tokens/s, which is 1.43x that of the GPU baseline. Our code is available at https://github.com/PKU-SEC-Lab/LightMamba.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。