FlashAttention Algorithm › Online Softmax Computation

Tiling creates an elegant solution for memory but introduces a subtle mathematical challenge. Softmax requires global information: to normalize any single attention score, you must know the sum of exponentials across the entire row. How can we compute softmax incrementally when we only see one block of keys at a time?

Consider the standard softmax formula:

$$\text{softmax}(x_i) = \frac{e^{x_i}}{\sum_{j=1}^{n} e^{x_j}}$$

where:

• $x_i$: the $i$-th attention score in a row (the score we want to normalize)
• $e^{x_i}$: the exponential of that score. This keeps a positive value
• $\sum_{j=1}^{n} e^{x_j}$: the sum of exponentials over all $n$ scores in the row, serving as the normalizing constant
• $n$: the total number of positions in the sequence

The denominator $\sum_{j=1}^{n} e^{x_j}$ presents our problem: we need to sum over all positions, but we're processing keys in blocks. After seeing only the first block, we don't know what scores the remaining blocks will contribute. How can we produce a valid softmax output without waiting until all blocks are processed?

The answer lies in a clever reformulation: rather than computing softmax in one pass, we maintain running statistics that allow us to incrementally update our computation as new blocks arrive. This technique, called online softmax, is the mathematical heart of FlashAttention.

Before diving into formulas, let's think about what information we need to maintain across blocks. Softmax has two key properties we must preserve:

1. Numerical stability: Exponentiating large values causes overflow. The standard trick is to subtract the maximum value before exponentiating: $e^{x_i - \max(x)}$ instead of $e^{x_i}$.

2. Normalization: The outputs must sum to 1, which requires knowing the total sum of exponentials.

This suggests we need to track two running statistics:

• The maximum score seen so far (for numerical stability)
• The running sum of exponentials (for normalization)

But here's the subtlety: when we see a new block with a larger maximum, we need to retroactively adjust our previous contributions. The exponentials we computed earlier were relative to the old maximum, and now they need to be relative to the new maximum.

After processing block $j$, we maintain two statistics:

• $m^{(j)}$: the maximum value seen so far (over all blocks 1 through $j$)
• $L^{(j)}$: the sum of exponentials $\sum_{i=1}^{j} e^{x_i - m^{(j)}}$, where each exponential is computed relative to the current maximum

When block $j+1$ arrives with new attention scores $\{x_i\}$, we update both statistics in sequence. First, we update the running maximum:

$$m^{(j+1)} = \max(m^{(j)}, \max_i(x_i))$$

where:

• $m^{(j+1)}$: the updated maximum after seeing block $j+1$
• $m^{(j)}$: the previous maximum (from blocks 1 through $j$)
• $\max_i(x_i)$: the largest score in the new block

This simply takes the larger of the old maximum and the new block's maximum. The maximum can only increase or stay the same, never decrease.

The running sum $L^{(j)}$ was computed relative to $m^{(j)}$, but we need all exponentials relative to $m^{(j+1)}$. We must rescale the previous sum:

$$L^{(j+1)} = L^{(j)} \cdot e^{m^{(j)} - m^{(j+1)}} + \sum_i e^{x_i - m^{(j+1)}}$$

where:

• $L^{(j+1)}$: the updated sum of exponentials after block $j+1$
• $L^{(j)}$: the previous sum (computed with respect to $m^{(j)}$)
• $e^{m^{(j)} - m^{(j+1)}}$: the rescaling factor that adjusts the old sum for the new maximum
• $\sum_i e^{x_i - m^{(j+1)}}$: the contribution from the new block, computed relative to the new maximum

The rescaling factor $e^{m^{(j)} - m^{(j+1)}}$ deserves careful attention. It's the key that makes online softmax mathematically exact.

Consider what happens in two cases:

Case 1: The new block has a larger maximum ($m^{(j+1)} > m^{(j)}$)

The exponent $m^{(j)} - m^{(j+1)}$ is negative, so the factor is less than 1. This correctly shrinks the previous contributions. Why? Because we originally computed $e^{x_i - m^{(j)}}$, but now we need $e^{x_i - m^{(j+1)}}$. By the properties of exponents:

$$e^{x_i - m^{(j+1)}} = e^{x_i - m^{(j)}} \cdot e^{m^{(j)} - m^{(j+1)}}$$

Case 2: The new block doesn't change the maximum ($m^{(j+1)} = m^{(j)}$)

The exponent is zero, so the factor equals 1. The previous contributions remain unchanged, which is correct since no rescaling is needed.

This rescaling mechanism is what allows us to process blocks incrementally while guaranteeing the exact same result as computing softmax over all values at once.

Both softmax sums equal 1.0, and the individual values match exactly. This confirms that online softmax is not an approximation. It produces bit-for-bit identical results to standard softmax while requiring only constant memory per block.

Online softmax handles the normalization challenge, but attention requires something more: we don't just compute softmax weights, we use those weights to compute a weighted sum of value vectors. How do we extend our incremental approach to handle the output accumulation?

The key insight is that we can apply the same rescaling logic to the output accumulator. Just as we rescale the running sum $L$ when the maximum changes, we also rescale the running output $O$.

For each query position, we maintain three running statistics:

• $m_i$: the maximum attention score seen so far
• $L_i$: the running sum of exponentials (for normalization)
• $o_i$: the running output (unnormalized weighted sum of values)

When processing a new key/value block, the algorithm performs these steps in sequence:

1. Compute attention scores: Calculate $s_{ij}$ for each query $i$ against each key $j$ in this block
2. Update the maximum: Find $m_{\text{new}} = \max(m_{\text{old}}, \max_j(s_{ij}))$ for each query
3. Rescale previous contributions: Multiply the output accumulator by $e^{m_{\text{old}} - m_{\text{new}}}$
4. Add the current block's contribution: Compute $\sum_j e^{s_{ij} - m_{\text{new}}} v_j$ and add it to the output
5. Update the running sum: Apply the same rescaling to $L_i$ and add the new block's exponentials

After processing all key/value blocks, the final output for query $i$ is simply $o_i / L_i$, the unnormalized weighted sum divided by the normalization constant.

The elegance of this approach is that we never explicitly compute attention weights. Instead, we accumulate weighted outputs and divide at the end. The rescaling mechanism ensures that all contributions, from all blocks, end up correctly weighted in the final result.

Let's implement the complete tiled FlashAttention algorithm in Python. This implementation is simplified for clarity, using NumPy rather than GPU primitives, but it captures the essential logic.