What is Flash Attention?

Author: Data Science in your pocket · Watch on YouTube

Video authored by the channel shown, not by IANIA; all rights belong to the creator. Transcript and summary generated automatically; the original-language content is preserved.

Original language: English

Original video: What is Flash Attention? · Channel: Data Science in your pocket · 11:33 · Watch on YouTube

Topics: Model inference and optimization, AI models

Summary

Flash Attention makes standard attention faster and more memory-efficient by reducing data movement between a GPU’s fast, small SRAM and its slower, larger HBM. It fuses operations into one kernel, processes attention in tiles, and uses online softmax to reproduce the exact result without storing the full score and probability matrices. For the backward pass, it recomputes probability tiles instead of keeping the full matrix in memory. These techniques reduce memory use from quadratic to linear in sequence length and can substantially speed up long-context workloads.

Key points

Standard attention can be bottlenecked by memory bandwidth and the repeated movement of intermediate data, rather than by arithmetic.
Kernel fusion combines attention operations so the full score and probability matrices are not written to HBM.
Tiling processes blocks of Q, K, and V that fit in SRAM instead of loading the entire attention matrix.
Online softmax tracks a running maximum and exponential sum, rescaling them as each new block arrives.
The tiled calculation produces the exact same attention result; it is not an approximation.
Flash Attention stores the output and small per-row statistics, then recomputes probability blocks during the backward pass.
The described memory use falls from O(N²) to O(N·D), linear in sequence length.
Flash Attention is especially useful for long sequences and memory-bound workloads; its benefits may be smaller for short sequences or tiny models.
The video presents Flash Attention 2 as improving parallelism and Flash Attention 3 as targeting Hopper GPUs with techniques including warp specialization.

Chapters

Start	Chapter	Summary
00:00	Why attention needs a better implementation	The video introduces attention’s quadratic memory cost and recaps how standard self-attention computes its output.
01:35	The memory-bandwidth bottleneck	It contrasts HBM and SRAM and explains how repeated data transfers slow standard attention.
02:38	Kernel fusion, tiling, and online softmax	The video describes fusing attention into one kernel, processing it in tiles, and using online softmax to handle whole-row normalization.
06:52	Backward pass and performance claims	It explains recomputing probability tiles during backpropagation, summarizes the memory-access analysis, and reviews reported improvements across Flash Attention versions.
09:29	When to use Flash Attention	The video discusses suitable workloads, ways to enable Flash Attention, and recaps its four main techniques.

Quotes

No approximation, no math shortcuts. — 00:30

You must stop the walking. — 02:38

Flash attention is exact attention. — 08:26

The result exact attention linear memory multiffold speedups from respecting the hardware instead of fighting it. — 11:02

Concepts and entities

Flash Attention, Self-attention, softmax, online softmax, tiling, kernel fusion, HBM, SRAM, A100, FlashAttention-2, FlashAttention-3, H100, PyTorch

Full transcript (original language)

00:00 Welcome to data science in your pocket. Here is a surprising fact. When you train a GPT scale model, the biggest eater of time and memory is attention itself. [music] The very mechanism that makes transformers smart. Attention’s cost grows quadratically with sequence length. [music] And worse, the standard implementation spends most of its time shuffling data around the GPU instead of computing. [music] So, here is the puzzle. Flash attention is dramatically faster, [music] uses linear instead of quadratic

00:30 memory, and computes the exact same answer. No approximation, no math shortcuts. How? By the end of this video, you will understand the trick down to the exponentials with a worked example you can verify by hand. Let’s go. Quick recap of self attention. Start with input X, one row per token, n rows of width D. [music] Multiply X by three learned matrices W sub Q, W subK, W subV to get Q, [music] K, and V. Then

