FlashAttention
In plain wordsA faster way for AI chips to compute attention, the step in which a language model relates every word to every other. It gets exactly the same answer while moving far less data around the chip.
A family of GPU kernels that compute exact scaled dot-product attentionscaled dot-product attention, n.: The standard formula for attention in transformers: score how well each question matches each label, scale the scores down, turn them into percentages, and use those to blend the stored information. without ever writing the full n × n matrix of attention scores to the GPU's large but slow high-bandwidth memoryhigh-bandwidth memory, n.: Very fast memory made by stacking memory chips on top of each other and placing them right next to the processor, so AI chips can be fed data quickly enough.; its authors called the design 'IO-aware' because it is organized around minimizing memory traffic rather than arithmetic. It splits queries, keys and values into tiles that fit in fast on-chip SRAMSRAM, abbr.: The fastest kind of computer memory, built directly into processor chips. It is quick but takes up a lot of space, so chips hold only a little of it., computes each tile's scores, softmaxsoftmax, n.: A formula that turns a list of raw scores into probabilities — all positive and adding up to one — with the biggest share going to the highest score. and weighted values in one fused kernel (kernel fusionkernel fusion, n.: Merging several small calculation steps into one, so a chip finishes them in a single pass instead of saving and reloading data in between, like cooking at the counter without returning ingredients to the fridge after each step.), and combines tiles with an online softmax that keeps running maxima and sums; the backward pass recomputes scores instead of storing them. Memory use becomes linear rather than quadratic in sequence length, and wall-clock speed improves several-fold because attention is otherwise dominated by memory traffic. FlashAttention-2 (2023) improved parallelism and work partitioning, reaching 50–73% of an A100's peak; FlashAttention-3 (2024) exploited asynchrony and FP8FP8, abbr.: An 8-bit number format that uses a quarter of the memory of the standard 32-bit format, letting new AI chips train and run models roughly twice as fast as with 16 bits. on H100H100, prop. n.: Nvidia's flagship AI chip of 2022–2024, so scarce and sought-after during the generative-AI boom that it became a unit of measure for AI computing power. GPUs, reaching up to 740 TFLOP/s in FP16. The approach underlies attention in most training and inference stacks and was a key enabler of long contextlong context, n.: The ability of an AI model to work with very long inputs — whole books, codebases or hours of transcripts — in one go, and the techniques that make this possible..
Swapping in FlashAttention let us train at 32k context without running out of memory.
Origin
Introduced by Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra and Christopher Ré of Stanford University and the University at Buffalo (arXiv, May 2022; NeurIPS 2022), building on the online-softmax trick (Maxim Milakov and Natalia Gimelshein, 2018) and memory-efficient attention (Markus Rabe and Charles Staats, 2021). Dao published FlashAttention-2 alone in July 2023; FlashAttention-3 followed in July 2024 with collaborators from Colfax, Meta and NVIDIA.
Usage note
Computes the same result as standard attention, up to floating-point rounding; it is not an approximation like sparse attentionsparse attention, n.: Attention in which each word looks at only a selected subset of the other words instead of all of them, so very long texts can be processed at a fraction of the cost. or linear attentionlinear attention, n.: A family of faster substitutes for attention whose cost grows in step with the length of the text rather than with its square, by keeping a running summary instead of comparing every pair of words.. Later versions target newer GPU generations.
Related
- scaled dot-product attentionscaled dot-product attention, n.: The standard formula for attention in transformers: score how well each question matches each label, scale the scores down, turn them into percentages, and use those to blend the stored information.
- kernel fusionkernel fusion, n.: Merging several small calculation steps into one, so a chip finishes them in a single pass instead of saving and reloading data in between, like cooking at the counter without returning ingredients to the fridge after each step.
- memory-boundmemory-bound, adj.: Describes a computer task slowed by how fast data can be fetched from memory rather than by how fast the chip can calculate — like a fast cook kept waiting for deliveries.
- SRAMSRAM, abbr.: The fastest kind of computer memory, built directly into processor chips. It is quick but takes up a lot of space, so chips hold only a little of it.
- high-bandwidth memoryhigh-bandwidth memory, n.: Very fast memory made by stacking memory chips on top of each other and placing them right next to the processor, so AI chips can be fed data quickly enough.
- long contextlong context, n.: The ability of an AI model to work with very long inputs — whole books, codebases or hours of transcripts — in one go, and the techniques that make this possible.
- sparse attentionsparse attention, n.: Attention in which each word looks at only a selected subset of the other words instead of all of them, so very long texts can be processed at a fraction of the cost.
- linear attentionlinear attention, n.: A family of faster substitutes for attention whose cost grows in step with the length of the text rather than with its square, by keeping a running summary instead of comparing every pair of words.
Sources
- Dao, T., Fu, D. Y., Ermon, S., Rudra, A. & Ré, C. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. NeurIPS 2022. arXiv:2205.14135. arxiv.org/abs/2205.14135
- Dao, T. (2023). FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning. arXiv:2307.08691. arxiv.org/abs/2307.08691
- Shah, J., Bikshandi, G., Zhang, Y., Thakkar, V., Ramani, P. & Dao, T. (2024). FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision. arXiv:2407.08608. arxiv.org/abs/2407.08608
Mentioned in
kernelkernel, n.: A word with several AI meanings: a similarity measure used by some learning methods; a small grid of numbers that scans an image inside a neural network; or a small program that runs on a GPU., long contextlong context, n.: The ability of an AI model to work with very long inputs — whole books, codebases or hours of transcripts — in one go, and the techniques that make this possible., memory bandwidthmemory bandwidth, n.: How fast data can be moved between a chip's memory and its calculators. For chatbots, this often decides how quickly words appear, more than raw calculating speed., scaled dot-product attentionscaled dot-product attention, n.: The standard formula for attention in transformers: score how well each question matches each label, scale the scores down, turn them into percentages, and use those to blend the stored information., arithmetic intensityarithmetic intensity, n.: How much calculating a task does for each piece of data it fetches from memory. Tasks with little calculation per fetch are slowed by memory, not by the chip's math speed., kernel fusionkernel fusion, n.: Merging several small calculation steps into one, so a chip finishes them in a single pass instead of saving and reloading data in between, like cooking at the counter without returning ingredients to the fridge after each step., memory wallmemory wall, n.: The growing gap between how fast chips can calculate and how fast they can fetch data from memory. Processors increasingly sit idle, waiting for numbers to arrive., memory-boundmemory-bound, adj.: Describes a computer task slowed by how fast data can be fetched from memory rather than by how fast the chip can calculate — like a fast cook kept waiting for deliveries., sparse attentionsparse attention, n.: Attention in which each word looks at only a selected subset of the other words instead of all of them, so very long texts can be processed at a fraction of the cost., SRAMSRAM, abbr.: The fastest kind of computer memory, built directly into processor chips. It is quick but takes up a lot of space, so chips hold only a little of it., TensorRT-LLMTensorRT-LLM, prop. n.: NVIDIA's free software library for running large language models as fast as possible on NVIDIA's own graphics chips, used by many companies that host AI models., TritonTriton, prop. n.: A free language from OpenAI for writing the fast inner routines of AI programs in Python-style code; its compiler handles hardware details that Nvidia's own tools leave to experts.