🤖 AI Summary
This work addresses the challenge of efficiently deploying Mamba models on modern hardware, which is hindered by complex data dependencies leading to excessive off-chip memory accesses. To tackle this, the authors introduce a cascade-of-Einsums abstraction to systematically characterize Mamba’s computational structure and devise cross-Einsum fusion strategies that minimize data movement. Building upon this foundation, they design Mambalaya, a reconfigurable accelerator that pioneers the use of an extended Einsum framework to uniformly fuse multiple operators within Mamba. Experimental results demonstrate that Mambalaya achieves a 4.9× speedup over MARCA during the prefill phase and a 1.9× speedup in the generation phase; moreover, in prefill-dominated scenarios, it delivers 1.5× higher performance than existing fused accelerators.
📝 Abstract
Mamba is an emerging, complex workload with various short-range and long-range dependencies, nonlinearities, and elementwise computations that are unable to run at near-peak speeds on modern hardware. Specifically, Mamba's complex dependency graph makes fusion across its full operator cascade difficult, leaving substantial inter-operator memory traffic on the table. To address these challenges, we propose Mambalaya, a novel reconfigurable accelerator that leverages fusion to overcome the limitations of Mamba. We use the recently proposed cascade-of-Einsums abstraction to characterize Mamba's full computational structure, then apply the extended Einsum framework to systematically explore inter-Einsum fusion opportunities. This principled approach yields a series of fusion mappings that reduce off-chip inter-Einsum traffic. These mappings are supported by the underlying Mambalaya architecture. Mambalaya achieves a layer performance speedup of 4.9$\times$ for prefill and 1.9$\times$ for generation over MARCA. In prefill-dominated scenarios, it achieves up to 1.5$\times$ over a recent fine-grained, memory-aware fusion accelerator for Mamba.