Online Softmax: Tiling Across a Reduction

Fusing an elementwise chain is easy: each output element depends on one input element, so a kernel can process any block independently.

Softmax is not elementwise. For a row of scores, each output is the exponential of that score divided by the sum of exponentials across the entire row. Every output depends on every input in its row.

That is what blocks naive tiling. Load a block of keys, compute the scores for that block, and you cannot normalise them, because the denominator involves keys you have not looked at yet. The obvious conclusion is that you must compute all the scores first, which means materialising the matrix, which is exactly what you were trying to avoid.

There is a second problem underneath. Softmax is implemented by subtracting the row maximum before exponentiating, because raw scores can be large and the exponential overflows in 16-bit at surprisingly modest values. The maximum is also a property of the whole row.

So tiling requires computing both a sum and a maximum over data that has not arrived. That looks impossible, and it is not.

The escape is an algebraic fact about the exponential, and it is the entire mathematical content of the algorithm.

Suppose you have computed a partial sum of exponentials using some running maximum, and a new block arrives with a larger maximum. You do not need to revisit the earlier data. Because the exponential turns subtraction into division, an exponential computed against the old maximum can be converted to one computed against the new maximum by multiplying by the exponential of the difference between them.

So the correction is a single scalar multiply per row.

Concretely, maintain two running quantities per row: the maximum seen so far, and the sum of exponentials relative to that maximum. When a block arrives, compute its own local maximum and local sum. The new global maximum is the larger of the two. Rescale the old running sum by the exponential of the old maximum minus the new one, rescale the block's sum the same way, and add them.

That is the whole update. Two scalars per row, corrected in constant time as each block arrives.

The statistics are only half of it. The output accumulator needs the same treatment, and this is the step that makes the algorithm work end to end.

As blocks of keys and values are processed, each contributes a weighted sum of value vectors, with weights from that block's unnormalised exponentials. Those contributions accumulate into an output vector per query.

When the running maximum changes, every previously accumulated contribution was computed against the old maximum and is now on the wrong scale. But the fix is the same scalar: multiply the accumulated output by the exponential of the old maximum minus the new one, then add the new block's contribution.

So the loop over key blocks maintains three things per query row: the running maximum, the running sum of exponentials, and the running output accumulator. Each is corrected by the same scalar whenever the maximum moves.

At the end, divide the accumulated output by the final sum, and the result is exactly what full softmax attention would have produced.

The score matrix was never assembled. Only its two summary statistics and a running output ever existed, and all three are small.

The algorithm is two nested loops, and knowing which is outer explains most of the implementation differences between versions.

The outer loop walks blocks of queries. Each iteration handles a set of query rows and produces their final outputs before moving on, so a query block's statistics and accumulator live in fast memory for the whole of its inner loop.

The inner loop walks blocks of keys and values. For each, it loads them into SRAM, computes the scores for this query block against this key block, finds the block's local maximum, updates the running maximum, rescales the running sum and the running accumulator, and adds this block's contribution.

When the inner loop finishes, the query block's output is normalised by the final running sum and written to HBM once.

Total HBM traffic for the score matrix is zero, because it was never a whole object. Q, K and V are read, the output is written, and the two per-row statistics are saved for the backward pass.

The original version put queries in the inner loop; FlashAttention-2 swapped them, which is what most of its speedup came from.

Block size is the parameter that determines whether the algorithm works at all, and it is set by a hard physical constraint.

Everything resident at once must fit in one streaming multiprocessor's SRAM: the query block, the key block, the value block, the score tile for that block pair, the running statistics and the output accumulator. On an A100 that is roughly 192 kilobytes.

That gives blocks on the order of 64 or 128 rows for typical head dimensions, though the exact figure depends on head dimension, data type and how much the compiler spills to registers.

The tension is in both directions. Larger blocks mean fewer inner-loop iterations and fewer rescaling operations, so less overhead. Too large and the working set spills out of SRAM, at which point the algorithm degenerates into the memory-bound version it was designed to replace.

Smaller blocks always fit and pay per-block overhead more often, and below a certain size the matrix multiplications become too small to use tensor cores efficiently.

This is why implementations carry hand-tuned block sizes per architecture and per head dimension, and why an unusual head dimension can fall off the fast path without any error being raised.

The forward pass is only half the training cost, and the backward pass has its own version of the same difficulty.

