FuseMax: Leveraging Extended Einsums to Optimize Attention Accelerator Design

TL;DR

FuseMax introduces a novel approach using Einsum cascades to optimize transformer attention acceleration, achieving near-perfect compute utilization and significant speedups without off-chip memory bottlenecks.

Contribution

The paper formalizes attention algorithms with Einsum cascades and proposes FuseMax, a new architecture that improves speed and energy efficiency in transformer attention acceleration.

Findings

FuseMax achieves 6.7x speedup over FLAT in attention tasks.

FuseMax reduces energy consumption to 79% of prior methods.

End-to-end transformer inference is 5.3x faster with FuseMax.

Abstract

Attention for transformers is a critical workload that has recently received significant "attention" as a target for custom acceleration. Yet, while prior work succeeds in reducing attention's memory-bandwidth requirements, it creates load imbalance between operators that comprise the attention computation (resulting in severe compute under-utilization) and requires on-chip memory that scales with sequence length (which is expected to grow over time). This paper ameliorates these issues, enabling attention with nearly 100% compute utilization, no off-chip memory traffic bottlenecks, and on-chip buffer size requirements that are independent of sequence length. The main conceptual contribution is to use a recently proposed abstraction -- the cascade of Einsums -- to describe, formalize, and taxonomize the space of attention algorithms that appear in the literature. In particular, we show how Einsum cascades can be used to infer non-trivial lower bounds on the number of passes a kernel must take through its input data, which has implications for either required on-chip buffer capacity or memory traffic. We show how this notion can be used to meaningfully divide the space of attention algorithms into several categories and use these categories to inform our design process. Based on the above characterization, we propose FuseMax -- a novel mapping and binding of attention onto a spatial array-style architecture. On attention, in an iso-area comparison, FuseMax achieves an average 6.7x speedup over the prior state-of-the-art, FLAT, while using 79\% of the energy. Similarly, on full end-to-end transformer inference, FuseMax achieves an average 5.3x speedup over FLAT using 83 of the energy.