mlsanityData to 14 Oct 2026

Flex Attention: A Programming Model for Generating Optimized Attention Kernels

arXiv 2412.05496 ↗First posted 7 Dec 2024cs.LG cs.PF cs.PLcited by 185

Not on the map. Why?

Cited by 46 in Diffusion models, but placed by its own abstract, which didn’t fit.

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.

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

13 of 22 references.

References · 13

  1. 2017Attention Is All You Needcited by 60,361
  2. 2019Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformercited by 14,599
  3. 2023Efficient Memory Management for Large Language Model Serving with PagedAttentioncited by 5,944
  4. 2022FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awarenesscited by 3,673
  5. 2020Longformer: The Long-Document Transformercited by 3,413
  6. 2023FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioningcited by 2,598
  7. 2023GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpointscited by 1,346
  8. 2021Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolationcited by 1,103
  9. 2024FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precisioncited by 519
  10. 2022Neighborhood Attention Transformercited by 239
  11. 2021Self-attention Does Not Need O(n2) Memorycited by 154
  12. 2024FlashMask: Efficient and Rich Mask Extension of FlashAttentioncited by 20
1 more
  1. 2024Faster Neighborhood Attention: Reducing the O(n^2) Cost of Self Attention at the Threadblock Levelcited by 14

Built on by

Cited by year (2026 to 14 Oct)

Citing papers per year: 2024: 1; 2025: 88; 2026: 96.

185 papers.

Sort:

By area:

  1. 2025Self Forcing: Bridging the Train-Test Gap in Autoregressive Video Diffusioncited by 625
  2. 2025Block Diffusion: Interpolating Between Autoregressive and Diffusion Language Modelscited by 409
  3. 2024Titans: Learning to Memorize at Test Timecited by 328
  4. 2025MAGI-1: Autoregressive Video Generation at Scalecited by 276
  5. 2025Kimi Linear: An Expressive, Efficient Attention Architecturecited by 196
  6. 2026DFlash: Block Diffusion for Flash Speculative Decodingcited by 135
  7. 2025Fast Video Generation with Sliding Tile Attentioncited by 128
  8. 2025Sparse VideoGen2: Accelerate Video Generation with Sparse Attention via Semantic-Aware Permutationcited by 106
  9. 2025VSA: Faster Video Diffusion with Trainable Sparse Attentioncited by 104
  10. 2025LaViDa: A Large Diffusion Language Model for Multimodal Understandingcited by 84
  11. 2025Long-Context State-Space Video World Modelscited by 64
  12. 2025VLASH: Real-Time VLAs via Future-State-Aware Asynchronous Inferencecited by 58
38 more
  1. 2025Edit Flows: Flow Matching with Edit Operationscited by 55
  2. 2025Any-Order Flexible Length Masked Diffusioncited by 51
  3. 2025LATTICE: Democratize High-Fidelity 3D Generation at Scalecited by 44
  4. 2025Relational Transformer: Toward Zero-Shot Foundation Models for Relational Datacited by 39
  5. 2026LoGeR: Long-Context Geometric Reconstruction with Hybrid Memorycited by 39
  6. 2025TiDAR: Think in Diffusion, Talk in Autoregressioncited by 37
  7. 2026Beyond Language Modeling: An Exploration of Multimodal Pretrainingcited by 37
  8. 2025PAN: A World Model for General, Actionable, and Long-Horizon World Simulationcited by 36
  9. 2025Training-Free Efficient Video Generation via Dynamic Token Carvingcited by 34
  10. 2025Esoteric Language Models: A Family of Any-Order Diffusion LLMscited by 34
  11. 2026LongLive-2.0: An NVFP4 Parallel Infrastructure for Long Video Generationcited by 30
  12. 2025Transformers are Graph Neural Networkscited by 29
  13. 2025Inference-Time Hyper-Scaling with KV Cache Compressioncited by 27
  14. 2025Eye, Robot: Learning to Look to Act with a BC-RL Perception-Action Loopcited by 25
  15. 2025ReasonEdit: Towards Reasoning-Enhanced Image Editing Modelscited by 25
  16. 2026Causal-rCM: A Unified Teacher-Forcing and Self-Forcing Open Recipe for Autoregressive Diffusion Distillation in Streaming Video Generation and Interactive World Modelscited by 25
  17. 2026DART: Diffusion-Inspired Speculative Decoding for Fast LLM Inferencecited by 24
  18. 2025Set Block Decoding is a Language Model Inference Acceleratorcited by 23
  19. 2025Macro-from-Micro Planning for High-Quality and Parallelized Autoregressive Long Video Generationcited by 21
  20. 2025FPSAttention: Training-Aware FP8 and Sparsity Co-Design for Fast Video Diffusioncited by 19
  21. 2025Sequential Diffusion Language Modelscited by 19
  22. 2025Cache What Lasts: Token Retention for Memory-Bounded KV Cache in LLMscited by 19
  23. 2025From Pixels to Words -- Towards Native Vision-Language Primitives at Scalecited by 18
  24. 2025Adaptive Self-improvement LLM Agentic System for ML Library Developmentcited by 16
  25. 2025BAgger: Backwards Aggregation for Mitigating Drift in Autoregressive Video Diffusion Modelscited by 16
  26. 2026HALO: A Unified Vision-Language-Action Model for Embodied Multimodal Chain-of-Thought Reasoningcited by 16
  27. 2026Thinking in Streaming Videocited by 16
  28. 2025The Automated LLM Speedrunning Benchmark: Reproducing NanoGPT Improvementscited by 15
  29. 20254Real-Video-V2: Fused View-Time Attention and Feedforward Reconstruction for 4D Scene Generationcited by 14
  30. 2025Deterministic Inference across Tensor Parallel Sizes That Eliminates Training-Inference Mismatchcited by 14
  31. 2026NVIDIA OmniDreams: Real-Time Generative World Model for Closed-Loop Autonomous Vehicle Simulationcited by 14
  32. 2026HunyuanOCR-1.5: Making Lightweight OCR VLMs Faster and Bettercited by 14
  33. 2025Towards Mechanistic Interpretability of Graph Transformers via Attention Graphscited by 13
  34. 2025Grouping First, Attending Smartly: Training-Free Acceleration for Diffusion Transformerscited by 13
  35. 2026ProphetKV: User-Query-Driven Selective Recomputation for Efficient KV Cache Reuse in Retrieval-Augmented Generationcited by 13
  36. 2025AdaSplash: Adaptive Sparse Flash Attentioncited by 12
  37. 2025Efficient Many-Shot In-Context Learning with Dynamic Block-Sparse Attentioncited by 12
  38. 2026Bernini: Latent Semantic Planning for Video Diffusioncited by 12

50 most cited of 185 All 185 (JSON)

Similar, not cited

Close in wording, no citation either way

  1. 2025Mixture of Sparse Attention: Content-Based Learnable Sparse Attention via Expert-Choice Routing
  2. 2025Attentions Under the Microscope: A Comparative Study of Resource Utilization for Variants of Self-Attention
  3. 2024SeerAttention: Learning Intrinsic Sparse Attention in Your LLMs
  4. 2026DualKV: Shared-Prompt Flash Attention for Efficient RL Training with Large Rollouts and Long Contexts
  5. 2025Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention
  6. 2025FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving
  7. 2025BlindSight: Harnessing Sparsity for Efficient Vision-Language Models
  8. 2025SeerAttention-R: Sparse Attention Adaptation for Long Reasoning