Rust in Everything
基于 candle-core 0.11.0

自定义算子

主线位置：算子主线的一条旁路，并且接到求导主线上。BackendStorage 里没有 softmax、norm、注意力（见 Storage 与后端 trait），这些融合算子由 candle-nn 用自定义算子实现：xs.apply_op2_no_bwd(alpha, &RmsNorm { eps }) → Storage::apply_op2 按设备对选 cpu_fwd 或 metal_fwd。带反向的版本还会在结果张量上记一个 Op::CustomOp1/2/3，backward 走到它时调用户写的 bwd（求导主线的整体见 backward 与 GradStore）。

flowchart TB
  u["x apply_op1 c"] --> s["Storage apply_op1"]
  s -->|Cpu| cf["c cpu_fwd"]
  s -->|Metal| mf["c metal_fwd"]
  s -->|Cuda| uf["c cuda_fwd"]
  cf --> t["新 Tensor"]
  mf --> t
  uf --> t
  t --> op["Op CustomOp1 x c"]
  op --> bw["backward 调 c bwd"]
  nb["apply_op1_no_bwd"] --> s
  nb -.->|不记 Op| t2["无梯度的 Tensor"]

代码在 candle-core/src/custom_op.rs（488 行）；用户在 candle-nn/src/ops.rs。

本页的实验 crate 是 exp-candle/backend-ops（candle-core / candle-nn 0.11.0，features metal accelerate，另外直接依赖 candle-metal-kernels 和 objc2-metal 来写 Metal 内核；M2 Max、macOS 27、rustc 1.99.0；release）。

三个 trait

CustomOp1 是一个输入的版本：

pub trait CustomOp1 {
    // Box<dyn> does not support const yet, so use a function to get the name.
    fn name(&self) -> &'static str;

    /// The forward pass, as run on a cpu device. Note that the storage can use arbitrary strides,
    /// offsets etc so the associated layout should be used to access it.
    fn cpu_fwd(&self, storage: &CpuStorage, layout: &Layout) -> Result<(CpuStorage, Shape)>;

    // ...
    fn cuda_fwd(&self, _storage: &CudaStorage, _layout: &Layout) -> Result<(CudaStorage, Shape)> {
        Err(crate::Error::Cuda(
            format!("no cuda implementation for {}", self.name()).into(),
        ))
    }

    // ...
    fn metal_fwd(
        &self,
        _storage: &MetalStorage,
        _layout: &Layout,
    ) -> Result<(MetalStorage, Shape)> {
        Err(crate::Error::Metal(
            format!("no metal implementation for {}", self.name()).into(),
        ))
    }

    /// This function takes as argument the argument `arg` used in the forward pass, the result
    /// produced by the forward operation `res` and the gradient of the result `grad_res`.
    /// The function should return the gradient of the argument.
    fn bwd(&self, _arg: &Tensor, _res: &Tensor, _grad_res: &Tensor) -> Result<Option<Tensor>> {
        Err(crate::Error::BackwardNotSupported { op: self.name() })
    }
}
  • cpu_fwd 必须写，其余三个有默认实现，默认就是报错。一个只写了 CPU 的算子放到 Metal 张量上，得到 no metal implementation for <name>。
  • 前向拿到的是后端自己的存储（CpuStorage 枚举、MetalStorage 缓冲区）加 Layout，返回新存储和输出形状。注释特意提醒：输入可以带任意步长和偏移，要按 Layout 读。处理不了的布局（比如非连续）自己报错。
  • 前向是「后端层」的代码，碰不到 Tensor；反向 bwd 正相反，拿到的是 Tensor（输入 arg、输出 res、输出的梯度 grad_res），用普通张量运算拼出输入的梯度。所以前向可以是手写的融合内核，反向不用再写内核。

CustomOp2 和 CustomOp3 形状相同，只是参数多一份、两份，bwd 返回元组，每个输入一个 Option<Tensor>。返回 None 表示这个输入不需要梯度。

怎么接进张量和求导

Tensor 上有两组方法。带反向的 apply_op1_arc：

    pub fn apply_op1_arc(&self, c: Arc<Box<dyn CustomOp1 + Send + Sync>>) -> Result<Self> {
        let (storage, shape) = self
            .storage()
            .apply_op1(self.layout(), c.as_ref().as_ref())?;
        let op = BackpropOp::new1(self, |s| Op::CustomOp1(s, c.clone()));
        Ok(from_storage(storage, shape, op, false))
    }

    pub fn apply_op1<C: 'static + CustomOp1 + Send + Sync>(&self, c: C) -> Result<Self> {
        self.apply_op1_arc(Arc::new(Box::new(c)))
    }

