2.2b: FlashAttention — Online Softmax

Community Article
Published February 3, 2026

The Softmax Formula: Where Everything Begins

Before we can understand online softmax, we need to deeply understand regular softmax. Let's start with the formula and build up from there.

Given a vector of scores S = [s₀, s₁, s₂, ..., sₙ₋₁], softmax converts these into probabilities:

softmax(S)i=esies0+es1++esn1 \text{softmax}(S)_i = \frac{e^{s_i}}{e^{s_0} + e^{s_1} + \cdots + e^{s_{n-1}}}

Each output is the exponential of that score divided by the sum of all exponentials. This guarantees all outputs are positive and sum to 1 — exactly what we need for attention weights.

Let's compute a concrete example:

S = [2, 4, 1, 3]

Numerators:   exp(2) = 7.39,  exp(4) = 54.60,  exp(1) = 2.72,  exp(3) = 20.09

Denominator:  7.39 + 54.60 + 2.72 + 20.09 = 84.80

Softmax:      [7.39/84.80, 54.60/84.80, 2.72/84.80, 20.09/84.80]
            = [0.087, 0.644, 0.032, 0.237]

This works, but there's a serious problem lurking here...


The Numerical Stability Problem: Why We Need "m"

What happens if our scores are large? Let's try S = [100, 102, 99, 101]:

exp(100) = 2.69 × 10⁴³
exp(102) = 1.99 × 10⁴⁴
exp(99)  = 9.89 × 10⁴²
exp(101) = 7.31 × 10⁴³

These numbers are astronomically large! And it gets worse — if scores reach 710 or higher, exp() returns infinity in standard floating-point arithmetic. Our computation completely breaks.

The solution: subtract the maximum value before taking exp().

Here's the key mathematical insight. For any constant c, we can show:

esijesj=esicjesjc \frac{e^{s_i}}{\sum_j e^{s_j}} = \frac{e^{s_i - c}}{\sum_j e^{s_j - c}}

Why? Because:

esicjesjc=esiecjesjec=esijesj \frac{e^{s_i - c}}{\sum_j e^{s_j - c}} = \frac{e^{s_i} \cdot e^{-c}}{\sum_j e^{s_j} \cdot e^{-c}} = \frac{e^{s_i}}{\sum_j e^{s_j}}

The exp(-c) terms cancel out! So we can subtract any constant without changing the result.

The smart choice: let c = max(S). This is our "m":

m = max(s₀, s₁, ..., sₙ₋₁) — This is what "m" represents m is the maximum score. We subtract it from all scores before taking exp() to prevent numerical overflow.

With m subtracted, the largest value we ever compute exp() on is 0 (when sᵢ = m), giving exp(0) = 1. All other values are negative, giving exp() results between 0 and 1. No overflow possible!

Let's redo our example with the stable formula:

S = [2, 4, 1, 3]

m = max(2, 4, 1, 3) = 4

Shifted scores: S - m = [2-4, 4-4, 1-4, 3-4] = [-2, 0, -3, -1]

Numerators:   exp(-2) = 0.135,  exp(0) = 1.000,  exp(-3) = 0.050,  exp(-1) = 0.368

Denominator:  0.135 + 1.000 + 0.050 + 0.368 = 1.553

Softmax:      [0.135/1.553, 1.000/1.553, 0.050/1.553, 0.368/1.553]
            = [0.087, 0.644, 0.032, 0.237]

Same answer as before, but now the intermediate values are all manageable (between 0 and 1).


The Sum of Exponentials: What "l" Represents

Look at the denominator in our softmax formula: 0.135 + 1.000 + 0.050 + 0.368 = 1.553. This sum appears in the denominator of every softmax output. It's so important that we give it a name: l (for "logsumexp" or just "sum").

l = exp(s₀ - m) + exp(s₁ - m) + ... + exp(sₙ₋₁ - m) l = Σⱼ exp(sⱼ - m) — This is what "l" represents l is the sum of all exponentials (after subtracting m). It's the normalizing constant that makes softmax sum to 1.


The Key Insight: m and l Are All You Need

Now here's the crucial observation. Once you know m and l, you can compute any softmax value:

