CPU 上 BF16 softmax 加起来不是 1

主线位置：算子主线在 CPU 后端的最底层，cpu/kernels.rs 里按 dtype 实现的 VecOps。softmax_last_dim 是 candle-nn 的融合算子（常用层），sum / sum_keepdim 是 candle-core 的归约（CPU 后端），两者在 CPU 上都落到同一个 vec_reduce_sum。

flowchart TB
  x["BF16 张量 一行 2048 个"] --> dev{"设备"}
  dev -->|CPU| acc["bf16 累加器"]
  acc --> stuck["加到 256 附近不再变"]
  stuck --> bad["分母偏小 行和大于 1"]
  dev -->|Metal| f32acc["F32 累加"]
  f32acc --> ok["行和为 1"]
  x --> fix["先 to_dtype F32"]
  fix --> ok

实验 crate exp-candle/pb-pit，bin bf16sm，release，M2 Max、macOS 27、rustc 1.99.0。

现象

4 行注意力分数，每行 2048 个键，数值用 randn(0, 0.5) 模拟（缩放后的 logits 通常就是这个量级），分别转成 F32、F16、BF16 做 softmax_last_dim，再把每行的概率转回 F32 加起来：

let scores = Tensor::randn(0f32, 0.5, (4, 2048), &Device::Cpu)?.to_device(&dev)?;
for dt in [DType::F32, DType::F16, DType::BF16] {
  let x = scores.to_dtype(dt)?;
  let p = ops::softmax_last_dim(&x)?;
  // row sums, computed in F32
}
let ones = Tensor::ones(4096, DType::BF16, &dev)?;
println!("4096 BF16 ones: sum_all = {}", ones.sum_all()?.to_dtype(DType::F32)?.to_scalar::<f32>()?);

[cpu] F32 softmax_last_dim row sums 1.0000 1.0000 1.0000 1.0000
[cpu] F16 softmax_last_dim row sums 0.9964 0.9818 0.9743 0.9871
[cpu] BF16 softmax_last_dim row sums 2.6053 2.7168 2.7638 2.5642
[cpu] 4096 BF16 ones: sum_all = 256
[metal] F32 softmax_last_dim row sums 1.0000 1.0000 1.0000 1.0000
[metal] F16 softmax_last_dim row sums 1.0000 1.0000 1.0001 1.0001
[metal] BF16 softmax_last_dim row sums 1.0006 1.0002 0.9997 0.9993
[metal] 4096 BF16 ones: sum_all = 4096

（randn 没有设种子，每次运行的数字略有不同，量级不变。）

• CPU 上 BF16 每行概率加起来是 2.6 到 2.8：分母只有真实值的三分之一多一点，每个概率都被放大了。注意力的输出是这些概率对 v 的加权和，整体被放大 2.7 倍，不报任何错。
• CPU 上 F16 偏小 1% 到 2.6%。
• 和 softmax 无关的 sum_all：4096 个 BF16 的 1 加起来是 256。
• Metal 上三种 dtype 都对（BF16 的误差在第四位，是结果存回 BF16 的舍入）。

出处

CPU 的融合 softmax（SoftmaxLastDim 的 cpu_fwd）对每一行先求最大值、逐个 exp，再：

                    let mut sum_exp = T::zero();
                    unsafe { T::vec_reduce_sum(dst.as_ptr(), &mut sum_exp, dim_m1) };

T 是 half::bf16 时，vec_reduce_sum 用的是 VecOps 的 默认实现：

    unsafe fn vec_reduce_sum(xs: *const Self, res: *mut Self, len: usize) {
        *res = Self::zero();
        for i in 0..len {
            *res += *xs.add(i)
        }
    }

