Online Softmax & FlashAttention · AI infra interview lesson

Why It Matters

Full attention produces N×N scores. Can small tiles produce the same result? Maintain a running maximum and normalization sum, then apply the same rescaling to the output.

The numbers: with N = 8192, the fp16 score matrix of a single head is 8192 * 8192 * 2 B = 128 MiB, and 32 heads make 4 GiB. That is S alone, for one layer and one sequence, and P needs another copy. FlashAttention only keeps m and l per row (or a single LSE), which is N fp32 values, 8192 * 32 * 4 B = 1 MiB across 32 heads. HBM accesses drop from Θ(N*d + N^2) for the standard implementation to Θ(N^2 * d^2 / M), where M is the SRAM size.

Walk the Shapes First

Q: [N=6, d=4], K, V: [Nk=6, d=4], Br = 4, Bc = 3:

Rules read from the table: a row needs just two numbers (m, l) to merge softmax across blocks, and the merge rescales the old statistics by the new max; when both sides are -inf the merge must be special-cased, otherwise NaN propagates all the way to the output.

Running the examples: execute this page in order in one Python environment. Later snippets may reuse earlier definitions. CPU examples validate calculations and state changes, not GPU timing. Download the notebook for the complete exercise workspace.

Core Implementations

Step 1: online softmax (m, l) and merging

New idea: the softmax of a row x depends only on m = max(x) and l = sum(exp(x - m)). Statistics of two segments merge as:

• m = max(m1, m2)
• l = exp(m1 - m)*l1 + exp(m2 - m)*l2

The merge is associative and commutative, so any chunking and any merge order works. That is what lets FlashAttention and Flash-Decoding run in parallel. A single-pass scan just merges in each element as (x_i, 1).

import torch, math
torch.manual_seed(0)
NEG_INF = float("-inf")

def online_softmax_stats(x):
    # single pass over a 1D tensor: running max m and running sum l
    m, l = NEG_INF, 0.0
    for xi in x.tolist():
        m_new = max(m, xi)
        if m_new == NEG_INF:              # nothing finite seen yet: skip, avoids -inf - -inf
            continue
        l = l * math.exp(m - m_new) + math.exp(xi - m_new)
        m = m_new
    return m, l

def merge_ml(m1, l1, m2, l2):
    m = torch.maximum(m1, m2)
    safe_m = torch.where(m == NEG_INF, torch.zeros_like(m), m)   # both sides -inf -> use 0
    return m, l1 * torch.exp(m1 - safe_m) + l2 * torch.exp(m2 - safe_m)

x = torch.tensor([1.0, 3.0, 2.0, 5.0])
m, l = online_softmax_stats(x)
assert m == 5.0 and abs(l - 1.2034379904932107) < 1e-6
assert torch.allclose(torch.exp(x - m) / l, torch.softmax(x, 0))

# arbitrary chunking + left fold == tree merge == direct computation
X = torch.randn(8, 37) * 10
X[1, :20] = NEG_INF                                   # partially masked row
X[2, :] = NEG_INF                                     # fully masked row
def chunk_stats(c):
    m = c.max(-1).values
    safe = torch.where(m == NEG_INF, torch.zeros_like(m), m)
    return m, torch.exp(c - safe[:, None]).sum(-1)
chunks = [chunk_stats(c) for c in X.split([5, 1, 13, 8, 10], dim=1)]
m_f, l_f = chunks[0]
for mc, lc in chunks[1:]:
    m_f, l_f = merge_ml(m_f, l_f, mc, lc)
left = merge_ml(*merge_ml(*chunks[0], *chunks[1]), *merge_ml(*chunks[2], *chunks[3]))
m_t, l_t = merge_ml(*left, *chunks[4])
ref_lse = torch.logsumexp(X, -1)
lse_f = torch.where(l_f > 0, m_f + torch.log(l_f), torch.full_like(m_f, NEG_INF))
assert torch.allclose(lse_f, ref_lse) and torch.allclose(m_t, m_f) and torch.allclose(l_t, l_f)
assert l_f[2] == 0 and m_f[2] == NEG_INF and not torch.isnan(l_f).any()

# the naive merge produces NaN on a fully masked row
m_bad = torch.tensor(NEG_INF)
naive_l = 0.0 * torch.exp(m_bad - torch.maximum(m_bad, m_bad))
assert torch.isnan(naive_l)

