Stable softmax numerics

One line of algebra

Softmax over a row is exp(x) / sum(exp(x)). Subtracting any constant from every element multiplies numerator and denominator by the same factor, so the output vector does not move. Choosing the row maximum as that constant puts every exponent at or below exp(0), which is 1, so no term can overflow and no division of infinities can occur.

That is the whole argument. Nothing about the distribution improves; the range of the exponent is what improves.

Where the range bites

f32 overflows at exp(88.7) and f16 at exp(11.1). A row of logits around 12 gives inf in half precision, and the naive form divides inf by inf, which is NaN. The same row through the max-subtracted form returns a normal distribution with the largest term at 1.

The other direction is quieter. A row near -1000 makes every exponent exp(0) after subtraction and returns a uniform row, while the naive form underflows every term to zero and divides zero by zero. Both cases belong in the generator, because a normal distribution never reaches either.

The loop shape

The maximum is a row reduction before any exponential, so the row is read twice unless it stays resident. At VLEN=256 an m4 group holds 32 f32 and the whole register file holds 256, so a 1024-column row is streamed: once for the max, then again for the exponential and the sum. Below about 512 columns the row fits beside its temporaries, and the loop becomes load, reduce, exponentiate in place, reduce, multiply by one reciprocal.

Things that break

A row that is entirely -inf, which is what a fully masked attention row is, gives max = -inf and then -inf - -inf is NaN before the exponential runs. Guard it and return zeros. The second failure is a maximum recomputed per tile without the terms from earlier tiles, which is right for a single-tile row and wrong for every longer one.