softmax: accumulate the f16 row sum in f32

[czoli1976]

softmax_inner_slice_f16 summed the exponentials into an f16 accumulator. Once a row is long enough that the running sum outgrows its own terms, the additions round away entirely and the row gets normalised by too small a divisor — so the output does not sum to 1.

The regression test added here fails on main with row sums to 1.2903354, expected 1.0 at row length 4096: attention weights inflated by 29%.

How bad, and where

Relative error of the f16 accumulator against the same sum in f32, over normally-distributed logits:

row length	sd 0.5	sd 1.0	sd 2.0
1024	0.2%	1.2%	2.0%
2048	2.1%	3.8%	4.6%
4096	15.4%	10.1%	9.0%
8192	38.2%	28.3%	14.2%

The mechanism is easiest to see in the flat case: with all logits equal, every term after the max subtraction is 1.0, so the sum should be row_len. f16 spacing at 2048 is 2.0, so from there on sum + 1.0 == sum and the total sticks at 2048 — 50% low for a 4096-long row.

This is the default path rather than an opt-in one. SoftmaxExp::default() is Libc, and ScaledMaskedSoftmax::eval constructs SoftmaxExp::Libc directly, so every f16 attention softmax reaches it. Error grows with sequence length, which is the direction long-context inference is going.

The fix

Accumulate in f32 and narrow once at the end, for both the softmax and log-softmax branches. Per-element values are untouched — the exponentials are still computed exactly as before and still stored as f16 — so the only thing that changes is the summation, and short rows are unaffected.

It should also be marginally cheaper: f16 + f16 widens to f32 internally in the half crate anyway, so this drops a narrowing per element.

Known, not fixed here

FastCompact has the same flaw one level down: HSoftMaxL2 in linalg/src/generic/reduce.rs accumulates with let mut sum = f16::zero() and reduces with reduce_two(a: f16, b: f16). Fixing that means changing the map-reduce accumulator type for f16, which touches the frame and the AVX-512 kernel alongside the generic one — a bigger change than this, and FastCompact is opt-in. Happy to follow up if you want it in the same series. This PR routes the FastCompact result through f32 for the reciprocal, but cannot repair a sum that was already lost inside the kernel.

Tests

tract-core 271, test-f16 2378, test-unit-core 816, tract-linalg 4414 — green. cargo fmt --all clean; cargo clippy clean (the two remaining warnings are pre-existing on main).

🍍