01:04 the scores S= Q * KRpose divided by the square root of D. Apply softmax row by row to get probabilities P and the output O= P [music] * V. Here is the catch. Sn byN and the standard implementation writes both in full [music] to GPU memory. Say N is 8,192 tokens. N squared is about 67 million entries and in 32bit floats that is roughly 268 megabytes for one attention

01:35 matrix [music] of one head of one layer. Real models have dozens of each [music] and you need the backward pass too. You can see where this goes. Here is the mental shift. On modern GPUs, attention’s bottleneck is not arithmetic. It is memory bandwidth. Your GPU has two kinds of memory. HBM is the big one, 40 [music] GB on an A100, but only about 1555 GB per second, sitting far from the compute units. SRAM is

02:08 tiny, roughly 20 megabytes total on an A100, but [music] blindingly fast, around 19 terabytes per second. Now watch standard attention. Read Q and K from HBM. Write S [music] back. Read S, write P, read P andV, write O. Round trip after round trip. Analogy. [music] The compute units are a brilliant lightning fast kitchen. And HBM is a warehouse across town. The chef barely cooks. [music] The day is wasted. Walking to the warehouse with

02:38 half-finish dishes. Making the kitchen faster changes [music] nothing. You must stop the walking. Everything flash attention does comes from that one insight. Idea one kernel fusion. Normally every operation matrix multiply scale softmax [music] runs as its own GPU kernel and each kernel reads inputs from HBM and writes outputs back. Fusion says write one single CUDA kernel that does all of attention end to end. Load Q, [music] K,

03:11 and V into SRAMM once. Compute everything on chip and write out only the final O. The N byN matrices S and P never exist in slow memory at all. One problem, [music] SRAMM is tiny. Each streaming multiprocessor on an A100 has about 192 kilob of shared memory. A hundreds of megabytes attention matrix will never fit in 192 kilob. We cannot load the whole problem onto the chip. [music] So we need to chew it in small pieces. That is idea two.

03:44 Idea two, tiling. Split Q into row blocks [music] and K and V into column blocks. The fused kernel loads one Q block into SRAMM and keeps it there. [music] Then loops over K and V blocks. Pull one pair into SRAM. Compute [music] a small tile of the score matrix. Update a running output. Move on. When the loop ends, that block of O is final. On an A100, typical blocks are 64 or 128 rows and columns, small enough to fit comfortably [music] in the available

04:15 shared memory. The giant attention matrix gets processed tile by tile, like reading a huge book one page at a time instead of photocopying it all onto your desk. But there is a snake in the grass. Softmax is not block local. Fixing that is the mathematical heart of flash attention. The core problem softmax needs the maximum and the exponential sum of the whole row but tiling shows us the row one block at a time. The fix is online softax from Milikov and Gimmelshine in

04:46 2018. Maintain two running statistics per [music] row m the max seen so far and l the running sum of exponential scaled to that max. When a new block arrives with block max m subblock update like this. The new max m prime is the maximum of m and m subblock. Compute [music] two risk scale factors. Alpha= e to the m minus m prime. Correcting all old exponentials and beta= e to the m subblock [music] minus m prime. Putting the new block on the

05:17 same scale. The new sum l prime equals alpha * l + beta * the block sum of e [music] to the x - m subblock. Rescale the output accumulator the same way. O ACC becomes alpha * occ plus beta * p [music] block * v block. At the end divide by l prime and it is numerically stable because every exponential is of a number at most zero. No overflow ever. Tiny example you can verify for scores

05:49 [music] 1 2 3 zero tiled into two blocks of two. Block one scores one and two. Max M1 equals [music] 2. Sum L1 is E the minus1 + E to the 0 about 0.368 + 1. So [music] 1.368. Now block 2 3 and zero. Block max three beats are running max. So m2 equals 3. Rescale the old sum by e to the 2 minus

06:20 3 about 0.368. New sum L2 equ= 0.368 * 1.368 [music] which is 0503 plus the new blocks contribution E to [music] the 0 + EUS 3 1.0498 total 1.553. Now the check textbook soft max denominator e to the 1 + e 2 + e 3 + e to the 0 31.19