算子对象要放进 Arc<Box<dyn CustomOp1 + Send + Sync>> 存到计算图上，所以 apply_op1 要求 'static + Send + Sync，并且按值拿走它。不带反向的 apply_op1_no_bwd 只借用：

    pub fn apply_op1_no_bwd<C: CustomOp1>(&self, c: &C) -> Result<Self> {
        let (storage, shape) = self.storage().apply_op1(self.layout(), c)?;
        Ok(from_storage(storage, shape, BackpropOp::none(), false))
    }

BackpropOp::none()：结果张量没有 op，在计算图里是一个起点。哪怕输入是 Var，梯度也到这里为止，而且不报错（BackpropOp 什么时候记、什么时候不记见计算图怎么记下来）。

Storage::apply_op1 就是按变体调 cpu_fwd / cuda_fwd / metal_fwd；apply_op2、apply_op3 先 same_device，不检查 dtype，dtype 不一致要算子自己处理。

求导这一边，backprop.rs 里 Op::CustomOp1 的分支只有几行：

                    Op::CustomOp1(arg, c) => {
                        if let Some(arg_grad) = c.bwd(arg, node, &grad)? {
                            let sum_grad = grads.or_insert(arg)?;
                            *sum_grad = sum_grad.add(&arg_grad)?
                        }
                    }

node 就是前向的结果张量，即 bwd 的 res 参数。反向本身是普通张量运算，所以「bwd 里再用一个自定义算子」也可以，只是它不会被再次求导，除非用 apply_op* 的带反向版本（二阶导）。

candle-nn 的融合算子

candle-nn/src/ops.rs 里有五个自定义算子：

函数算子trait反向
sigmoidSigmoidCustomOp1有，apply_op1
softmax_last_dimSoftmaxLastDimCustomOp1无，apply_op1_no_bwd
rms_normRmsNormCustomOp2无，apply_op2_no_bwd
layer_normLayerNormCustomOp3无，apply_op3_no_bwd
sdpaSdpaCustomOp3无；只有 Metal 实现，cpu_fwd 报错

以 SoftmaxLastDim 为例，CPU 前向：

    fn cpu_fwd(&self, storage: &CpuStorage, layout: &Layout) -> Result<(CpuStorage, Shape)> {
        fn softmax<T: candle::WithDType + num_traits::Float>(
            src: &[T],
            layout: &Layout,
        ) -> Result<(CpuStorage, Shape)> {
            let src = match layout.contiguous_offsets() {
                None => candle::bail!("input has to be contiguous"),
                Some((o1, o2)) => &src[o1..o2],
            };
            let el_count = layout.shape().elem_count();
            let dims = layout.shape().dims();
            let dim_m1 = dims[dims.len() - 1];
            let mut dst = vec![T::zero(); el_count];
            src.par_chunks(dim_m1)
                .zip(dst.par_chunks_mut(dim_m1))
                .for_each(|(src, dst)| {
                    let mut max = T::neg_infinity();
                    unsafe { T::vec_reduce_max(src.as_ptr(), &mut max, dim_m1) };
                    for (s, d) in src.iter().zip(dst.iter_mut()) {
                        *d = (*s - max).exp();
                    }
                    let mut sum_exp = T::zero();
                    unsafe { T::vec_reduce_sum(dst.as_ptr(), &mut sum_exp, dim_m1) };
                    for d in dst.iter_mut() {
                        *d /= sum_exp
                    }
                });
            let storage = candle::WithDType::to_cpu_storage_owned(dst);
            Ok((storage, Shape::from_dims(dims)))
        }

        match storage {
            CpuStorage::BF16(slice) => softmax::<half::bf16>(slice, layout),
            CpuStorage::F16(slice) => softmax::<half::f16>(slice, layout),
            CpuStorage::F32(slice) => softmax::<f32>(slice, layout),
            CpuStorage::F64(slice) => softmax::<f64>(slice, layout),
            _ => candle::bail!("unsupported dtype for softmax {:?}", storage),
        }
    }

每一行一个 rayon 任务，行内用 VecOps 的 SIMD 求最大值和总和（见 CPU 后端），一次遍历算完，不产生中间张量。组合版的 softmax 是 max_keepdim → broadcast_sub → exp → sum_keepdim → broadcast_div，五个 op、四个中间张量。Metal 前向调 candle_metal_kernels::call_last_softmax，内核在 reduce.metal 里，也是一个 dispatch。这就是「融合」的全部含义：手写一个内核代替一串 op。candle 没有自动融合，想省中间张量只能这样一个个手写。

