Softmax on f16/bf16 is numerically broken on long axes (rows sum to ~2.0, not 1.0)

Severity

Silent correctness bug (wrong numerical output, no error/panic) on the attention hot path. Affects any f16/bf16 model doing softmax over a long axis (e.g. attention over long sequences). The CPU result also disagrees with the Metal backend.

Summary

candle_nn::ops::softmax_last_dim runs the whole computation in the tensor's dtype. For f16/bf16 the exp-denominator is summed sequentially in the narrow dtype, which saturates on a long axis: an f16 running sum of ones stalls at 2048 (once the sum's ULP exceeds 1, each +1 rounds back). So the row is normalized by 2048 instead of its true length, and the "probabilities" no longer sum to 1.

Repro (CPU, no GPU needed)

use candle_core::{Device, Tensor, DType};
use candle_nn::ops::softmax_last_dim;

fn main() -> candle_core::Result<()> {
    // 4096 equal logits in f16 => every exp = 1, denominator should be 4096.
    let t = Tensor::zeros((1, 4096), DType::F16, &Device::Cpu)?;
    let sm = softmax_last_dim(&t)?;
    // Measure in f32 so the check itself doesn't saturate.
    let row_sum = sm.to_dtype(DType::F32)?.sum_all()?.to_scalar::<f32>()?;
    let first  = sm.to_dtype(DType::F32)?.flatten_all()?.get(0)?.to_scalar::<f32>()?;
    println!("row_sum = {row_sum}   first_prob = {first}");
    Ok(())
}

Expected vs actual

Expected: row_sum = 1.0        first_prob = 1/4096 = 0.000244
Actual:   row_sum = 2.0        first_prob = 1/2048 = 0.000488

Every probability is ~2x too large; the softmax is not normalized. bf16 is worse (8-bit mantissa saturates even sooner). f32/f64 are fine.

Root cause

The f16/bf16 path sums the exponentials in the narrow dtype. The Metal kernel already keeps the softmax denominator in f32 (MD::d is float), so the backends disagree.

Fix

Compute the f16/bf16 path in f32 (max, exp, denominator, normalization) and cast back, matching the Metal kernel; f32/f64 paths unchanged. I've opened #3717 with this fix + regression tests. After it, an f16 softmax over a length-4096 axis matches the Metal backend to max-abs 1.2e-7 (both rows sum to 1.0).

Environment

candle main; reproduces on CPU (Apple M4, macOS), backend-independent.

Start at candle_nn::ops::softmax_last_dim and run the CPU reproduction for a length-4096 f16 axis. Review the regression tests and fix described in pull request #3717, then verify that the CPU result agrees with the Metal backend and that the row sum is 1.0.