softmax(S)i=esiml \text{softmax}(S)_i = \frac{e^{s_i - m}}{l}

You don't need to store all the intermediate exp() values. You just need:

  • m (the maximum) — so you can compute exp(sᵢ - m) for any score sᵢ
  • l (the sum) — to divide by

Let's verify this with our example:

m = 4,  l = 1.553

softmax(S)₀ = exp(2 - 4) / 1.553 = exp(-2) / 1.553 = 0.135 / 1.553 = 0.087 ✓
softmax(S)₁ = exp(4 - 4) / 1.553 = exp(0) / 1.553  = 1.000 / 1.553 = 0.644 ✓
softmax(S)₂ = exp(1 - 4) / 1.553 = exp(-3) / 1.553 = 0.050 / 1.553 = 0.032 ✓
softmax(S)₃ = exp(3 - 4) / 1.553 = exp(-1) / 1.553 = 0.368 / 1.553 = 0.237 ✓

This is the foundation of online softmax: if we can compute the correct m and l without seeing all scores at once, we can compute the correct softmax.


The Challenge: Computing m and l in Blocks

Here's the problem FlashAttention faces. The scores vector might have millions of elements (one per key in the sequence). We can't fit them all in fast memory at once. We need to process them in blocks.

Full scores:  S = [2, 4, 1, 3]

Block A:  [2, 4]    ← We see this first
Block B:  [1, 3]    ← We see this second

Can we compute the correct global m and l by processing one block at a time?

Let's think about what we need:

Global m = max(2, 4, 1, 3) = 4
Global l = exp(2-4) + exp(4-4) + exp(1-4) + exp(3-4) = 1.553

Computing m incrementally is easy: just keep track of the maximum seen so far, and update it when you see a larger value.

After Block A:  m = max(2, 4) = 4
After Block B:  m = max(4, max(1, 3)) = max(4, 3) = 4  ✓

Computing l incrementally is trickier. The naive approach fails:

After Block A:  l_A = exp(2-4) + exp(4-4) = 0.135 + 1.000 = 1.135
After Block B:  l_B = exp(1-3) + exp(3-3) = 0.368 + 1.000 = 1.368

Naive sum: l_A + l_B = 1.135 + 1.368 = 2.503  ✗ (should be 1.553!)

What went wrong? Block B used its own local maximum (3) instead of the global maximum (4). The exponentials in l_B are computed with the wrong reference point!


The Solution: Rescaling When the Maximum Changes

Here's the key insight that makes online softmax work. We can correct values computed with one maximum to be as if they were computed with a different maximum.

The mathematical relationship:

exp(s - m_old)  and  exp(s - m_new)  are related by:

exp(s - m_new) = exp(s - m_old) × exp(m_old - m_new)

Let's verify this:

exp(s - m_new) = exp(s - m_old + m_old - m_new)
               = exp(s - m_old) × exp(m_old - m_new)  ✓

This means: if we have values computed with m_old, we can convert them to m_new by multiplying by exp(m_old - m_new).

This factor exp(m_old - m_new) is what we call the rescaling factor.


Worked Example: Processing Blocks Correctly

Let's redo our example, this time handling the maximum change properly.

Initial state:

m = -∞  (no scores seen yet)
l = 0   (no exponentials summed yet)

Process Block A: [2, 4]

BLOCK A: scores = [2, 4]

  1. Find this block's maximum: m_block = max(2, 4) = 4
  2. Update global maximum: m_new = max(m, m_block) = max(-∞, 4) = 4
  3. Rescale previous l if maximum changed: scale = exp(m - m_new) = exp(-∞ - 4) = 0l = l × scale = 0 × 0 = 0 (initialization, so rescaling has no effect)
  4. Compute this block's exponentials (using m_new): exp([2, 4] - 4) = [exp(-2), exp(0)] = [0.135, 1.000]
  5. Add to running sum: l = 0 + 0.135 + 1.000 = 1.135
  6. Update m for next iteration: m = m_new = 4

State after Block A: m = 4, l = 1.135

Process Block B: [1, 3]

BLOCK B: scores = [1, 3]

