FlashAttention saves the matrix, not the quadratic work

A transformer compares every token with every other token. Standard implementations write the resulting square table of attention scores to high-bandwidth GPU memory, or HBM, then read it back for softmax and the weighted sum. On long sequences, moving that table can cost more time than doing the arithmetic.

FlashAttention changes the schedule, not the answer. It loads small blocks into fast on-chip memory, updates a running softmax, and never materializes the full score matrix. That made exact attention both faster and much less memory-hungry in Dao and colleagues' 2022 paper.

Softmax normally looks global: subtract the largest score, exponentiate every score, then divide by their sum. The trick is to carry three sufficient statistics between tiles—the largest score seen so far, the rescaled exponential sum, and the rescaled weighted-value accumulator. When a later tile contains a larger score, one exponential factor rescales everything already accumulated.

def merge_block(scores, values, m, total, acc):
    block_max = max(scores)
    new_max = max(m, block_max)
    old_scale = 0 if m == -inf else exp(m - new_max)
    total *= old_scale
    acc = [x * old_scale for x in acc]
    for score, value in zip(scores, values):
        weight = exp(score - new_max)
        total += weight
        acc = [a + weight * v for a, v in zip(acc, value)]
    return new_max, total, acc

The charts and equivalence check run the tested TypeScript core. Python and C++ above are faithful translations of the same online update. The site version loops over one query at a time for clarity; the paper's CUDA kernel tiles both query and key dimensions and fuses the operations.

With 4,096 tokens, one attention head has 16,777,216 pairwise scores. A 64-by-64 tile holds 4,096 scores—4,096 times fewer score slots at once. This is an analytical count of intermediate score elements, not a claim about total GPU bytes: queries, keys, values, outputs, statistics, and kernel bookkeeping still occupy memory.

Here is the skim-missed limit. Exact dense attention still evaluates every query-key pair. Counting the two matrix products, one head needs roughly 2N²d multiply-adds for sequence length Nand head width d. Doubling the context therefore quadruples this work even though the score workspace is tiled.

The speedups below are the authors' reported end-to-end results on their hardware and software stack, not outputs of our toy model. They reflect kernel fusion and fewer HBM transfers as well as the online-softmax math.

The paper estimates up to nine times fewer HBM accesses for typical head dimensions and on-chip memory sizes. It also states the engineering boundary plainly: the method requires custom low-level CUDA kernels, and performance may not transfer unchanged across GPU architectures.

• Measure wall-clock time while varying tile size and head width.
• Separate kernel-launch, arithmetic, and HBM time with a profiler.
• Compare exact and block-sparse attention at matched output error.

To see the same query-key-value mechanism without the GPU scheduling, poke at the interactive transformer page.