From TeAAL to FuseMax: Separation of Concerns for Attention Accelerator Design

Attention for transformers has recently received significant ‘attention’ as a target for custom acceleration. Our prior work, TeAAL [4], proposes a new accelerator design methodology that allows architects to reason about and optimize their designs iteratively. With a focus on attention, this work makes contributions to both the theory and practice of TeAAL’s methodology. On the theory side, we propose a set of analyses that can be applied using only the algorithm specification—the cascade of Einsums. On the practice side, we use the new analyses to analyze and taxonomize the space of attention algorithms and to iteratively build up an efficient, high-utilization accelerator. Our resulting design, FuseMax, achieves an average 6.7× speedup on attention and 5.3× speedup on end-to-end transformer inference over the prior state-of-the-art, FLAT [3], while using 79% and 83% of the energy, respectively.