Current state: m = 4, l = 1.135

  1. Find this block's maximum: m_block = max(1, 3) = 3

  2. Update global maximum: m_new = max(m, m_block) = max(4, 3) = 4 (The global max stays at 4 — Block A had the max)

  3. Rescale previous l if maximum changed: scale = exp(m - m_new) = exp(4 - 4) = exp(0) = 1l = l × scale = 1.135 × 1 = 1.135 (Maximum didn't change, so no rescaling needed)

  4. Compute this block's exponentials (using m_new = 4): exp([1, 3] - 4) = [exp(-3), exp(-1)] = [0.050, 0.368]

    IMPORTANT: We use m_new = 4, not m_block = 3! This is what the naive approach got wrong.

  5. Add to running sum: l = 1.135 + 0.050 + 0.368 = 1.553

  6. Update m for next iteration: m = m_new = 4

State after Block B: m = 4, l = 1.553

Final result: m = 4, l = 1.553 — exactly what we computed with standard softmax!


When Rescaling Actually Happens

In the previous example, the maximum didn't change between blocks (Block A had score 4, Block B's max was only 3). Let's see an example where rescaling is actually needed.

New example where the maximum is in the second block:

S = [1, 2, 5, 3]

Block A: [1, 2]  (local max = 2)
Block B: [5, 3]  (local max = 5)

Global max = 5 (in Block B!)

Process Block A: [1, 2]

m_block = 2
m_new = max(-∞, 2) = 2
scale = exp(-∞ - 2) = 0
l = 0

exp([1, 2] - 2) = [exp(-1), exp(0)] = [0.368, 1.000]
l = 0 + 0.368 + 1.000 = 1.368

State after Block A:  m = 2,  l = 1.368

So far, l = 1.368 represents exp(1-2) + exp(2-2) = exp(-1) + exp(0).

Process Block B: [5, 3]

BLOCK B: scores = [5, 3]

Current state: m = 2, l = 1.368

  1. Find this block's maximum: m_block = max(5, 3) = 5

  2. Update global maximum: m_new = max(2, 5) = 5THE MAXIMUM CHANGED! (from 2 to 5)

  3. Rescale previous l: scale = exp(m - m_new) = exp(2 - 5) = exp(-3) = 0.050l = l × scale = 1.368 × 0.050 = 0.068

    What just happened? Our old l was: l_old = exp(1-2) + exp(2-2) = 1.368

    But with the correct global max of 5, it should be: l_correct = exp(1-5) + exp(2-5) = exp(-4) + exp(-3) = 0.018 + 0.050 = 0.068

    The rescaling gave us exactly 0.068!

  4. Compute this block's exponentials (using m_new = 5): exp([5, 3] - 5) = [exp(0), exp(-2)] = [1.000, 0.135]

  5. Add to running sum: l = 0.068 + 1.000 + 0.135 = 1.203

State after Block B: m = 5, l = 1.203

Verification:

Standard computation:
  m = max(1, 2, 5, 3) = 5
  l = exp(1-5) + exp(2-5) + exp(5-5) + exp(3-5)
    = exp(-4) + exp(-3) + exp(0) + exp(-2)
    = 0.018 + 0.050 + 1.000 + 0.135
    = 1.203 ✓

The rescaling step corrected our running sum to account for the new maximum.


Why Rescaling Works: Visual Intuition

Let me show you visually what the rescaling does:

RESCALING: CORRECTING THE REFERENCE POINT

After Block A (m = 2): Our exponentials are computed relative to max = 2: Score 1: exp(1 - 2) = exp(-1) = 0.368 Score 2: exp(2 - 2) = exp(0) = 1.000 l = 1.368

Block B arrives, and we discover the true max is 5. We NEED exponentials relative to max = 5: Score 1: exp(1 - 5) = exp(-4) = 0.018 Score 2: exp(2 - 5) = exp(-3) = 0.050 l should be 0.068

The relationship: 0.018 = 0.368 × exp(2 - 5) = 0.368 × 0.050 0.050 = 1.000 × exp(2 - 5) = 1.000 × 0.050 0.068 = 1.368 × exp(2 - 5) = 1.368 × 0.050

Multiplying by exp(m_old - m_new) converts everything!


Important Properties of the Rescaling Factor

The rescaling factor exp(m_old - m_new) has two important properties:

Property 1: It's always ≤ 1

Since m_new = max(m_old, m_block), we know m_new ≥ m_old. Therefore:

  • m_old - m_new ≤ 0
  • exp(m_old - m_new) ≤ exp(0) = 1

We're always scaling down, never up. This prevents numerical overflow.

Property 2: Scaling down is correct

When the new block has larger scores, the old scores become relatively less important. Scaling them down reflects this correctly. In the extreme case where the new block has much larger scores, the old contribution becomes negligible — which is exactly right for softmax (attention should focus on the highest scores).


Extending to Attention: Adding the Output Accumulation

So far we've computed m and l, which are enough to reconstruct softmax values. But for attention, we need the actual output — a weighted sum of values:

Output = Σᵢ softmax(S)ᵢ × Vᵢ = Σᵢ [exp(sᵢ - m) / l] × Vᵢ

We can rewrite this as:

Output = (1/l) × Σᵢ exp(sᵢ - m) × Vᵢ

Let's define the unnormalized output:

O_unnorm = Σᵢ exp(sᵢ - m) × Vᵢ This is the weighted sum BEFORE dividing by l. At the end, we compute: Output = O_unnorm / l

The key insight: O_unnorm needs rescaling just like l does! When the maximum changes, all the exponential weights change, so the accumulated output must be rescaled too.

The complete state we maintain:

  • m: running maximum of all scores seen
  • l: running sum of exp(score - m) for all scores seen
  • O: running unnormalized output (weighted sum with exp weights)

When the maximum changes, we rescale both l and O by the same factor exp(m_old - m_new).


Complete Worked Example: Attention with Two Blocks

Let's trace through a complete attention computation.

Setup:

Scores: S = [2, 4, 1, 3]

Values: V = [[1, 0],    ← value for key 0
             [0, 1],    ← value for key 1
             [1, 1],    ← value for key 2
             [0, 0]]    ← value for key 3

Block A: S_A = [2, 4],  V_A = [[1, 0], [0, 1]]
Block B: S_B = [1, 3],  V_B = [[1, 1], [0, 0]]

What should the output be?

Softmax: P = [0.087, 0.644, 0.032, 0.237]

Output = 0.087×[1,0] + 0.644×[0,1] + 0.032×[1,1] + 0.237×[0,0]
       = [0.087, 0] + [0, 0.644] + [0.032, 0.032] + [0, 0]
       = [0.119, 0.676]

Now let's compute this incrementally:

Initialization:

m = -∞
l = 0
O = [0, 0]

Process Block A: S_A = [2, 4], V_A = [[1, 0], [0, 1]]

BLOCK A

  1. m_block = max(2, 4) = 4, m_new = max(-∞, 4) = 4
  2. Rescale (initialization, scale = 0): l = 0 × 0 = 0, O = [0, 0] × 0 = [0, 0]
  3. Compute exponential weights: P_block = exp([2, 4] - 4) = [0.135, 1.000]
  4. Update l: l = 0 + 0.135 + 1.000 = 1.135
  5. Accumulate weighted values into O: O = [0, 0] + 0.135 × [1, 0] + 1.000 × [0, 1] = [0.135, 1.000]
  6. m = 4

State: m = 4, l = 1.135, O = [0.135, 1.000]

Process Block B: S_B = [1, 3], V_B = [[1, 1], [0, 0]]

BLOCK B

Current state: m = 4, l = 1.135, O = [0.135, 1.000]

  1. m_block = max(1, 3) = 3, m_new = max(4, 3) = 4 (no change)
  2. Rescale (scale = exp(4-4) = 1, no change): l = 1.135, O = [0.135, 1.000]
  3. Compute exponential weights (using m_new = 4): P_block = exp([1, 3] - 4) = [0.050, 0.368]
  4. Update l: l = 1.135 + 0.050 + 0.368 = 1.553
  5. Accumulate weighted values into O: O = [0.135, 1.000] + 0.050 × [1, 1] + 0.368 × [0, 0] = [0.185, 1.050]
  6. m = 4

State: m = 4, l = 1.553, O = [0.185, 1.050]

Final normalization:

Output = O / l = [0.185, 1.050] / 1.553 = [0.119, 0.676] ✓

This matches our expected output exactly!


Example with Rescaling: When Maximum Changes

Let's see rescaling in action with values too.

Setup:

Scores: S = [1, 2, 5, 3]  (max is in Block B!)

Values: V = [[1, 0], [0, 1], [2, 0], [0, 2]]

Block A: S_A = [1, 2],  V_A = [[1, 0], [0, 1]]
Block B: S_B = [5, 3],  V_B = [[2, 0], [0, 2]]

Process Block A:

m_block = 2, m_new = 2
P_block = exp([1, 2] - 2) = [0.368, 1.000]
l = 0.368 + 1.000 = 1.368
O = 0.368 × [1, 0] + 1.000 × [0, 1] = [0.368, 1.000]

State: m = 2, l = 1.368, O = [0.368, 1.000]

Process Block B: (Rescaling happens!)

BLOCK B (with rescaling)

Current state: m = 2, l = 1.368, O = [0.368, 1.000]

  1. m_block = 5, m_new = max(2, 5) = 5

  2. RESCALE! scale = exp(2 - 5) = exp(-3) = 0.050

    l = 1.368 × 0.050 = 0.068

    O = [0.368, 1.000] × 0.050 = [0.018, 0.050]

    Both l and O are scaled down because the old scores (1 and 2) are now much less significant compared to the new high score (5).

  3. P_block = exp([5, 3] - 5) = [1.000, 0.135]

  4. l = 0.068 + 1.000 + 0.135 = 1.203

  5. O = [0.018, 0.050] + 1.000 × [2, 0] + 0.135 × [0, 2] = [2.018, 0.320]

State: m = 5, l = 1.203, O = [2.018, 0.320]

Final normalization:

Output = O / l = [2.018, 0.320] / 1.203 = [1.678, 0.266]

Verification:

m = 5
exp([1,2,5,3] - 5) = [0.018, 0.050, 1.000, 0.135]
l = 1.203
P = [0.015, 0.042, 0.831, 0.112]

Output = 0.015×[1,0] + 0.042×[0,1] + 0.831×[2,0] + 0.112×[0,2]
       = [0.015, 0] + [0, 0.042] + [1.662, 0] + [0, 0.224]
       = [1.677, 0.266] ✓ (matches within rounding)

Scaling to Multiple Queries

In attention, we have multiple queries, each producing its own output. The beautiful thing is: each query row is completely independent.

Each query maintains its own (m, l, O):

  • Query 0 has m[0], l[0], O[0]
  • Query 1 has m[1], l[1], O[1]
  • ... and so on

When processing a key-value block, we update each query's state independently. Different queries might have different maximums — that's fine, they don't interfere with each other.

MULTIPLE QUERIES: INDEPENDENT PROCESSING

  • Q = [q₀, q₁] (2 queries)
  • K = [k₀, k₁, k₂, k₃] (4 keys, processed in blocks of 2)

Attention scores (each query has its own row):

  • S = [[s₀₀, s₀₁, s₀₂, s₀₃], ← Query 0's scores
  • [s₁₀, s₁₁, s₁₂, s₁₃]] ← Query 1's scores

State for Query 0: m[0], l[0], O[0]

State for Query 1: m[1], l[1], O[1]

These are updated independently. Query 0 might have its max in Block A while Query 1 has its max in Block B.


The Complete Online Softmax Algorithm

Now we can state the full algorithm:

ONLINE SOFTMAX FOR ATTENTION (ONE QUERY)

Input: s₀, s₁, ..., sₙ₋₁ (scores, processed in blocks), v₀, v₁, ..., vₙ₋₁ (values, processed in blocks)

Output: Attention output = Σᵢ softmax(S)ᵢ × vᵢ

Initialize: m = -∞ (running maximum), l = 0 (running sum of exponentials), O = 0 (running unnormalized output)

For each block of scores and values (S_block, V_block):

  1. m_block = max(S_block)

  2. m_new = max(m, m_block)

  3. Rescale if maximum changed:

    scale = exp(m - m_new)

    l = l × scale

    O = O × scale

  4. Process this block:

    P_block = exp(S_block - m_new)

    l = l + sum(P_block)

    O = O + P_block · V_block

  5. m = m_new

