Flex Attention: A Programming Model for Generating Optimized Attention Kernels
This addresses a bottleneck for researchers and practitioners in machine learning by simplifying the development and optimization of attention mechanisms, though it is incremental as it builds on existing attention optimization concepts.
The paper tackles the problem of efficiently implementing diverse attention variants in deep learning, which is hindered by the monolithic nature of optimized kernels like FlashAttention, by introducing FlexAttention, a compiler-driven programming model that enables implementing most attention variants in a few lines of PyTorch code and achieves competitive performance compared to handwritten kernels.
Over the past 7 years, attention has become one of the most important primitives in deep learning. The primary approach to optimize attention is FlashAttention, which fuses the operation together, drastically improving both the runtime and the memory consumption. However, the importance of FlashAttention combined with its monolithic nature poses a problem for researchers aiming to try new attention variants -- a "software lottery". This problem is exacerbated by the difficulty of writing efficient fused attention kernels, resisting traditional compiler-based approaches. We introduce FlexAttention, a novel compiler-driven programming model that allows implementing the majority of attention variants in a few lines of idiomatic PyTorch code. We demonstrate that many existing attention variants (e.g. Alibi, Document Masking, PagedAttention, etc.) can be implemented via FlexAttention, and that we achieve competitive performance compared to these handwritten kernels. Finally, we demonstrate how FlexAttention allows for easy composition of attention variants, solving the combinatorial explosion of attention variants.