dtype mismatch in rope BF16 F32 F32

主线位置：推理主线第四步「带 KV 缓存生成」，模型换成 bf16，缓存那一行没改。

版本：PyTorch 2.15.0 / transformers 5.22.0（CPU）、candle 0.11.0（CPU，accelerate；Metal）

flowchart TB
  vb["VarBuilder BF16"] --> q["q k 是 BF16"]
  c["Cache new F32"] --> cs["cos sin 是 F32"]
  q --> rope["rope 内核"]
  cs --> rope
  rope -->|Metal| e1["dtype mismatch"]
  rope -->|CPU| e2["unsupported"]

实验代码：PyTorch 侧 kv.py，candle 侧 crate pitb-exp 的 bin p22_cachedt（release）。模型是本地的 SmolLM2-135M-Instruct 权重。平台 M2 Max、macOS 27、rustc 1.99.0。

PyTorch 的写法

transformers 里换精度只改 from_pretrained 的 dtype，缓存对象不用动：

m16 = AutoModelForCausalLM.from_pretrained(model_dir, dtype=torch.bfloat16, attn_implementation="eager").eval()
c16 = DynamicCache(config=m16.config)
m16(ids, past_key_values=c16, use_cache=True)
print("bf16 model: cache keys", c16.layers[0].keys.dtype, "rotary inv_freq", m16.model.rotary_emb.inv_freq.dtype)

bf16 model: cache keys torch.bfloat16 rotary inv_freq torch.float32

DynamicCache 的构造函数没有 dtype 参数，第一次写入时直接存下传进来的 k、v，dtype 跟着模型走。RoPE 的 inv_freq 在模型的 rotary_emb 里，留在 float32，每次前向现算 cos / sin 再转成输入的 dtype。

搬到 candle 的现象

从 F32 改成 BF16 时，只改了 VarBuilder 那一行（Metal，CPU 上 BF16 不能 matmul，见 unsupported dtype BF16 for op matmul）：

let dt = DType::BF16;
let model = Llama::load(st_vb(dt, &dev)?, &cfg)?;
let mut cache = Cache::new(true, DType::F32, &cfg, &dev)?;
let x = Tensor::new(&PROMPT, &dev)?.unsqueeze(0)?;
match model.forward(&x, 0, &mut cache) {
  Ok(l) => println!("F32 cache: ok {:?}", l.dims()),
  Err(e) => println!("F32 cache: Err({e})"),
}

F32 cache: Err(dtype mismatch in rope BF16 F32 F32)

CPU 上反过来（模型 F32、Cache 用 BF16 建），报的是另一句话：

cpu F32 model, BF16 cache: Some("unsupported dtype for rope F32 BF16 BF16")

报错里没有出现 Cache，三个 dtype 分别是 q（或 k）、cos、sin。

原因

candle 的 Cache 除了 KV，还装着 RoPE 的 cos / sin 表和 mask。Cache::new 在 F32 里算好整张表（max_position_embeddings 行），最后转成传进来的 dtype：

        let cos = idx_theta.cos()?.to_dtype(dtype)?;
        let sin = idx_theta.sin()?.to_dtype(dtype)?;

前向时 apply_rotary_emb 把 q、k 和这两张表一起交给 candle_nn::rotary_emb::rope，这是一个融合算子，不做 dtype 转换。Metal 实现 RotaryEmb::metal_fwd 开头检查三者一致：

        if cos.dtype() != src.dtype() || sin.dtype() != src.dtype() {
            candle::bail!(
                "dtype mismatch in rope {:?} {:?} {:?}",
                src.dtype(),
                cos.dtype(),
                sin.dtype()
            )
        }

CPU 实现 RotaryEmb::cpu_fwd 用 match (s1, s2, s3) 只列了四种三者相同的组合，其余落到 unsupported dtype for rope。两个后端措辞不同，说的是同一件事。transformers 把表放在模型里、每次按输入转换，所以没有这个坑；candle 把表放在缓存里、一次转好，于是缓存也有了 dtype。

改法

VarBuilder 和 Cache::new 用同一个 dt 变量：

let dt = DType::BF16;
let model = Llama::load(st_vb(dt, &dev)?, &cfg)?;
let mut cache = Cache::new(true, dt, &cfg, &dev)?;

BF16 cache: logits max|diff|=6.02e-1 (max|ref| 20.6), greedy same prefix 20/20

这是 Metal BF16 的正常结果：和 transformers f32 差 0.602，bf16 该有的量级（半精度）。cos / sin 表在 BF16 里存；transformers 那边在 float32 里算出 cos / sin，再 cos.to(dtype=x.dtype) 转成 bf16（transformers 5.22.0 的 LlamaRotaryEmbedding.forward），两边最后交给旋转的都是 bf16 的表。

怎么检查

• dtype 只定义一次，VarBuilder、Cache::new、自己建的 mask 或常量张量都用它。
• 看到 rope 加三个 dtype 的报错，第一个是激活的 dtype，后两个是缓存里的表：去找 Cache::new 那一行。
• Cache 的字段是私有的，读不出它的 dtype；能检查的只有自己传进去的那个变量。

小结

• transformers 的 DynamicCache 没有 dtype 参数，跟着写入的 k、v 走；RoPE 的表在模型里，每次按输入转换。
• candle llama.rs 的 Cache 装着 RoPE 的 cos / sin 表，Cache::new 的 dtype 必须和模型一致。不一致时 Metal 报 dtype mismatch in rope BF16 F32 F32，CPU 报 unsupported dtype for rope F32 BF16 BF16。
• 改法：同一个 dt 传给 VarBuilder 和 Cache::new；Metal BF16 跑通，差 0.602。