In this article
Before you start
This page belongs to the deep learning path and assumes only the ideas below.
Why interviewers ask this
Online softmax matters in interviews because it is the idea that makes FlashAttention possible, and a candidate who can only say "FlashAttention is faster" has not explained why. A question that starts from attention cost, graphics processing unit (GPU) memory or numerically stable softmax can lead here within a follow-up or two, and candidates report variants of it in machine learning systems rounds.
A weak answer calls online softmax an approximation or a trick for streaming data. A strong answer says it is exact, names the running maximum and the running sum, and writes the rescaling factor. A Staff-level answer adds why it saves time: softmax is limited by how much data moves through memory, not by arithmetic, so removing a pass or a whole matrix write is what pays off. The deep learning interview questions page lists neighbouring questions of the same kind.
Online softmax from first principles
- Online softmax
- Online softmax is an algorithm that computes the maximum and the sum of exponentials of a vector in a single pass, keeping a running maximum and a running sum that is rescaled whenever the maximum grows, and it returns exactly the same result as ordinary softmax.
- Running normalizer
- A running normalizer is the partial sum of exponentials of the elements read so far, measured relative to the largest value read so far, and it is the number that the final division of softmax needs.
In plain words
Imagine a teacher who records each exam score as a share of the best score so far. When a new best score appears, every earlier share is now too big. The teacher does not reread the old exams. She multiplies the running total by one number, and the total is correct again. The analogy breaks in one place: the real shares are exponentials of the scores, so the single correcting number is also an exponential.
In one precise sentence, online softmax keeps a pair (m, d) for the elements read so far, where m is their maximum and d is the sum of exp(x - m), and it updates the pair one element at a time.
Softmax itself maps a vector x to weights:
Step 1: Naive softmax needs two passes, and safe softmax needs three
Naive softmax first adds up exp(x) over the whole vector, then divides each exp(x) by that sum, so it reads the vector twice. Milakov and Gimelshein (2018) count three memory accesses per element for it: two loads and one store. It also breaks on large scores, because exp(x) overflows in 32-bit floating point once x passes about 88.7.1 Subtracting the largest value first fixes that and changes no output, because the same factor cancels in the ratio.
That version is called safe softmax. It needs three passes (maximum, sum, outputs) and four memory accesses per element, and the same paper notes that all major deep learning frameworks used it when it was written. For why the subtraction is safe, see why softmax needs numerical stabilization. What you now know: stability costs an extra pass over the data, and the next step removes it.
Step 2: A new maximum rescales the running sum
Online softmax reads the vector once and keeps two numbers, the running maximum m and the running normalizer d. It starts with m equal to minus infinity and d equal to zero. For each new element x it takes the new maximum, shrinks the old sum, and adds the new term:
When a new maximum appears, the running sum is multiplied by the exponential of the old maximum minus the new one. That single multiplication is valid because every earlier term was measured against the old maximum, and exp(x - old) times exp(old - new) equals exp(x - new). The factor is never above one, so it cannot overflow. The paper also shows that d stays between 1 and j after j elements, so the sum stays in a safe range.
Input: one element and the pair (m, d). Output: the updated pair. The idea comes from the online way of computing variance, and "online" here means one element at a time, not real-time or streaming inference.
Step 3: A second pass turns the final pair into outputs
After the single pass, m is the maximum of the whole vector and d is the sum of exp(x - m) over all of it. A second loop computes each output as exp(x_i - m) / d. Online softmax needs two passes where safe softmax needs three, and three memory accesses per element where safe softmax needs four, which is a quarter fewer accesses.1
The normalizer comes from one pass, but the outputs still need a second one. So online softmax does not make softmax a one-pass operation. Softmax is memory-bound, so fewer accesses is what matters. The paper's measurements on a V100 GPU show about 1.15 times speedup for vectors of at least 1,000 elements and up to about 1.3 times at large batch sizes with at least 4,000 elements. Fused with a top-5 selection over 25,000 elements, the paper reports up to 5 times, of which 2.5 times came from fusion and 2 times from online softmax itself (Milakov and Gimelshein, 2018).
Step 4: Partial results merge in any order
The pair (m, d) of one block of elements can be merged with the pair of another block: take the larger maximum, then rescale each block's sum to that maximum and add them. The paper shows this merge operator is associative, so a vector can be split across threads or blocks and combined in any grouping. Written out, with m for the merged maximum:
The merge is not a plain sum of two pairs: each block's sum is first rescaled to the larger maximum.
This is the bridge to attention. FlashAttention is an exact attention algorithm that processes the score matrix block by block in fast on-chip memory, using this merge so the full matrix is never written out. Dao et al. (2022) say softmax "couples columns of K", so they decompose it and keep extra statistics (m, ℓ) to compute it one block at a time.
Step-by-step description
- Safe softmax reads the vector three times.
- A new maximum shrinks the running sum.
- A second pass divides by the final sum.
- Partial pairs combine in any order.
- Online softmax needs two passes where safe softmax needs three.
- When a new maximum appears the running sum is multiplied by the exponential of the old maximum minus the new one.
- Merging partial results gives the same sum in any grouping.
Knowing it is half the job. Saying it is the other half.
Interview Clozo is an AI interview copilot for real interview rounds: it hears the question and puts a strong answer on your screen, and you say it in your own words. See the features or how it works.
Start Free Trial →Worked example
Take the illustrative vector [2, 1, 3, 0.5] and read it one element at a time, keeping the pair (m, d).1
| Element read | Running maximum m | Running normalizer d |
|---|---|---|
| start | minus infinity | 0 |
| 2 | 2 | 1.0 |
| 1 | 2 | 1.3679 |
| 3 | 3 | 1.5032 |
| 0.5 | 3 | 1.5853 |
The row for element 3 is the interesting one. The element 3 raises the maximum from 2 to 3, so the old sum 1.3679 is multiplied by exp(2 - 3), about 0.3679, giving 0.5032, and the new term exp(3 - 3) equals 1 is added, which gives 1.5032.1 The final outputs exp(x - 3) / 1.5853 are about 0.2321, 0.0854, 0.6308 and 0.0518, and they sum to one. A two-pass safe softmax gives the same sum, 1.5853, and across 1,000 random vectors (seeded, illustrative) the largest difference between the two methods was below one part in a trillion.1
Merging works the same way. Split the vector into the blocks [2, 1] and [3, 0.5]. Block one has the pair (2, 1.3679) and block two has (3, 1.0821). The merged maximum is 3, and the merged sum is 1.3679 times exp(2 - 3) plus 1.0821, which is 1.5853 again.1
Naive softmax makes three memory accesses per element, safe softmax makes four and online softmax makes three.
Text description of the diagram
- Naive softmax makes three memory accesses per element, safe softmax makes four and online softmax makes three.
- Online softmax keeps safe softmax's stability but cuts memory accesses from four to three per element.
Why this matters for attention: for a sequence of N tokens, one attention head scores every token against every other, so it holds N times N scores. At N equal to 1,024 that is 1,048,576 scores, about 0.002 gibibytes in 16-bit floating point (fp16); at 16,384 it is 0.5 gibibytes; at 65,536 it is 8 gibibytes, per head, per layer and per sequence in the batch.1 Standard implementations "materialize the matrices S and P to high-bandwidth memory (HBM)", which takes O(N^2) memory, where high-bandwidth memory (HBM) is the large, slower memory of a GPU (Dao et al., 2022).
Trade-offs and edge cases
FlashAttention's use of online softmax trades extra arithmetic for far fewer memory accesses. In the measurement of Dao et al. (2022) on a generative pre-trained transformer (generative pre-trained transformer (GPT)-2) medium model with sequence length 1,024 on an A100, forward and backward together, standard attention took 66.6 gigaflops, 40.3 gigabytes of HBM traffic and 41.7 milliseconds, while FlashAttention took 75.2 gigaflops, 4.4 gigabytes and 7.3 milliseconds. Floating-point operations (FLOPs) went up, traffic went down by about nine times, and the run time followed the traffic.
Text description of the diagram
- FlashAttention does more arithmetic than standard attention.
- FlashAttention moves less data through high-bandwidth memory than standard attention and runs faster.
| Method | Passes over the vector | Memory accesses per element | Can overflow |
|---|---|---|---|
| Naive softmax | 2 | 3 | yes |
| Safe softmax | 3 | 4 | no |
| Online softmax | 2 | 3 | no |
Edge cases and limits:
- Hardware dependence. The 1.15 and 1.3 times figures are from a V100 benchmark; the paper reports no central processing unit (CPU) tests, so the speedup is not universal.
- Exact, not approximate. FlashAttention returns exact attention with O(N) extra memory and O(N²d) arithmetic, where d is the head size; only its block-sparse extension is approximate.
- Backward pass. FlashAttention stores the output and the statistics (m, ℓ) and recomputes the score blocks on-chip, which costs more arithmetic and saves HBM traffic.
- Neighbouring work. Rabe and Staats (2022) described a practical implementation, built on a chunked, maximum-tracking idea, that needs memory of the order of the square root of n; FlashAttention targets memory accesses and fuses the work into one kernel. FlashAttention-2 keeps the algorithm, cuts non-matrix-multiply operations, and reports around 2 times speedup over FlashAttention.
- Where the merge is used today. The Flash-Decoding write-up (Dao, Haziza, Massa and Sizov, October 2023) splits the keys and values into chunks, computes attention for each chunk together with its log-sum-exp, then runs a final reduction that rescales and combines the chunks, which applies the merge of Step 4 across the sequence. FlashAttention-3 (Shah et al., July 2024) targets Hopper-generation GPUs, uses asynchrony and 8-bit floating point (FP8), and reports 1.5 to 2.0 times speedup over FlashAttention-2 in 16-bit precision. Libraries differ in which kernels they use and change defaults over time, so as of October 2026 check the documentation of the version you run.
- Number formats. Narrow formats overflow at much smaller values than 32-bit floating point, so the range of the running pair is a design choice; the fp16 versus bf16 page explains how formats set overflow limits.
How to answer it in an interview
What is online softmax, and why does FlashAttention need it?
Online softmax computes the softmax normalizer in a single pass. I keep a running maximum and a running sum of exponentials. When a new element raises the maximum, I multiply the old sum by the exponential of the old maximum minus the new maximum, then add the new term. It stays exact and stable, and it cuts softmax from three passes to two. The merge step is associative, so blocks combine in any order, which is how FlashAttention avoids writing the full attention matrix to memory.
The spoken structure: say what the pair (m, d) is, give the rescaling factor, state that it is exact, then connect it to blocks and memory traffic.
Common mistakes
- "Online softmax is an approximation." (wrong) The paper proves it equals safe softmax. (right) "It is exact; it only changes the order of the computation."
- "Online softmax removes the need to subtract the maximum." (wrong) It does the subtraction incrementally with the running maximum. (right) "It keeps max-subtraction, applied as the maximum changes."
- "It makes softmax a one-pass operation." (wrong) The normalizer takes one pass but the outputs need a second. (right) "Two passes instead of three."
- "FlashAttention is fast because it does fewer FLOPs." (wrong) It does slightly more. (right) "It does far fewer reads and writes to HBM, and softmax is memory-bound."
- "Online means streaming or real-time." (wrong) The word refers to updating one element at a time, as in online variance. (right) "Online means incremental update."
Follow-up questions
How does flash attention use online softmax?
For each row of scores it keeps running statistics (m, ℓ) and a partial output. When a new block of keys and values arrives, it rescales the old statistics and the old output by the exponential of the old maximum minus the new maximum, then adds the new block's contribution. The full matrix never goes to HBM.
What is the difference between safe softmax and online softmax?
Both subtract the maximum and give identical results. Safe softmax finds the maximum in a separate pass first; online softmax finds it while summing, so it needs two passes instead of three.
Is online softmax exact or an approximation?
It is exact up to ordinary floating-point rounding. Dense FlashAttention is also exact; only its block-sparse variant approximates.
Why does the merge step have to be associative?
Associativity lets blocks be combined in any grouping, so work can be split across threads and blocks and the result does not depend on how it was split. See also more deep learning interview questions.
Why is FlashAttention faster if it does more arithmetic?
Attention is limited by memory traffic, not by arithmetic. FlashAttention moves about a ninth of the data of the standard version in the measurement above, so it finishes sooner despite extra work.
- Online softmax finds the maximum and the normalizer in one pass using a running pair (m, d).
- When the maximum grows, multiply the old sum by
exp(m_old - m_new)before adding the new term. - The result is exact, and it cuts softmax from three passes to two and from four memory accesses to three.
- The merge of two partial pairs is associative, which allows block-by-block attention.
- FlashAttention gains speed from fewer HBM reads and writes, not from fewer arithmetic operations.
Sources
- Online normalizer calculation for softmax (arXiv 1805.02867) — arXiv (Milakov and Gimelshein, NVIDIA), 28 Jul 2018
- FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (arXiv 2205.14135) — arXiv / NeurIPS 2022 (Dao, Fu, Ermon, Rudra, Re), 23 Jun 2022
- FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning (arXiv 2307.08691) — arXiv (Dao), 17 Jul 2023
- Self-attention Does Not Need O(n^2) Memory (arXiv 2112.05682) — arXiv (Rabe and Staats, Google Research), 10 Oct 2022
- Flash-Decoding for long-context inference — Stanford CRFM (Dao, Haziza, Massa, Sizov), 12 Oct 2023
- FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision (arXiv 2407.08608) — arXiv (Shah, Bikshandi, Zhang, Thakkar, Ramani, Dao), 11 Jul 2024
