Merge normalized attention without averaging tile outputs

Let scores be finite dimensionless numbers on the allowed positions of one row, with one value vector per score. Subtracting the row maximum before exponentiation keeps every exponential at most one, avoiding overflow from large positive scores. The numerator is a vector, while the denominator is a scalar. Dividing them gives a weighted average in the same units as the values. Storing only that average discards how much exponential mass supported it.

Two tiles can each produce a normalized output of ordinary magnitude while their total masses differ by orders of magnitude. An unweighted average of those outputs therefore assigns the wrong influence to their keys. Even a mass-weighted average needs masses expressed relative to the same maximum. This is why the maximum belongs in the summary: it is the reference level that makes the other two fields interpretable, rather than an optional diagnostic statistic.

For disjoint nonempty tiles, choose the larger maximum and multiply each mass and weighted numerator by the exponential of its old maximum minus the new maximum. Both rescaling factors lie between zero and one. Add the rescaled masses and numerators, then normalize once. Expanding the definitions shows this is the same sum that would have been formed over all allowed keys at once. In exact arithmetic, grouping or ordering these merges does not change the answer.

Floating-point execution changes the reduction order, so exact attention here means the same mathematical operator rather than bitwise identity. A tree reduction, sequential tile loop, and materialized softmax can round differently. Keep sufficiently accurate accumulation and compare errors against a higher-precision reference. Large score differences can underflow tiny contributions; the stable maximum subtraction addresses overflow but cannot promise that every real-number contribution survives a finite numerical format.

A fully masked tile contributes no keys. Represent it as an explicit empty summary with zero mass and zero numerator, and let merging an empty summary return the other summary unchanged. A sentinel maximum of negative infinity can be useful, but evaluating negative infinity minus negative infinity creates NaN. Branch around the both-empty case instead of assuming a later multiplication by zero will repair undefined intermediates. Masked logits are exclusions, not large finite negative guesses.

If the entire row is masked, the usual normalized softmax expression has zero denominator and no probability distribution. Specify an extension such as zero output, or reject the input, and make forward and backward implementations agree. FlashAttention 2.8.3 documents zero output for its all-masked causal rows; another API must be checked independently. This explicit policy also matters when a sparse pattern, ragged batch, or unequal query/key alignment creates empty rows unexpectedly.