Reproduce “FLAT: An Optimized Dataflow for Mitigating Attention Bottlenecks” on OpenResearch
Turn your favorite coding agents into research agents that review literature, develop hypotheses, and run experiments.
Universal app for Apple silicon and Intel. Requires macOS 11 or later.
About this paper
Sheng-Chun Kao, Suvinay Subramanian, Gaurav Agrawal, Amir Yazdanbakhsh, Tushar Krishna
An optimized dataflow named FLAT addresses memory and computational bottlenecks in Transformer attention mechanisms, enabling efficient processing of input sequences up to 64K tokens while reducing memory footprint growth from quadratic to linear, leading to performance and energy efficiency improvements on accelerators.
Key insights
- Inter-operator fusion of Logit, Softmax, and Attend is crucial for attention efficiency, requiring specialized handling due to Softmax's many-to-many dependencies.
- Tiling at a 'row-granularity' is key to managing the intermediate O(N²) tensor, allowing the active memory footprint to scale linearly with sequence length.
- Unlike activation-weight operations, finer tiling granularities are beneficial for activation-activation operations like attention, effectively controlling on-chip memory without compromising data reuse.
Summary by alphaXiv.
Tip: change arxiv.org to autoarxiv.org in any arXiv abstract link to open the paper here.