# Flex Attention: A Programming Model for Generating Optimized Attention Kernels

arXiv: https://arxiv.org/abs/2412.05496 · HTML: https://mlsanity.com/paper/2412.05496/

## Key facts

- Cited by 185 papers in the map (data through 14 Oct 2026).
- Builds on 13 of its 22 references that are in the map.
- First posted on arXiv 7 Dec 2024.

## Builds on (most cited first, 13 of 13)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [Efficient Memory Management for Large Language Model Serving with PagedAttention](https://mlsanity.com/paper/2309.06180/) (2023)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [Longformer: The Long-Document Transformer](https://mlsanity.com/paper/2004.05150/) (2020)
- [FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning](https://mlsanity.com/paper/2307.08691/) (2023)
- [GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints](https://mlsanity.com/paper/2305.13245/) (2023)
- [Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation](https://mlsanity.com/paper/2108.12409/) (2021)
- [FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision](https://mlsanity.com/paper/2407.08608/) (2024)
- [Neighborhood Attention Transformer](https://mlsanity.com/paper/2204.07143/) (2022)
- [Self-attention Does Not Need O(n2) Memory](https://mlsanity.com/paper/2112.05682/) (2021)
- [FlashMask: Efficient and Rich Mask Extension of FlashAttention](https://mlsanity.com/paper/2410.01359/) (2024)
- [Faster Neighborhood Attention: Reducing the O(n^2) Cost of Self Attention at the Threadblock Level](https://mlsanity.com/paper/2403.04690/) (2024)

## Built on by (most cited first, 30 of 185)

- [Self Forcing: Bridging the Train-Test Gap in Autoregressive Video Diffusion](https://mlsanity.com/paper/2506.08009/) (2025)
- [Block Diffusion: Interpolating Between Autoregressive and Diffusion Language Models](https://mlsanity.com/paper/2503.09573/) (2025)
- [Titans: Learning to Memorize at Test Time](https://mlsanity.com/paper/2501.00663/) (2024)
- [MAGI-1: Autoregressive Video Generation at Scale](https://mlsanity.com/paper/2505.13211/) (2025)
- [Kimi Linear: An Expressive, Efficient Attention Architecture](https://mlsanity.com/paper/2510.26692/) (2025)
- [DFlash: Block Diffusion for Flash Speculative Decoding](https://mlsanity.com/paper/2602.06036/) (2026)
- [Fast Video Generation with Sliding Tile Attention](https://mlsanity.com/paper/2502.04507/) (2025)
- [Sparse VideoGen2: Accelerate Video Generation with Sparse Attention via Semantic-Aware Permutation](https://mlsanity.com/paper/2505.18875/) (2025)
- [VSA: Faster Video Diffusion with Trainable Sparse Attention](https://mlsanity.com/paper/2505.13389/) (2025)
- [LaViDa: A Large Diffusion Language Model for Multimodal Understanding](https://mlsanity.com/paper/2505.16839/) (2025)
- [Long-Context State-Space Video World Models](https://mlsanity.com/paper/2505.20171/) (2025)
- [VLASH: Real-Time VLAs via Future-State-Aware Asynchronous Inference](https://mlsanity.com/paper/2512.01031/) (2025)
- [Edit Flows: Flow Matching with Edit Operations](https://mlsanity.com/paper/2506.09018/) (2025)
- [Any-Order Flexible Length Masked Diffusion](https://mlsanity.com/paper/2509.01025/) (2025)
- [LATTICE: Democratize High-Fidelity 3D Generation at Scale](https://mlsanity.com/paper/2512.03052/) (2025)
- [Relational Transformer: Toward Zero-Shot Foundation Models for Relational Data](https://mlsanity.com/paper/2510.06377/) (2025)
- [LoGeR: Long-Context Geometric Reconstruction with Hybrid Memory](https://mlsanity.com/paper/2603.03269/) (2026)
- [TiDAR: Think in Diffusion, Talk in Autoregression](https://mlsanity.com/paper/2511.08923/) (2025)
- [Beyond Language Modeling: An Exploration of Multimodal Pretraining](https://mlsanity.com/paper/2603.03276/) (2026)
- [PAN: A World Model for General, Actionable, and Long-Horizon World Simulation](https://mlsanity.com/paper/2511.09057/) (2025)
- [Training-Free Efficient Video Generation via Dynamic Token Carving](https://mlsanity.com/paper/2505.16864/) (2025)
- [Esoteric Language Models: A Family of Any-Order Diffusion LLMs](https://mlsanity.com/paper/2506.01928/) (2025)
- [LongLive-2.0: An NVFP4 Parallel Infrastructure for Long Video Generation](https://mlsanity.com/paper/2605.18739/) (2026)
- [Transformers are Graph Neural Networks](https://mlsanity.com/paper/2506.22084/) (2025)
- [Inference-Time Hyper-Scaling with KV Cache Compression](https://mlsanity.com/paper/2506.05345/) (2025)
- [Eye, Robot: Learning to Look to Act with a BC-RL Perception-Action Loop](https://mlsanity.com/paper/2506.10968/) (2025)
- [ReasonEdit: Towards Reasoning-Enhanced Image Editing Models](https://mlsanity.com/paper/2511.22625/) (2025)
- [Causal-rCM: A Unified Teacher-Forcing and Self-Forcing Open Recipe for Autoregressive Diffusion Distillation in Streaming Video Generation and Interactive World Models](https://mlsanity.com/paper/2606.25473/) (2026)
- [DART: Diffusion-Inspired Speculative Decoding for Fast LLM Inference](https://mlsanity.com/paper/2601.19278/) (2026)
- [Set Block Decoding is a Language Model Inference Accelerator](https://mlsanity.com/paper/2509.04185/) (2025)

## Abstract (arXiv)

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.