Step 2: FlashAttention V1 (KV outer, Q inner)

New idea: extend the merge in Step 1 from a scalar l to an output vector O. O is already normalized, so the merge multiplies back the old l, adds the new block's contribution, and divides by the new l:

• O_new = (l_old * exp(m_old - m_new) * O_old + exp(m_ij - m_new) * (P_ij @ V_j)) / l_new

V1 loops over KV blocks outside and Q blocks inside: once a KV block is in SRAM, it sweeps every Q block, and O, m, l are read from HBM, updated and written back each time.

import torch, math
import torch.nn.functional as F
torch.manual_seed(0)
NEG_INF = float("-inf")

def merge_ml(m1, l1, m2, l2):
    m = torch.maximum(m1, m2)
    safe_m = torch.where(m == NEG_INF, torch.zeros_like(m), m)
    return m, l1 * torch.exp(m1 - safe_m) + l2 * torch.exp(m2 - safe_m)

def flash_attn_v1(q, k, v, Br=4, Bc=3):
    N, d = q.shape
    Nk = k.shape[0]
    scale = 1 / math.sqrt(d)
    O = torch.zeros(N, d)                             # these three live in HBM
    l = torch.zeros(N)
    m = torch.full((N,), NEG_INF)
    for j0 in range(0, Nk, Bc):                       # outer: load one KV block
        Kj, Vj = k[j0:j0 + Bc], v[j0:j0 + Bc]
        for i0 in range(0, N, Br):                    # inner: sweep every Q block
            S = q[i0:i0 + Br] @ Kj.T * scale          # [Br, Bc]
            m_ij = S.max(-1).values
            P = torch.exp(S - m_ij[:, None])
            l_ij = P.sum(-1)
            m_old, l_old = m[i0:i0 + Br], l[i0:i0 + Br]
            m_new, l_new = merge_ml(m_old, l_old, m_ij, l_ij)
            O[i0:i0 + Br] = ((l_old * torch.exp(m_old - m_new))[:, None] * O[i0:i0 + Br]
                             + torch.exp(m_ij - m_new)[:, None] * (P @ Vj)) / l_new[:, None]
            m[i0:i0 + Br], l[i0:i0 + Br] = m_new, l_new
    return O

for N, Nk, d in [(6, 6, 4), (10, 7, 8), (1, 1, 2), (5, 13, 4)]:
    q, k, v = torch.randn(N, d), torch.randn(Nk, d), torch.randn(Nk, d)
    ref = F.scaled_dot_product_attention(q[None], k[None], v[None])[0]
    assert torch.allclose(flash_attn_v1(q, k, v), ref, atol=1e-5)

The cost of V1: the inner loop reads and writes O every time and divides by l_new for every block. V2 swaps the loops and defers the division to the end.

Step 3: FlashAttention V2 (Q outer, delayed rescale, returns LSE)

New idea: the outer loop is over Q blocks (one Q block per thread block, fully independent of each other) and the inner loop is over KV. (Pro lesson)

Step 4: causal block classes and fully masked rows

New idea: under causal attention a (Q block, KV block) pair falls into three classes: entirely above the diagonal, skipped (skip); (Pro lesson)

Apply causal visibility to scores

Classify each tile as fully visible, partly visible or invisible. Mask QK scores for partly visible tiles. (Pro lesson)

Step 5: Flash-Decoding (split-KV + LSE merge)

New idea: in decode the query is a single token, so splitting over Q gives one block and not enough parallelism on a GPU. (Pro lesson)

In Real Systems

FlashAttention library (Dao-AILab/flash-attention) : flash_attn_func , flash_attn_varlen_func and flash_attn_with_kvcache are the usual entr… (Pro lesson)

Common Pitfalls

Both sides are -inf in the online softmax merge → exp(m1 - m) with both m1 and m equal to -inf (a thread got 0 elements, or a whole block is… (Pro lesson)

Say It in the Interview

Softmax only needs two running statistics per row, the max m and the rescaled sum l, and two partial states merge associatively with m = max… (Pro lesson)

Hands-on Notebook

The notebook has 7 exercises, from online softmax all the way to Flash-Decoding. (Pro lesson)