impl VecOps for f32 覆盖了这个方法；half::f16 和 half::bf16 只覆盖了 vec_dot 等几个，没有覆盖 vec_reduce_sum，所以累加器 *res 本身就是 BF16。BF16 只有 8 位有效位：累加器在 64 到 128 之间时相邻可表示数相差 0.5，再加一个小于 0.25 的数就被舍掉；到 256 时连加 1 也舍回 256（相邻可表示数相差 2）。上面每行先减去最大值再 exp，平均每项约 0.2，真实的和是四五百，累加器却在一两百处停住，比值正好是 2.6 左右（用 Python 按 BF16 舍入逐项模拟同样的数据，得到 2.25 到 2.75）。F16 有 11 位有效位，停得晚，误差小得多。sum / sum_all 在 CPU 上走 ReduceSum，同样调 T::vec_reduce_sum，所以 4096 个 1 的和是 256。

Metal 的 softmax 内核用 float 累加（常用层 里有它的在线 softmax），求和内核的累加器也比 BF16 宽，所以没有这个问题。拼出来的 ops::softmax 在 CPU 上同样经过 sum_keepdim，一样错；常用层 的动脑筋用全零输入把 256、257、4096 三种长度都量过了。

这是 0.11.0 的上游 bug：同一个算子在 CPU 和 Metal 上给出不同的数学结果，CPU 那一边是错的。

修法

在 CPU 上做 softmax 和求和之前转成 F32，算完再转回去：

let p = ops::softmax_last_dim(&x.to_dtype(DType::F32)?)?.to_dtype(dt)?;
let s = ones.to_dtype(DType::F32)?.sum_all()?;

[cpu] F32 softmax_last_dim row sums 1.0000 1.0000 1.0000 1.0000
[cpu] F16 softmax_last_dim row sums 1.0000 1.0000 1.0000 1.0000
[cpu] BF16 softmax_last_dim row sums 0.9999 1.0000 0.9999 1.0000
[cpu] 4096 BF16 ones: sum_all = 4096
[metal] F32 softmax_last_dim row sums 1.0000 1.0000 1.0000 1.0000
[metal] F16 softmax_last_dim row sums 1.0000 1.0000 1.0000 1.0000
[metal] BF16 softmax_last_dim row sums 1.0000 0.9999 1.0000 0.9999
[metal] 4096 BF16 ones: sum_all = 4096

CPU 后端本来就不支持 BF16 的 matmul（DType 与 Device），整个 BF16 模型在 CPU 上跑不起来，所以这个坑出现在两种地方：BF16 张量上单独做的归约和 softmax（统计、损失、取样前的处理），以及 F16 模型（不开 accelerate 时 CPU 的 F16 matmul 可用，softmax 偏 1% 到 3%）。candle-transformers 里的模型做法不一：llama.rs 的注意力先把 q、k、v 转成 F32 再乘、再 softmax，最后转回输入 dtype；mistral.rs 直接在输入 dtype 上 softmax_last_dim。CPU 上最省事的是整个模型用 F32；要用半精度时，自己写的层照 llama.rs 的做法。

怎么确认没中招

• CPU 上凡是 BF16 / F16 张量的 sum、sum_keepdim、sum_all、mean（除以个数前先求和）、softmax_last_dim、ops::softmax，归约长度超过一两百时都可疑。在代码里搜这些调用，看输入的 dtype。
• 直接测：对 softmax 的输出转 F32 再按最后一维求和，和 1 比；或者对一个 BF16 的全 1 张量 sum_all，结果应当等于元素个数。
• 同一个 BF16 模型在 CPU 和 Metal 上输出差别很大时，先怀疑这一条。

小结

• CPU 后端的 vec_reduce_sum 对 BF16 / F16 用半精度累加器；BF16 加到 256 左右就加不上去了。
• BF16 的 softmax_last_dim 一行 2048 个时每行和约 2.7，sum_all 4096 个 1 得 256；F16 偏差约 1% 到 3%。Metal 用 F32 累加，结果正确。
• CPU 上归约之前转 F32；CPU 上的模型最好整个用 F32。