FlexAttention: The Flexibility of PyTorch with the Performance of FlashAttention

FlexAttention is a new PyTorch API designed to enhance the flexibility and performance of attention mechanisms in machine learning. Traditional attention implementations, while optimized for performance, often lack the flexibility needed for researchers to experiment with new attention variants without writing custom kernels. FlexAttention addresses this issue by allowing users to implement various attention variants using a few lines of PyTorch code. It compiles these implementations into a fused FlashAttention kernel, which maintains competitive performance without additional memory overhead. The API supports a range of attention variants, including Causal, Relative Positional Embeddings, and more, by enabling users to modify attention scores through a user-defined function. This flexibility allows for the exploration of complex combinations of attention mechanisms while leveraging existing PyTorch infrastructure. Additionally, FlexAttention can take advantage of sparsity in attention masks, leading to significant performance improvements. The API is expected to facilitate innovation in attention mechanisms, limited only by the user's creativity. Examples and applications of FlexAttention can be found in the Attention Gym repository.

- FlexAttention enhances flexibility in implementing various attention mechanisms in PyTorch.

- It compiles user-defined attention variants into efficient FlashAttention kernels.

- The API allows for modifications of attention scores and supports multiple attention variants.

- It leverages sparsity in attention masks for improved performance.

- FlexAttention aims to foster innovation in machine learning research.