LightMamba: Efficient Mamba Acceleration on FPGA with Quantization and Hardware Co-design

📄 arXiv: 2502.15260v2 📥 PDF

作者: Renjie Wei, Songqiang Xu, Linfeng Zhong, Zebin Yang, Qingyu Guo, Yuan Wang, Runsheng Wang, Meng Li

分类: cs.CL

发布日期: 2025-02-21 (更新: 2025-10-10)

备注: Accepted by DATE 2025

🔗 代码/项目: GITHUB

💡 一句话要点

LightMamba：通过量化和硬件协同设计，在FPGA上高效加速Mamba模型推理。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: Mamba模型 FPGA加速 量化 硬件协同设计 状态空间模型 后训练量化 低功耗计算

📋 核心要点

  1. Mamba等状态空间模型（SSM）虽然计算复杂度与序列长度呈线性关系，但其激活值异常分散和计算依赖复杂，导致现有LLM加速器效率低下。
  2. LightMamba通过协同设计量化算法和FPGA加速器架构，利用旋转辅助量化和二的幂次SSM量化，将计算量化到4比特，并优化硬件利用率和内存效率。
  3. 实验结果表明，LightMamba在FPGA上实现了显著的能效提升和吞吐量，优于GPU基线，验证了所提出方法的有效性。

📝 摘要（中文）

本文提出LightMamba，一种协同设计量化算法和FPGA加速器架构的方法，用于高效的Mamba模型推理。针对Mamba模型中激活值异常分散和计算依赖复杂的问题，LightMamba首先提出一种FPGA友好的后训练量化算法，该算法采用旋转辅助量化和二的幂次SSM量化，将大部分计算降低到4比特。其次，设计了一种FPGA加速器，通过部分展开Mamba计算来平衡效率和硬件成本。通过计算重排序以及细粒度的分块和融合，加速器的硬件利用率和内存效率得到显著提高。在Xilinx Versal VCK190 FPGA上的实验结果表明，LightMamba的能效比GPU基线高4.65倍至6.06倍。在Alveo U280 FPGA上评估时，LightMamba达到93 tokens/s，是GPU基线的1.43倍。代码已开源。

🔬 方法详解

问题定义：Mamba模型虽然在长序列建模上具有优势，但其激活值异常分散，计算依赖性复杂，导致在现有LLM加速器上难以高效加速。现有的加速器无法有效处理Mamba的这些特性，导致性能瓶颈。

核心思路：LightMamba的核心思路是通过量化和硬件协同设计来克服Mamba模型的加速挑战。量化算法旨在降低计算复杂度，同时保持模型精度；硬件加速器则针对量化后的模型进行优化，提高硬件利用率和内存效率。通过软硬件协同优化，实现Mamba模型在FPGA上的高效推理。

技术框架：LightMamba的技术框架主要包括两个部分：FPGA友好的后训练量化算法和FPGA加速器架构。量化算法首先采用旋转辅助量化来减少量化误差，然后使用二的幂次SSM量化将大部分计算降低到4比特。FPGA加速器则通过部分展开Mamba计算来平衡效率和硬件成本，并通过计算重排序以及细粒度的分块和融合来提高硬件利用率和内存效率。

关键创新：LightMamba的关键创新在于量化算法和硬件架构的协同设计。旋转辅助量化和二的幂次SSM量化能够有效地降低计算复杂度，同时保持模型精度。FPGA加速器通过计算重排序和细粒度的分块和融合，充分利用了FPGA的并行计算能力和片上存储资源。这种软硬件协同优化的方法与现有方法有本质区别，能够更好地适应Mamba模型的特性。

关键设计：在量化算法方面，旋转辅助量化的具体实现方式和二的幂次SSM量化的比例是关键设计参数。在FPGA加速器方面，部分展开的程度、计算重排序的具体策略、分块的大小和融合的方式等都是关键设计参数，需要根据具体的硬件平台和模型大小进行调整。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

LightMamba在Xilinx Versal VCK190 FPGA上实现了显著的能效提升，比GPU基线高4.65倍至6.06倍。在Alveo U280 FPGA上，LightMamba达到了93 tokens/s的吞吐量，是GPU基线的1.43倍。这些实验结果表明，LightMamba在FPGA上能够实现高效的Mamba模型推理。

🎯 应用场景

LightMamba在FPGA上的高效加速能力使其适用于资源受限的边缘设备和嵌入式系统，例如智能摄像头、机器人和自动驾驶等。通过降低计算复杂度和功耗，LightMamba可以支持在这些设备上部署更大规模的Mamba模型，从而提高其智能化水平和应用范围。

📄 摘要（原文）

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.