06:52 scale by e to the minus3 and you get [music] 1.553 exactly the same exact soft max never holding more than two scores at once. The backward pass is where standard attention gets punished twice. >> [music] >> Back prop through softmax needs the full N byN matrix P kept in HBM from the forward pass. Flash attention throws P away. It stores only the output O and the pero statistics M and L two floats per row order of N memory instead of

07:24 order of N squared. During backward, it recomputes [music] P block by block in SRAMM using the same tiling loop. Wasteful. Remember the kitchen, the chef’s [music] time is cheap. The warehouse trips are expensive. Recomputing a tile costs flops, [music] but reading P from HBM costs far more time than those flops. Doing more arithmetic makes backward faster and the training time memory footprint [music] collapses right where you need the headroom. Now the theory, because the authors

07:55 proved it, count HBM accesses. Standard attention moves on the order of N * D + [music] N^ 2 bytes. The n^ 2 term dominates. Flash attention with blocks tuned to SRAMM of size M needs only on the order of N^ 2 * D ^ 2 / M with head dimension D of 64 to 128 and M around 100 kilobyt that is roughly [music] 10 to 20 times fewer bytes crossing the slow boundary. And since S and P are never materialized, [music] memory drops from

08:26 O of N^ 2 to O of N * D, linear in sequence length. Tattoo this on your brain. Flash attention is exact attention. Not sparse, not low rank, not a kernel approximation. Same [music] function, same gradients, same quality, purely a smarter schedule of the same arithmetic engineered around the memory hierarchy. A free lunch that is actually free. What does it buy in practice? These are figures reported by the papers, so treat

08:58 them as approximate. The original flash attention from Try Dow and colleagues in 2022 ran attention roughly two to four times faster [music] and reported up to a 7.6 time speed up on GPT2 training end to end because attention dominated at long sequences. Linear memory pushed context windows from 16,000 tokens to 64,000 beyond. Flash Attention 2 in 2023 roughly [music] doubled speed again by parallelizing over the sequence length dimension and cutting non-matrix

09:29 multiply [music] flops feeding the tensor cores better. Flash Attention 3 targeted Hopper GPUs [music] like the H100 using warp specialization to overlap compute with data movement plus [music] FP8 reaching up to about 75% of theoretical peak throughput. Attention went from bottleneck to one of the most optimized kernels in deep learning. When should you use it? Long sequences 4,000 tokens and up where the quadratic wall bites. Training or fine-tuning

10:00 large language models [music] where saved memory becomes bigger batches or longer contexts. Memory bound inference with long prompts. Honestly, any transformer with a big context. [music] When does it not matter? short sequences under a thousand tokens where everything fits in cache anyway and tiny models where tiling overhead can eat [music] the gains and you rarely implement anything in PyTorch. Just call scaled.product attention and it auto selects a flashbackend when possible or

10:31 install the flash attention package from try dowo in hugging face pass attention implementation equals flash attention too when loading the model one line. It has become the default way to run attention which is why understanding the machinery is worth your time for ideas make flash attention work. One kernel [music] fusion run attention in a single kernel so the giant intermediates never touch HBM. Two [music] tiling chop Q K and V into

11:02 SRAM sized blocks and loop. Three, online soft max. [music] A running max and running sum with rescaling exact even though you see one block at a time. Four, recomputation. Store only the output and tiny per row stats. Rebuild P on the fly during backward [music] because onchip math beats memory reads. The result exact attention linear memory multiffold speedups from respecting the hardware instead of [music] fighting it. If this clicked, the natural next stops are the KV cache for fast [music]

11:33 generation and page detention, which applies the same memory thinking to inference serving. Both are linked on screen. Subscribe to Data Science in Your Pocket for more deep dives where we actually do the math. Thanks for [music] watching.