Rust in Everything
基于 candle 0.11.0

SDPA has no cpu impl

主线位置：推理主线第二步「建模块」，手写注意力时照 PyTorch 的习惯调融合的 sdpa。

版本：PyTorch 2.15.0（CPU）、candle 0.11.0（CPU，accelerate；Metal）
flowchart TB
  sd["ops sdpa"] --> dev{"设备"}
  dev -->|CPU| e1["no cpu impl"]
  dev -->|Metal| hd{"head dim 在名单里"}
  hd -->|否| e2["head dim 不支持"]
  hd -->|"512 且 F32 多 token"| e3["q dims 不支持"]
  hd -->|是| ok["融合内核"]
  e1 --> hand["手写 attention"]
  e2 --> hand
  e3 --> hand

实验代码：PyTorch 侧 sdpa.py，candle 侧 crate pitb-exp 的 bin p24_sdpa（release）。平台 M2 Max、macOS 27、rustc 1.99.0。candle 有哪些融合算子、sdpa 在其中的位置，见什么时候选 candle的「算子覆盖」和源码卷的常用层。

PyTorch 的写法

transformers 的注意力默认走 F.scaled_dot_product_attention，自己写模型时也常这样调。四种 head dim，CPU 上：

for d in [4, 48, 64, 512]:
  q = torch.arange(4 * d, dtype=torch.float32).reshape(1, 1, 4, d) / (4 * d)
  o = F.scaled_dot_product_attention(q, q, q, is_causal=True)

cpu head_dim   4: out (1, 1, 4, 4) last row[:2] [0.50061, 0.56311]
cpu head_dim  48: out (1, 1, 4, 48) last row[:2] [0.681584, 0.686792]
cpu head_dim  64: out (1, 1, 4, 64) last row[:2] [0.698097, 0.702003]
cpu head_dim 512: out (1, 1, 4, 512) last row[:2] [0.748214, 0.748702]

PyTorch 的 sdpa 是一个入口：在每个设备上都有实现，内部按形状和后端选融合内核或退回数学实现，调用方不用管。

搬到 candle 的现象

candle-nn 里有一个同名的函数，参数也相近（q, k, v, mask, do_causal, scale, softcapping）：

let c = candle_nn::ops::sdpa(&q, &q, &q, None, true, scale, 1.0);
let qm = q.to_device(&metal)?;
match candle_nn::ops::sdpa(&qm, &qm, &qm, None, true, scale, 1.0) { /* ... */ }

head_dim   4 cpu   sdpa: Err(SDPA has no cpu impl)
head_dim   4 metal sdpa: Err(Meta SDPA does not support q head dim 4: q dims [1, 1, 4, 4], k dims [1, 1, 4, 4], v dims [1, 1, 4, 4].)
head_dim  48 cpu   sdpa: Err(SDPA has no cpu impl)
head_dim  48 metal sdpa: Err(Meta SDPA does not support q head dim 48: q dims [1, 1, 4, 48], k dims [1, 1, 4, 48], v dims [1, 1, 4, 48].)
head_dim  64 cpu   sdpa: Err(SDPA has no cpu impl)
head_dim  64 metal sdpa: ok, max|diff| vs torch 5.96e-8
head_dim 512 cpu   sdpa: Err(SDPA has no cpu impl)
head_dim 512 metal sdpa: Err(Meta SDPA does not support q dims [1, 1, 4, 512], k dims [1, 1, 4, 512], v dims [1, 1, 4, 512].)
head_dim 512 metal sdpa F16: ok

八次调用只有两次成功。CPU 上一律失败；Metal 上 head dim 4、48 不在名单里；512 在名单里，但 F32 的多 token 输入被拒，换成 F16 才行。好在都是 Err，不是 panic，也不是静默算错：唯一成功的 F32 那一次和 PyTorch 差 5.96e-8。

原因

candle_nn::ops::sdpa 是一个自定义算子 Sdpa，它的 cpu_fwd 只有一句：

        candle::bail!("SDPA has no cpu impl")

metal_fwd 按 head dim 和形状选两种 MLX 风格的内核：

        let supported_head_dim = q_head == 32
            || q_head == 64
            || q_head == 72
            || q_head == 80
            || q_head == 96
            || q_head == 128
            || q_head == 256
            || q_head == 512;

        let supports_sdpa_full_mask = self.mask.is_none() || q_seq <= k_seq;
        // F32 full attention at head_dim=512 exceeds 32KB Metal threadgroup memory
        let supports_sdpa_full_dtype = !(q_head == 512 && q.dtype() == DType::F32);
        let supports_sdpa_full =
            q_seq > 1 && supported_head_dim && supports_sdpa_full_mask && supports_sdpa_full_dtype;
        let supports_sdpa_vector = q_seq == 1 && supported_head_dim && q_seq <= k_seq;
  • head dim 不在这 8 个里：does not support q head dim。
  • 多于一个查询（prefill）走 full 内核，head dim 512 的 F32 超过 Metal 32KB 的 threadgroup 内存（源码注释），被排除；单个查询（decode）走 vector 内核，没有这条限制。所以同一个 head dim 512 的 F32 模型，decode 一步能用 sdpa，prefill 不能。

