MARCA: Mamba Accelerator with Reconfigurable Architecture
Jinhao Li, Shan Huang, Jianrong Xu, J Liu, Li Ding, Ningyi Xu, Guohao Dai
Abstract
State space model (SSM) especially Mamba has demonstrated remarkable capabilities in various domains. Compared to Transformers, Mamba reduces the quadratic computational complexity and achieves a higher algorithm accuracy (e.g., the accuracy of Mamba-2.8b is higher than OPT-6.7b). However, challenges still exist in accelerating Mamba computations. (1) Incompatibility between element-wise operations and Tensor Core. Linear operations (matrix multiplications) and element-wise operations are the two dominating operations in Mamba. The time proportion of element-wise operations escalates significantly (e.g., >60% with 2048 input length). These operations do not need reduction, which is not compatible with the existing Tensor Core-based architectures (e.g., 1/16 normalized speed). (2) Large area overhead for nonlinear function unit. The optimized nonlinear function unit like exponential unit still occupies >30% of the processing element (PE) area. (3) Large memory access but limited data sharing for element-wise operations. Linear and element-wise operations in Mamba exhibit large compute intensity variance (e.g., ~3 orders of magnitude) and large read/write ratio variance (e.g., >3 orders). Due to the limited data sharing in element-wise operations, it is useless to apply the existed methods like tiling to element-wise operations.