Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

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

TL;DR

This work tackles accelerating Mamba on FPGA by co-designing a rotation-assisted post-training quantization and an FPGA accelerator. It introduces a PoT-based, FPGA-friendly quantization for the SSM, enabling low-bit computations, and a three-module accelerator (MMU, SSMU, HTU) with computation reordering and fine-grained tiling/fusion. The approach yields substantial gains, achieving up to 93 tokens/s on Alveo U280 and 7.21 tokens/s on Versal VCK190, with energy efficiency improvements of up to 6.06x over a RTX 2070 GPU baseline and 4.65x over a 4090. These results demonstrate practical, end-to-end FPGA acceleration for Mamba, enabling efficient autoregressive generation with reduced memory and compute costs. The methods are implemented in LightMamba and publicly available at the project repository.

Abstract

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.

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

TL;DR

This work tackles accelerating Mamba on FPGA by co-designing a rotation-assisted post-training quantization and an FPGA accelerator. It introduces a PoT-based, FPGA-friendly quantization for the SSM, enabling low-bit computations, and a three-module accelerator (MMU, SSMU, HTU) with computation reordering and fine-grained tiling/fusion. The approach yields substantial gains, achieving up to 93 tokens/s on Alveo U280 and 7.21 tokens/s on Versal VCK190, with energy efficiency improvements of up to 6.06x over a RTX 2070 GPU baseline and 4.65x over a 4090. These results demonstrate practical, end-to-end FPGA acceleration for Mamba, enabling efficient autoregressive generation with reduced memory and compute costs. The methods are implemented in LightMamba and publicly available at the project repository.

Abstract

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.

Paper Structure

This paper contains 17 sections, 1 equation, 10 figures, 4 tables.

Table of Contents

  1. Introduction
  2. related works
  3. LLM Quantization
  4. FPGA-based LLM Accelerators
  5. Motivation
  6. FPGA-friendly Quantization Algorithm
  7. Rotation-assisted Linear Layer Quantization
  8. FPGA-friendly SSM Quantization
  9. FPGA-based Mamba Acceleration
  10. Overall Hardware and Hadamard Transform Unit
  11. Computation Reordering
  12. Fine-grained Tiling and Fusion
  13. experiments
  14. Experiment Setup
  15. Evaluation Result
  16. ...and 2 more sections

Figures (10)

  • Figure 1: The model architecture of Mamba2 and the detailed computation graph of the SSM layer.
  • Figure 2: Activation distribution in Mamba2-2.7B before and after rotation.
  • Figure 3: The hardware cost of different operations in the SSM layer with naive Non-PoT quantization and PoT quantization.
  • Figure 4: (a) The proposed rotation-assisted quantization algorithm. Both $Q$ and $H$ are Hadamard matrices to ensure computation correctness. (b) Quantization error of the output projection weight after only rotation or fusion and rotation.
  • Figure 5: Diagram of (a) the overall architecture, (b) SSMU, (c) MMU, (d) 128-point HTU, and (e) 40-point HTU.
  • ...and 5 more figures