eMamba: Efficient Acceleration Framework for Mamba Models in Edge Computing

TL;DR

eMamba is a hardware acceleration framework that optimizes Mamba models for edge devices, achieving high efficiency and accuracy with significantly reduced resource consumption.

Contribution

This work introduces eMamba, the first end-to-end hardware acceleration framework specifically designed for Mamba models on edge platforms, including novel approximation and NAS techniques.

Findings

eMamba achieves 1.63-19.9× fewer parameters while maintaining accuracy.

It demonstrates 4.95-5.62× lower latency and 2.22-9.95× higher throughput on FPGA and ASIC.

eMamba generalizes well to large-scale natural language tasks with stable perplexity.

Abstract

State Space Model (SSM)-based machine learning architectures have recently gained significant attention for processing sequential data. Mamba, a recent sequence-to-sequence SSM, offers competitive accuracy with superior computational efficiency compared to state-of-the-art transformer models. While this advantage makes Mamba particularly promising for resource-constrained edge devices, no hardware acceleration frameworks are currently optimized for deploying it in such environments. This paper presents eMamba, a comprehensive end-to-end hardware acceleration framework explicitly designed for deploying Mamba models on edge platforms. eMamba maximizes computational efficiency by replacing complex normalization layers with lightweight hardware-aware alternatives and approximating expensive operations, such as SiLU activation and exponentiation, considering the target applications. Then, it performs an approximation-aware neural architecture search (NAS) to tune the learnable parameters used during approximation. Evaluations with Fashion-MNIST, CIFAR-10, and MARS, an open-source human pose estimation dataset, show eMamba achieves comparable accuracy to state-of-the-art techniques using 1.63-19.9$\times$ fewer parameters. In addition, it generalizes well to large-scale natural language tasks, demonstrating stable perplexity across varying sequence lengths on the WikiText2 dataset. We also quantize and implement the entire eMamba pipeline on an AMD ZCU102 FPGA and ASIC using GlobalFoundries (GF) 22 nm technology. Experimental results show 4.95-5.62$\times$ lower latency and 2.22-9.95$\times$ higher throughput, with 4.77$\times$ smaller area, 9.84$\times$ lower power, and 48.6$\times$ lower energy consumption than baseline solutions while maintaining competitive accuracy.