The block-size below which a flash-attention kernel stops being accurate, and an audit showing the shipped Apple M4 kernel stays on the safe side of it.
Flash attention never materializes the attention matrix. It walks the keys in
blocks of size Bc and carries the softmax denominator l and the output o as
per-row scalars, rescaling them by the running row max each block. That carry is
a recursive scalar sum, and recursive summation in low precision stagnates.
In a format with t stored mantissa bits, the running sum l_k = fl(l_{k-1} + a_k)
drops the term a_k entirely whenever
a_k / l_{k-1} < 2^-(t+1) .
A softmax row of post-max-subtracted terms has a typical magnitude mu, so the
partial sum after k terms is about k*mu and every term past index 2^(t+1)
falls below the drop threshold and is lost. The accumulator saturates and
underestimates the denominator, with a relative error that climbs toward 1 as the
row length n grows. The drop index is 256 for bfloat16 (t = 7), 2048 for
fp16, and 16.7M for fp32. Measured, a sequential bf16 sum of n positive terms
has relative error 1.3% at n = 64, 50% at 1024, and 97% at 16384,
while fp32 stays at 1e-7.
The failure is not a property of bf16. The same bf16 terms summed by a pairwise
(tree) reduction keep full accuracy, because no single accumulator absorbs all the
mass: at n = 16384, pairwise bf16 has 6e-4 relative error while sequential
bf16 has 97%, a factor of 1510. Flash attention cannot use a tree reduction
across key blocks, because l and o are per-row scalars updated once per block,
so its cross-block reduction is forced sequential in the number of blocks n/Bc.
Putting the two together: a flash kernel that accumulates its carry in a format
with t mantissa bits is accurate while
n / Bc < 2^(t+1) ,
and cliffs beyond it. Measured on a bf16 carry at n = 8192: relative denominator
error is flat near 0.3% for n/Bc up to about 256, then rises to 2.3% at
n/Bc = 256, 38% at 1024, and 73% at 8192 (one element per block). The
knob a kernel author reaches for to save registers and SRAM, a smaller carry
precision or a smaller block, walks straight into this cliff once n/Bc passes
the format's drop index.
Do the attention kernels people actually run avoid this? Torch's
scaled_dot_product_attention, including the fused Metal path on the M4 GPU, has
a bf16 error that is flat across sequence length: 5.85e-3 at n = 256 and
6.16e-3 at n = 8192, versus a swamping ceiling above 70% for a naive
bf16-carry flash kernel at the same length. Flat error means the kernel keeps its
softmax carry and its P @ V accumulation in fp32. The measured flat floor is
exactly what that fp32 carry buys over the naive ceiling, and it is what a
low-precision-accumulate variant would give up.
stagnation.py: the drop-index lemma and the measured swamping curve by format.carry.py: sequential versus pairwise on identical terms, and then/Bc < 2^(t+1)block-size envelope.audit.py: thescaled_dot_product_attentionaudit on the M4 GPU, against the naive-carry ceiling.test_flashswamp.py: the lemma, the envelope, and the audit as tests.
python stagnation.py
python carry.py
python audit.py
python test_flashswamp.py
Keep the online softmax carry (l and o) in fp32 regardless of the matmul
precision, or keep n/Bc below the accumulator format's 2^(t+1) drop index.
The error is set by the number of sequential block updates, not by the precision
of the scores.