unsupported dtype BF16 for op matmul

主线位置：算子主线的后端一步。Storage::matmul 只检查两边 dtype 一致，不检查后端有没有实现这个 dtype；CpuStorage::matmul 按 CpuStorage 的变体分派，MetalStorage::matmul 按 dtype 选 GEMM 内核，没实现的组合在这里才失败。全表见 DType 与 Device。

flowchart TB
  op["matmul 或 exp"] --> same{"两边 dtype 一致"}
  same -->|是| be{"后端有实现"}
  be -->|CPU 没有| e1["unsupported dtype"]
  be -->|CPU 整数一元| p["todo! panic"]
  be -->|Metal 没有内核| e2["Error while loading function"]
  be -->|有| ok["计算"]
  fix["CPU 上 to_dtype 再算"] --> ok

实验 crate exp-candle/pa-pit（features metal accelerate）和 exp-candle/pa-noacc（同样的代码，只开 metal），bin p03_dtype，debug；M2 Max、macOS 27、rustc 1.99.0。

现象

从 safetensors 读出 BF16 权重（SmolLM2 的 model.safetensors 就是 BF16），在 CPU 上直接乘；再把同一个 (4, 4) 矩阵换成几种 dtype 在两个设备上试：

cpu BF16 matmul: Err(unsupported dtype BF16 for op matmul)
cpu F16 matmul: Err(the accelerate backend does not support f16 matmul)
cpu F32 matmul: Ok [4, 4] F32
metal BF16 matmul: Ok [4, 4] BF16
metal F16 matmul: Ok [4, 4] F16
metal F32 matmul: Ok [4, 4] F32
metal F64 matmul: Err(Metal error mlx matmul doesn't support F64)
metal I32 matmul: Err(Metal error mlx matmul doesn't support I32)
metal I32 add: Err(Metal error Error while loading function: badd_i32)
metal F64 add: Err(Metal error Error while loading function: badd_f64)
metal F64 to_dtype(F32): Err(Metal contiguous to_dtype F64 F32 not implemented)

thread 'main' (3927703) panicked at ~/.cargo/registry/src/index.crates.io-1949cf8c6b5b557f/candle-core-0.11.0/src/op.rs:584:1:
not yet implemented: no unary function for u32
note: run with `RUST_BACKTRACE=1` environment variable to display a backtrace
cpu U32 exp: PANIC (message above)

不开 accelerate（pa-noacc）时前两行变成：

cpu BF16 matmul: Err(unsupported dtype BF16 for op matmul)
cpu F16 matmul: Ok [4, 4] F16

四种失败形式：unsupported dtype、accelerate 的专门报错、Metal 找不到内核、panic。同一段代码，换一个设备或一个 feature，结果就从能跑变成报错。

出处

CPU 的 matmul 是 impl Map2 for MatMul 里按 feature 选的 f。不开 accelerate 时是 gemm 那一支，开头就是白名单：

        match T::DTYPE {
            DType::F16 | DType::F32 | DType::F64 => {}
            _ => Err(Error::UnsupportedDTypeForOp(T::DTYPE, "matmul").bt())?,
        }

开了 accelerate 是另一个 f，只有 F32（sgemm）、F64（dgemm）两支，F16 单独 bail!("the accelerate backend does not support f16 matmul")，其余落到 UnsupportedDTypeForOp。两种构建都没有 BF16。

Metal 的 MetalStorage::matmul 只把 F32、F16、BF16 映射到 GemmDType，其余报 mlx matmul doesn't support ..。逐元素运算按 dtype 拼出内核名（badd_i32），到已编译的库里找，binary.metal 里没有这个实例化时就是 Error while loading function。to_dtype 按 (源, 目标) 对列举内核，没列到的组合 bail!("Metal contiguous to_dtype {left:?} {right:?} not implemented")：F64 上传之后在 Metal 上连转回 F32 都做不到，只能 to_device(&Device::Cpu) 下载。

整数的 exp 不经过任何 dtype 检查。一元运算由 op.rs 的 unary_op! 宏生成，整数分支是 todo!：

            fn u32(_: u32) -> u32 {
                todo!("no unary function for u32")
            }

panic 信息里的 op.rs:584:1 是宏调用 unary_op!(Exp, "exp", ..) 那一行，不是 todo! 所在的行；同一个宏生成的 log、sin、neg、sqr、sqrt 等在整数上都一样。

修法

在 CPU 上转好 dtype 再算。 BF16 权重在 CPU 上先转 F32，乘完需要的话再转回：

let y = x.to_dtype(DType::F32)?.matmul(&w.to_dtype(DType::F32)?)?.to_dtype(DType::BF16)?;

cpu BF16 -> F32 matmul -> BF16: Ok [4, 4] BF16
cpu U32 -> F32 exp: Ok [3] F32
I32 on cpu -> F32 -> metal add: Ok [4, 4] F32

加载整个模型时不用逐个转：VarBuilder::from_mmaped_safetensors(.., dtype, ..) 的 dtype 参数让每个张量取出时就转好（见 VarBuilder）；device.bf16_default_to_f32() 在 CPU 上返回 F32、在 Metal 上返回 BF16，可以拿来选这个参数。

几条规则：

• CPU 上的 matmul 用 F32（开 accelerate 时 F16 也不行）；BF16 只用来存。
• Metal 上算东西只用 F32、F16、BF16；U8、U32、I64 够做掩码和下标。I16、I32、F64 的数据在 CPU 上转成 F32 或 I64 再上传，Metal 上的转换内核没有它们。
• 整数张量不要做 exp、sqrt 一类一元运算，先转浮点。

怎么确认没中招

• 报错里出现 unsupported dtype、doesn't support、not implemented、Error while loading function 中的任何一个，先打印参与运算的每个张量的 dtype() 和 device()。
• 换设备或换 feature 后第一次跑，用一个小输入把模型完整前向一遍：dtype 问题都在第一次调用时暴露，不依赖数据。
• panic 位置在 candle-core-0.11.0/src/op.rs，信息是 no unary function for <整数类型>：找整数张量上的一元运算。

小结

• dtype 一致只是前提，后端实现了才算数；没实现的报法有四种，整数一元运算是 panic。
• CPU matmul：gemm 支持 F16 / F32 / F64，accelerate 只有 F32 / F64；两者都不支持 BF16。
• Metal：F32 / F16 / BF16 能算；I16 / I32 / F64 只能上传下载，转换也要在 CPU 上做。