融合的代价写在表的最后一列。grad_of 程序对同一个 Var 用不同的算子，loss = y.sqr().sum_all()，再 backward：

fn grad_of(label: &str, f: impl Fn(&Tensor) -> candle_core::Result<Tensor>) -> anyhow::Result<()> {
  let x = Var::new(&[[-1f32, 0., 1.], [2., 3., 4.]], &Device::Cpu)?;
  let y = f(x.as_tensor())?;
  let loss = y.sqr()?.sum_all()?;
  let grads = loss.backward()?;
  match grads.get(x.as_tensor()) {
    Some(g) => println!("{label:<34} grad = {:?}", g.to_vec2::<f32>()?),
    None => println!("{label:<34} grad = None"),
  }
  Ok(())
}

candle_nn::ops::sigmoid            grad = [[0.1057542, 0.25, 0.28746969], [0.18495615, 0.08606823, 0.034690097]]
candle_nn::ops::softmax_last_dim   grad = None
candle_nn::ops::softmax(D::Minus1) grad = [[-0.075717956, -0.1301048, 0.20582275], [-0.075717956, -0.1301048, 0.20582275]]
candle_nn::ops::rms_norm           grad = None

softmax_last_dim 和 rms_norm 后面，x 的梯度是 None，backward 不报错。candle-nn 的层也受影响：RmsNorm 层的 forward 在输入连续时直接调 ops::rms_norm，LayerNorm 层在 remove_mean、输入连续且有偏置时调 ops::layer_norm。要训练，就用组合版：ops::softmax(xs, D::Minus1)、RmsNorm::forward_diff、ops::rms_norm_slow、ops::layer_norm_slow。这几个层怎么用见常用层。

❗IMPORTANT
动脑筋：candle_nn::RmsNorm 层前向一个 Var，backward 之后输入 x 的梯度是什么？换成 forward_diff 呢？

rmsnorm_layer 程序：

let norm = candle_nn::RmsNorm::new(Tensor::ones(3, DType::F32, &Device::Cpu)?, 1e-5);
for (label, diff) in [("RmsNorm::forward", false), ("RmsNorm::forward_diff", true)] {
  let x = Var::new(&[[-1f32, 0., 1.], [2., 3., 4.]], &Device::Cpu)?;
  let y = if diff { norm.forward_diff(x.as_tensor())? } else { norm.forward(x.as_tensor())? };
  let w = Tensor::new(&[1f32, 2., 3.], &Device::Cpu)?;
  let grads = y.broadcast_mul(&w)?.sum_all()?.backward()?;
  // 打印 grads.get(x.as_tensor())
}

写一个自定义算子

custom_op 程序实现 softplus，y = ln(1 + eˣ)，导数是 sigmoid(x)。CPU 前向只处理连续的 f32：

struct Softplus;

impl CustomOp1 for Softplus {
  fn name(&self) -> &'static str {
    "softplus"
  }

  fn cpu_fwd(&self, storage: &CpuStorage, layout: &Layout) -> candle_core::Result<(CpuStorage, Shape)> {
    let xs = storage.as_slice::<f32>()?;
    let (start, end) = match layout.contiguous_offsets() {
      Some(o) => o,
      None => candle_core::bail!("softplus: input has to be contiguous"),
    };
    let ys: Vec<f32> = xs[start..end].iter().map(|&x| (1.0 + x.exp()).ln()).collect();
    Ok((CpuStorage::F32(ys), layout.shape().clone()))
  }

  fn bwd(&self, arg: &Tensor, _res: &Tensor, grad_res: &Tensor) -> candle_core::Result<Option<Tensor>> {
    // d/dx softplus(x) = sigmoid(x)
    let s = candle_nn::ops::sigmoid(arg)?;
    Ok(Some(grad_res.mul(&s)?))
  }
  // metal_fwd 见下
}

as_slice::<f32>() 在 dtype 不是 f32 时返回错误，所以别的 dtype 不会被当成 f32 读。Metal 前向要自己带一段 .metal 源码，第一次用时编译、存进一个 OnceLock，之后按 candle-nn 的写法取缓冲区和编码器：