Final normalization: Output = O / l

image

For multiple queries, run this algorithm independently for each query (in parallel).


Why This Algorithm is Exact

Online softmax produces mathematically identical results to standard softmax. It's not an approximation.

The proof follows from our observation that m and l fully determine softmax. At each step:

  • m is the true maximum of all scores seen so far
  • l is the true sum of exponentials (relative to m) for all scores seen so far
  • O is the true unnormalized output for all scores seen so far

The rescaling step maintains these invariants when the maximum changes. After processing all blocks, we have the correct global m and l, and O/l gives the correct output.


Summary

The softmax formula:

softmax(S)i=esiml \text{softmax}(S)_i = \frac{e^{s_i - m}}{l}

where:

  • m = max(S) — maximum score (for numerical stability)
  • l = Σⱼ exp(sⱼ - m) — sum of exponentials (normalizing constant)

The key insight: m and l are sufficient to compute softmax. If we can compute them incrementally, we can compute softmax without seeing all scores at once.

The challenge: When processing in blocks, different blocks might have different local maximums. Values computed with one maximum are "wrong" if a later block has a larger maximum.

The solution: When the maximum changes from m_old to m_new, multiply all previous results by exp(m_old - m_new). This rescales them to be correct with respect to the new maximum.

For attention: We also maintain O (unnormalized output). It gets rescaled along with l whenever the maximum changes. At the end, Output = O / l.

The result: Exact attention computation, processing scores and values in small blocks, without ever materializing the full N×N attention matrix.


Check Your Understanding

  1. Why do we subtract the maximum before computing exp()? What would happen if we didn't?
    • Purpose: Numerical stability. Without subtracting the max, exp() can overflow.
    • What happens without it: For scores like [100, 102, 99], exp(100) ≈ 2.69 × 10⁴³. At scores ≥ 710, exp() returns infinity in standard floating-point.
    • With max subtraction: The largest value we ever compute is exp(0) = 1 (when sᵢ = m). All other values are exp(negative) ∈ (0, 1). No overflow possible.
    • Mathematical justification: Subtracting any constant c from all scores doesn't change the softmax result — the exp(-c) terms cancel in numerator and denominator.
  2. If m = 3, l = 2.5, and a new block arrives with maximum 7, what is the rescaling factor? What happens to l?
    • Rescaling factor: scale = exp(m_old - m_new) = exp(3 - 7) = exp(-4) ≈ 0.0183
    • New l: l = l × scale = 2.5 × 0.0183 ≈ 0.046
    • Interpretation: The old scores (with max 3) are now 4 units below the new max (7). Their exponentials shrink by a factor of ~55×. This is correct — softmax should now heavily favor the new high-scoring elements.
  3. Suppose we process scores [10, 20, 30] in three blocks of one element each. What are m and scale after each block?
    • Block 1 (score = 10):
      • m_new = max(-∞, 10) = 10
      • scale = exp(-∞ - 10) = 0 (initialization, no effect)
      • After: m = 10
    • Block 2 (score = 20):
      • m_new = max(10, 20) = 20
      • scale = exp(10 - 20) = exp(-10) ≈ 0.0000454
      • After: m = 20
    • Block 3 (score = 30):
      • m_new = max(20, 30) = 30
      • scale = exp(20 - 30) = exp(-10) ≈ 0.0000454
      • After: m = 30
    • Pattern: Each new block has a higher max, so we rescale every time. The cumulative effect: earlier scores become exponentially less significant.
  4. Why do we rescale O by the same factor as l? What would go wrong if we only rescaled l?
    • O contains weighted values: O = Σ exp(sᵢ - m) × Vᵢ. Those exp() weights were computed with the old m.
    • If we only rescale l: The weights in O would be "too large" relative to the new l. The ratio O/l would be wrong.
    • Example: Old m=2, new m=5. Old weights used exp(s-2), new weights use exp(s-5). If l is scaled down by exp(-3) but O isn't, then O/l is 20× too large.
    • Both must use the same reference: Rescaling both by exp(m_old - m_new) maintains the invariant that O and l are computed with the same m.

Community

Sign up or log in to comment