Computing gradients through attention requires the softmax output, the n-by-n matrix of normalised weights. Standard implementations save it during the forward pass and read it back during the backward pass, which is precisely the storage the forward algorithm just eliminated.

Saving it would defeat the purpose. Memory would be quadratic again, and the traffic would return.

The resolution is recomputation, the same technique the distributed training path used for activations. Do not save the matrix. Save the two per-row statistics, the final maximum and the final sum, which together occupy a few bytes per row rather than n entries.

During the backward pass, recompute the scores block by block from Q and K, which are already stored, and normalise them using the saved statistics. Because the final maximum and sum are known, no running update is needed on the way back: each block can be normalised correctly in one pass.

So the backward pass performs extra arithmetic and reads almost nothing. On a memory-bound kernel that is a straightforward win.

Putting figures on what the statistics replace makes the trade concrete.

At 8,192 tokens, one attention head's score matrix in 16-bit is 134.2 MB. The two saved statistics for the same head are two 32-bit values per row, which is 65.5 kilobytes.

Scale to a realistic training step: batch 8, 32 heads, 8,192 tokens. The score matrices are 34.4 gigabytes. The statistics are 16.8 megabytes, a factor of roughly two thousand smaller.

The scaling is the more important part. Standard attention's intermediate storage grows with the square of sequence length; the statistics grow linearly. At 2,048 tokens per head the comparison is 8.4 MB against 0.02 MB. At 32,768 it is 2.1 gigabytes against 0.26 megabytes.

This is what made long-context training possible. Not the speedup, which is welcome, but the removal of a quadratic memory term that made 32k-token training impossible on any hardware at any batch size.

And the arithmetic paid for it is modest: one extra pass over the scores in the backward direction, on a kernel that was waiting on memory anyway.

A fair question at this point is whether all that rescaling introduces error, and the answer needs care.

Mathematically the algorithm is exact. The rescaling identity is an equality, not an approximation, so in exact arithmetic the output is identical to standard attention.

In floating point the operations are performed in a different order and with different intermediate roundings, so individual values can differ in their last bits. That is reassociation, the same thing that makes summing an array in a different order give a slightly different result, and it is not error accumulation in any meaningful sense.

In one respect the algorithm is better behaved. Because it tracks a running maximum and rescales, the exponentials computed at each step are bounded by construction, so the classic overflow that motivated max-subtraction in the first place cannot occur.

The accumulators are kept in 32-bit even when inputs are 16-bit, which is standard practice and keeps the running sum well-conditioned over many blocks.

The practical upshot is that a model trained with a standard implementation and one trained with a fused kernel produce the same results within normal floating point tolerance, which is why swapping them requires no revalidation.

Language models use causal attention: each position attends only to itself and earlier positions. In a materialised implementation that is done by adding negative infinity to the upper triangle before the softmax, which costs a full extra pass over the matrix.

In a tiled implementation it becomes an optimisation instead.

Blocks are indexed by query range and key range. If a key block lies entirely after a query block, every entry is masked, and the block can be skipped without computing anything. If a key block lies entirely before, no masking is needed and the block is processed with no mask logic at all. Only blocks straddling the diagonal require element-level masking.

Since roughly half the block pairs fall in the skip category, causal attention runs at close to half the cost of full attention, which is exactly what the mathematics implies but is not what a masked dense implementation delivers.

The same reasoning extends to structured sparsity generally. Sliding windows, block-sparse patterns and document-boundary masks all become block-level skip decisions, which is why fused kernels and structured masking compose so well.

The specific algorithm matters less than the shape of the argument, which recurs whenever a reduction blocks fusion.

The reduction looked like it required all the data at once. It did not. It required a small summary of the data seen so far, plus a rule for correcting that summary when new data arrives.

For softmax the summary is a maximum and a sum, and the correction is one multiply. The same structure appears elsewhere: computing a mean and variance in one pass rather than two, which is what fused normalisation layers do; maintaining running statistics for a loss; any associative reduction with a correction term.

The general question to ask of a reduction is whether it admits an incremental form. If a partial result can be updated rather than recomputed, the operation can be fused into whatever produced its input, and a round trip to main memory disappears.

That is a question worth asking of your own code. The chain of tensor operations with a reduction in the middle is an extremely common pattern, and the reduction is usually the reason nobody fused it.

The next lesson is about writing such kernels without dropping to raw CUDA.