它不是 PyTorch 那种「哪里都能调、内部选实现」的入口，而是一个 Metal 专用的加速内核。origin/main 上 ops.rs 和 0.11.0 相同。

📝NOTE
冷知识：0.10 里 4 个 token 走的是另一个内核

PR #3479（「metal: route q_seq > 1 through full SDPA kernel, not vector」，首次进入的 tag 是 0.11.0）之前，metal_fwd 的分界是 8：supports_sdpa_full = q_seq > 8 && ...、supports_sdpa_vector = q_seq <= 8 && ...，2 到 8 个查询也交给为 decode 写的 vector 内核。这个 PR 的 diff 只改了这两个条件，变成上面的 q_seq > 1 和 q_seq == 1，提交说明里写的是「The functional fix collapses to two predicate edits in candle-nn/src/ops.rs.」。512 的 F32 排除条件只挂在 full 上（它来自加入 head dim 512 的 PR #3417，首次进入 0.10.0），所以按 diff 推，0.10 里本页 4 个 token、head dim 512 的 F32 调用会走 vector 内核、不报错；0.11.0 起报错。本页没有编译 0.10 去跑。

改法

自己拼注意力：matmul、按因果关系 where_cond、softmax_last_dim、matmul。每个设备、每种 head dim 都能跑：

fn attention(q: &Tensor, k: &Tensor, v: &Tensor, scale: f64) -> candle_core::Result<Tensor> {
  let (n, m) = (q.dim(D::Minus2)?, k.dim(D::Minus2)?);
  let w = (q.matmul(&k.t()?)? * scale)?;
  // query i (at position m - n + i) sees keys 0..=m - n + i
  let mask: Vec<u8> = (0..n).flat_map(|i| (0..m).map(move |j| u8::from(j <= m - n + i))).collect();
  let mask = Tensor::from_vec(mask, (n, m), q.device())?.broadcast_as(w.shape())?;
  let neg = Tensor::new(f32::NEG_INFINITY, q.device())?.to_dtype(w.dtype())?.broadcast_as(w.shape())?;
  candle_nn::ops::softmax_last_dim(&mask.where_cond(&w, &neg)?)?.matmul(v)
}

head_dim   4 cpu   hand attention: max|diff| vs torch 2.98e-8
head_dim   4 metal hand attention: max|diff| vs torch 5.96e-8
head_dim  48 cpu   hand attention: max|diff| vs torch 1.79e-7
head_dim  48 metal hand attention: max|diff| vs torch 1.19e-7
head_dim  64 cpu   hand attention: max|diff| vs torch 5.96e-8
head_dim  64 metal hand attention: max|diff| vs torch 5.96e-8
head_dim 512 cpu   hand attention: max|diff| vs torch 5.96e-8
head_dim 512 metal hand attention: max|diff| vs torch 1.19e-7

输出都在 [0, 1] 之间，最大差 1.79e-7，容差 1e-6，是 f32 求和顺序的量级。mask 写成「查询 i 看到前 m - n + i + 1 个键」，带 KV 缓存时 n < m 也对，和 PyTorch is_causal=True 不同：PyTorch 按左上角对齐，1 个查询、3 个键（值分别是 0、1、2）时输出 0.0，只看到了第 0 个键，不带 mask 时是 1.0。所以带缓存 decode 时，PyTorch 一侧也不能用 is_causal=True。candle_transformers::models::llama 不开 flash-attn 时就是这四步（对齐 logits）。

要在 Metal 上用 sdpa 提速，可以保留手写版本做退路：sdpa 返回 Err 时调 attention。

怎么检查

  • 在 CPU 上跑一遍单元测试：sdpa 只要被调用就会 Err，最容易暴露。
  • 只在 Metal 上跑的代码，测试要覆盖 prefill（多个查询）和 decode（一个查询）两种形状，并确认模型的 head dim 在上面 8 个之内；head dim 512 的 F32 模型 prefill 不能用。
  • 把手写注意力和 sdpa 的输出对一次（这里是 5.96e-8），确认缩放、mask 的约定一致。

小结

  • PyTorch 的 F.scaled_dot_product_attention 在每个设备、每种 head dim 上都能调。
  • candle 的 candle_nn::ops::sdpa 只有 Metal 实现：CPU 上 Err(SDPA has no cpu impl)；Metal 上 head dim 只认 32、64、72、80、96、128、256、512，512 配 F32 时多个查询不行。
  • 手写 matmul + where_cond + softmax_last_dim + matmul，CPU 和 Metal 上和 PyTorch 差不超过 1.79e-7。

💬 关联讨论

发起讨论
加载中...