Online Softmax (the FlashAttention Trick) - LLM Internals

Online Softmax: the trick behind FlashAttention

What you'll build

Two functions that compute a softmax in one streaming pass over blocks: online_softmax for a 1-D vector, and online_attention, which computes exact attention block by block over the keys without ever forming the full score matrix.

Background

Standard attention materializes the (n_q, n_k) score matrix, then takes a row softmax. For long sequences that matrix doesn't fit in fast on-chip memory (SRAM), and moving it to and from slow memory (HBM) dominates the runtime. FlashAttention (Dao et al., 2022) walks over the keys in blocks and keeps three running statistics per query. It is exact, not an approximation; it saves memory traffic, not FLOPs. The enabling idea is the online softmax (Milakov & Gimelshein, 2018): when the running max increases, rescale everything accumulated so far by exp(m_old - m_new).

Formula: online_softmax(x, block_size)

Start with a running max m = -inf and running sum l = 0. For each block b of x (the last block may be shorter):

m_new = max(m, max(b))
l     = l * exp(m - m_new) + sum(exp(b - m_new))     # rescale the old sum to the new max
m     = m_new

After the pass: softmax(x) = exp(x - m) / l.

Formula: online_attention(Q, K, V, block_size)

Keep per query row a running max m (n_q, 1), a running denominator l (n_q, 1) and an unnormalized output acc (n_q, d_v). For each block (K_b, V_b) of block_size keys:

S     = Q @ K_b.T / sqrt(d)                # (n_q, B)
m_new = maximum(m, rowmax(S))              # (n_q, 1)
scale = exp(m - m_new)                     # shrink the old statistics
P     = exp(S - m_new)                     # (n_q, B)
l     = l * scale + rowsum(P)
acc   = acc * scale + P @ V_b              # (n_q, d_v)
m     = m_new

At the end, output = acc / l, which equals softmax(Q K^T / sqrt(d)) @ V up to float rounding.

Inputs

• x: 1-D float array; block_size: positive int
• Q (n_q, d), K (n_k, d), V (n_k, d_v); block_size keys per block

Expected Return Format

• online_softmax(x, block_size): 1-D array, same shape as x, sums to 1
• online_attention(Q, K, V, block_size): array (n_q, d_v)

Edge cases and conventions

• exp(-inf - m_new) = 0, so the first block needs no special case.
• Any block_size (1, a divisor of n_k, or larger than n_k) must give the same answer.
• Logits around 1000 must not overflow.

Constraints

NumPy only. Process the keys block by block; don't take a softmax over a whole row of scores at once.

Follow-up

• What is the memory footprint of the score matrix for n = 32k tokens and 32 heads in fp16? What does FlashAttention store instead?
• FlashAttention-2 delays the division by l to the very end and parallelises over query blocks. Why does that help?

Input: online_softmax(np.array([1000.0, 1001.0, 1002.0]), block_size=2)
Output: [0.09, 0.2447, 0.6652]

Input: online_softmax(np.array([0.0, 1.0, 5.0, 2.0, 9.0]), block_size=2)
Output: [0.0001, 0.0003, 0.018, 0.0009, 0.9807]

import numpy as np

def online_softmax(x, block_size):
    """
    Softmax of a 1-D array computed with one streaming pass over blocks
    (running max + running sum), then one normalization pass.

    Returns:
        probs: same shape as x, sums to 1
    """
    # Your code here
    pass


def online_attention(Q, K, V, block_size):
    """
    Attention softmax(Q K^T / sqrt(d)) V computed block by block over the keys,
    FlashAttention style.

    Args:
        Q: (n_q, d), K: (n_k, d), V: (n_k, d_v)
        block_size: number of keys per block

    Returns:
        output: (n_q, d_v)
    """
    # Your code here
    pass