# FlashAttention on a Napkin: A Diagrammatic Approach to Deep Learning IO-Awareness

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

## Key facts

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

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [High-Resolution Image Synthesis with Latent Diffusion Models](https://mlsanity.com/paper/2112.10752/) (2021)
- [SDXL: Improving Latent Diffusion Models for High-Resolution Image Synthesis](https://mlsanity.com/paper/2307.01952/) (2023)
- [Scaling Rectified Flow Transformers for High-Resolution Image Synthesis](https://mlsanity.com/paper/2403.03206/) (2024)
- [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)
- [A Survey of Quantization Methods for Efficient Neural Network Inference](https://mlsanity.com/paper/2103.13630/) (2021)
- [Fast Transformer Decoding: One Write-Head is All You Need](https://mlsanity.com/paper/1911.02150/) (2019)
- [FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision](https://mlsanity.com/paper/2407.08608/) (2024)
- [Data Movement Is All You Need: A Case Study on Optimizing Transformers](https://mlsanity.com/paper/2007.00072/) (2020)
- [AI and Memory Wall](https://mlsanity.com/paper/2403.14123/) (2024)
- [Backprop as Functor: A compositional perspective on supervised learning](https://mlsanity.com/paper/1711.10455/) (2017)
- [Categorical Foundations of Gradient-Based Learning](https://mlsanity.com/paper/2103.01931/) (2021)
- [Benchmarking and Dissecting the Nvidia Hopper GPU Architecture](https://mlsanity.com/paper/2402.13499/) (2024)
- [Stream-K: Work-centric Parallel Decomposition for Dense Matrix-Matrix Multiplication on the GPU](https://mlsanity.com/paper/2301.03598/) (2023)
- [An Introduction to String Diagrams for Computer Scientists](https://mlsanity.com/paper/2305.08768/) (2023)
- [Markov Categories and Entropy](https://mlsanity.com/paper/2212.11719/) (2022)
- [A Case Study in CUDA Kernel Fusion: Implementing FlashAttention-2 on NVIDIA Hopper Architecture using the CUTLASS Library](https://mlsanity.com/paper/2312.11918/) (2023)
- [Reducing shared memory footprint to leverage high throughput on Tensor Cores and its flexible API extension library](https://mlsanity.com/paper/2308.15152/) (2023)

## Built on by (most cited first, 7 of 7)

- [Agentic Deep Graph Reasoning Yields Self-Organizing Knowledge Networks](https://mlsanity.com/paper/2502.13025/) (2025)
- [In-situ graph reasoning and knowledge expansion using Graph-PReFLexOR](https://mlsanity.com/paper/2501.08120/) (2025)
- [Graph-Aware Isomorphic Attention for Adaptive Dynamics in Transformers](https://mlsanity.com/paper/2501.02393/) (2025)
- [AutoChemSchematic AI: Agentic Physics-Aware Automation for Chemical Manufacturing Scale-Up](https://mlsanity.com/paper/2505.24584/) (2025)
- [Accelerating Machine Learning Systems via Category Theory: Applications to Spherical Attention for Gene Regulatory Networks](https://mlsanity.com/paper/2505.09326/) (2025)
- [Weaves, Wires, and Morphisms: Formalizing and Implementing the Algebra of Deep Learning](https://mlsanity.com/paper/2604.07242/) (2026)
- [Graphical Design of Interpretable Architectures](https://mlsanity.com/paper/2608.18936/) (2026)

## Abstract (arXiv)

Optimizing deep learning algorithms currently requires slow, manual derivation, potentially leaving much performance untapped. Methods like FlashAttention have achieved a x6 performance improvement over native PyTorch by avoiding unnecessary data transfers, but required three iterations over three years to be developed. Automated compiled methods have consistently lagged behind. This paper extends Neural Circuit Diagrams for deep learning models to consider resource usage and the distribution of tasks across a GPU hierarchy. We show how diagrams can use simple relabellings to derive high-level streaming and tiling optimization strategies along with performance models. We show how this high-level performance model allows the effects of quantization and multi-level GPU hierarchies to be readily considered. We develop a methodology for representing intermediate-level pseudocode with diagrams, allowing hardware-aware algorithms to be derived step-by-step. Finally, we show how our methodology can be used to better understand existing techniques like FlashAttention. This work uses a theoretical framework to link assumptions about GPU behaviour to claims about performance. We aim to lay the groundwork for a scientific approach to GPU optimization where experiments can address clear hypotheses rather than post-hoc rationalizations.