const SRC: &str = r#"
#include <metal_stdlib>
using namespace metal;
kernel void softplus_f32(
  constant size_t &n [[buffer(0)]],
  device const float *x [[buffer(1)]],
  device float *y [[buffer(2)]],
  uint i [[thread_position_in_grid]]
) {
  if (i >= n) return;
  y[i] = log(1.0f + exp(x[i]));
}
"#;

    let n = layout.shape().elem_count();
    let out = device.new_buffer_builder().with_size_for(n, DType::F32).build()?;
    let encoder = device.command_encoder()?;
    let enc: &candle_metal_kernels::metal::ComputeCommandEncoder = encoder.as_ref();
    enc.set_compute_pipeline_state(pipeline);
    let input = BufferOffset {
      buffer: storage.buffer(),
      offset_in_bytes: layout.start_offset() * 4,
    };
    candle_metal_kernels::set_params!(enc, (n, &input, Output::new(&out)));
    let width = pipeline.max_total_threads_per_threadgroup().min(n);
    let grid = objc2_metal::MTLSize { width: n.div_ceil(width), height: 1, depth: 1 };
    let group = objc2_metal::MTLSize { width, height: 1, depth: 1 };
    enc.dispatch_thread_groups(grid, group);
    drop(encoder);
    Ok((MetalStorage::new(out, device.clone(), n, DType::F32), layout.shape().clone()))

三处要照 candle 的规矩来：

  • 输出缓冲区从 new_buffer_builder() 取，进 candle 的缓冲池，张量 drop 后可以复用。
  • 编码器用 device.command_encoder()，往共享的并发编码器里追加，不自己建命令缓冲区，这样提交和等待都由 candle 管。
  • 输出要包成 Output::new(&out)，让编码器把它登记为输出；candle 靠这个登记决定要不要插内存屏障（见 Metal 后端）。裸 &Buffer 会被当成输入。本页曾故意把输出当输入登记、紧接着让一个 add 读它，20 次运行结果都对，没能复现出错，但正确性不该靠这个运气。

OnceLock 里的管线绑在第一次用到的那个设备上；多设备的程序要按设备缓存，或者像 candle-nn 那样把源码交给 device.kernels()（Kernels 只认 candle-metal-kernels 自带的 Source，外部源码进不去，所以这里自己缓存）。

运行结果：

Cpu softplus [0.12692805, 0.474077, 0.6931472, 0.974077, 3.0485873]
Cpu max |diff| vs exp/log: 0.000000059604645
Metal { gpu_id: 4294968354 } softplus [0.12692788, 0.4740769, 0.6931472, 0.97407687, 3.048587]
Metal { gpu_id: 4294968354 } max |diff| vs exp/log: 0
apply_op1(Softplus)                grad = [[0.16849808, 0.6931472, 1.9201424], [3.746784, 5.808011, 7.891757]]
apply_op1_no_bwd(&Softplus)        grad = None
transposed: Err(softplus: input has to be contiguous)
  • 和组合版 (x.exp() + 1).log() 比，CPU 差 6e-8（一个 ulp 量级），Metal 完全一致（同样的 exp、log，同样开着快速数学）。
  • 梯度对得上：loss = Σ y²，∂loss/∂x = 2·y·sigmoid(x)；x = 0 处 y = ln 2、sigmoid = 0.5，梯度 ln 2 = 0.6931472。
  • 同一个算子换成 apply_op1_no_bwd，梯度是 None。
📝NOTE
没有蠢问题

问： 两个融合算子都要求输入连续。注意力里常见 q.matmul(&k.t()?)? 之后接 softmax，会不会报错？

答： 要看 softmax 的输入本身。matmul 的输出是新分配的连续张量，没问题；直接对一个转置视图调就会报错，而且 CPU 和 Metal 的报错不一样：

softmax_last_dim(cpu, transposed): Err(input has to be contiguous)
softmax_last_dim(metal, transposed): Err(Non contiguous softmax-last-dim is not implemented)

CPU 那句没提算子名，在一长串前向里不容易定位。先 .contiguous()? 再调，或者用组合版 softmax（逐元素和归约 op 都能处理步长）。

原地算子

InplaceOp1 / InplaceOp2 / InplaceOp3 的前向拿 &mut CpuStorage（或 &mut MetalStorage）、返回 ()，没有 bwd，注释写明「These ops work in place and as such back-prop is unsupported」。Tensor::inplace_op1 接收的是 &self：

    pub fn inplace_op1<C: InplaceOp1>(&self, c: &C) -> Result<()> {
        self.storage_mut().inplace_op1(self.layout(), c)
    }

    /// Applies a unary custom op in place (for the first tensor).
    pub fn inplace_op2<C: InplaceOp2>(&self, rhs: &Self, c: &C) -> Result<()> {
        self.storage_mut()
            .inplace_op2(self.layout(), &rhs.storage(), rhs.layout(), c)
    }

storage_mut() 是 Arc<RwLock<Storage>> 的写锁（Tensor 的结构）。所以「原地」改的是整块共享存储：这个张量的所有 clone、所有视图都会看到变化，计算图里存着的旧张量也会。

inplace_op2 先取 self 的写锁，再取 rhs 的读锁。如果 rhs 和 self 是同一块存储（同一个张量、它的 clone 或视图），同一线程在持有写锁时再要读锁，std::sync::RwLock 不会成功。inplace_self 程序（debug）：

let a = Tensor::new(&[1f32, 2.], &Device::Cpu)?;
let b = Tensor::new(&[10f32, 20.], &Device::Cpu)?;
a.inplace_op2(&b, &AddInto)?;
println!("a after a += b: {:?}", a.to_vec1::<f32>()?);
let view = a.clone();
println!("clone shares storage: {:?}", view.to_vec1::<f32>()?);
println!("calling a.inplace_op2(&a, ..)");
a.inplace_op2(&a, &AddInto)?;

a after a += b: [11.0, 22.0]
clone shares storage: [11.0, 22.0]
calling a.inplace_op2(&a, ..)
still running after 5 s, killing 70088

（最后一行是外面的脚本打印的：进程 5 秒没有退出，被杀掉。）没有 panic，没有 Err，程序就停在那里。a.inplace_op2(&a, ..) 想表达「a += a」，得先 let a2 = a.copy()? 再传 &a2（在 Metal 上 copy 不够，见 Metal 后端）。

candle-core 自己唯一的原地算子实现是 UgIOp1：用 ug 这个小型内核生成器把一段 SSA 形式的内核编成 Metal / CUDA 代码，再作为 InplaceOp1 跑，cpu_fwd 直接报「ug ops are only supported on metal/cuda at the moment」。

📝NOTE
冷知识：ug 曾经是 candle-core 的必选依赖

PR #3268（「feat!: Make ug dependency optional」，首次进入 0.10.0）之前，candle-core 的 Cargo.toml 无条件依赖 ug，开 cuda / metal 时再拉 ug-cuda / ug-metal，lib.rs 无条件导出 UgIOp1。这个 PR 新建了 candle-ug crate 把三者包起来，在 candle-core 里加了 ug = ["dep:candle-ug"] feature，UgIOp1 和 MetalDevice::compile 都加上 #[cfg(feature = "ug")]。标题里的 ! 是破坏性变更的标记：以前直接用 candle_core::UgIOp1 的代码，升级后要显式打开 ug feature。

动脑筋答案

RmsNorm::forward       grad = None
RmsNorm::forward_diff  grad = [[2.449453, 2.4494715, 2.4494898], [-0.12199855, -0.022181034, 0.07763654]]
  • forward 是 None：输入连续，层调 ops::rms_norm，它用 apply_op2_no_bwd，结果张量没有 op，backward 从 loss 往回走到它就停了，x 根本不在拓扑排序的节点里。没有任何报错或警告。
  • forward_diff 调的是层内部 LayerNorm（remove_mean 为假）的组合路径：sqr → sum_keepdim → 除法 → 加 eps → sqrt → broadcast_div → broadcast_mul，每一步都记了 Op，梯度能传到 x；代价是多出几个中间张量和几个内核。
  • 这里的损失用了权重 [1, 2, 3]：如果用 y.sqr().sum_all()，RMS 归一化之后每行平方和几乎恒等于 3，梯度会接近 0，看不出区别。
  • 同一个层有两个入口，训练要用 forward_diff。它的文档注释写的是「Faster variant of the forward kernel, this can only be used on contiguous tensors though」，和代码对不上（更快、只限连续输入的是 forward 走的融合路径），不要按注释选。搞混时不报错，只会让上游参数一直不更新。

小结

  • CustomOp1/2/3 每个后端一个前向（只有 cpu_fwd 必写），前向操作后端存储加 Layout；bwd 操作 Tensor，用普通运算写反向。
  • apply_op* 把算子放进 Op::CustomOp* 记到图上，backward 调 bwd；apply_op*_no_bwd 不记，梯度在这里静默中断。
  • candle-nn 的 softmax_last_dim、rms_norm、layer_norm 是融合算子，CPU 上一行一个 rayon 任务、Metal 上一个内核，但都没有反向，而且要求输入连续；训练用组合版。
  • 自定义 Metal 前向要用 candle 的缓冲池和共享编码器，输出用 Output::new 登记。
  • InplaceOp* 改的是整块共享存储，没有反向；t.inplace_op2(&t, ..) 会在 RwLock 上死锁。

💬 关联讨论

发起讨论
加载中...