常用层

主线位置：推理主线的 forward 里每一层用到的东西。quantized_llama 的一次前向（见量化 llama 的一次前向）用了本页的 Embedding、ops::rms_norm（经 quantized_nn::RmsNorm）、rotary_emb::rope_i、ops::softmax_last_dim、ops::silu，Metal 上单 token 时还有 ops::sdpa；线性层换成了量化的 QMatMul（见 QTensor 与 QMatMul）。BERT 用的是 Linear、LayerNorm，见 BERT。融合算子怎么接进 CustomOp 机制在自定义算子，本页只看它们和拼出来的版本差在哪。

flowchart TB
  f["RmsNorm forward"] --> c{"输入连续"}
  c -->|是| fused["ops rms_norm"]
  c -->|否| comp["LayerNorm forward"]
  fused --> op2["CustomOp2 无反向"]
  op2 --> cpu["CPU 一个函数"]
  op2 --> mtl["Metal 一个内核"]
  comp --> many["sqr sum sqrt div mul"]
  many --> grad["每步记 Op 可求导"]

本页的实验是 nn-exp 的 layers（release），平台 M2 Max、macOS 27、rustc 1.99.0，依赖开了 metal 和 accelerate。计时先跑一次（「首次」，Metal 上含内核编译），再预热 5 次，然后取 100 次（规模扫描是 200 次）的平均；Metal 上每次读时钟前都 synchronize()。

Linear：权重是 [out, in]

Linear 只有 weight 和可选的 bias。linear 向 VarBuilder 要的形状是 (out_dim, in_dim)，和 PyTorch nn.Linear.weight 的约定一样，所以 Hugging Face 的权重不用转置就能装进来：上一页读到的 k_proj.weight 是 [192, 576]，就是输出 192、输入 576。burn 的 Linear 默认是 [in, out]，装 PyTorch 权重时要靠适配器转置（见 burn 的存储与加载）。

代价落在前向：y = x @ w.t() + b。forward 的三维分支：

            [bsize, m, k] => {
                if x.is_contiguous() {
                    let w = self.weight.t()?;
                    x.reshape((bsize * m, k))?
                        .matmul(&w)?
                        .reshape((bsize, m, ()))?
                } else {
                    let w = self.weight.broadcast_left(bsize)?.t()?;
                    x.matmul(&w)?
                }
            }

t() 只交换布局里的两个维度，不搬数据：

Linear (3,2) on (1,2): [[210.0, 430.0, 650.0]]
w.t() shape [2, 3] stride [1, 2] contiguous false

步长 [1, 2] 的右操作数直接交给 matmul，gemm 和 BLAS 都能按「转置」的方式读它，见 CPU 后端。连续的三维输入先压成二维再乘一次，源码注释说广播的批量 matmul 在 CUDA 和 CPU 上慢得多。在 CPU 上量一下 576 → 1536、输入 (4, 128, 576)：

Linear 576->1536 on (4,128,576): contiguous 0.430 ms, non-contiguous (broadcast path) 0.501 ms, contiguous() then forward 1.072 ms

非连续输入走广播分支只慢了 16%；为了走快路径先 contiguous() 反而是两倍多，复制的开销比 matmul 本身还大。

Embedding：index_select

Embedding::forward 把 id 张量摊平，index_select(&indexes, 0) 取出对应的行，再 reshape 成 id 的形状加一维 hidden_size。没有别的逻辑，id 的 dtype 要求就是 index_select 的要求（整数类型）。embedding 新建时用标准正态初始化。

归一化：LayerNorm 和 RmsNorm

RMSNorm 是不减均值的 LayerNorm，candle 就是这么实现的：RmsNorm 是 LayerNorm 的新类型包装，里面 remove_mean: false、没有偏置。两者都有两条路：

impl Module for RmsNorm {
    fn forward(&self, xs: &Tensor) -> Result<Tensor> {
        if xs.is_contiguous() {
            crate::ops::rms_norm(xs, &self.0.weight, self.0.eps as f32)
        } else {
            self.0.forward(xs)
        }
    }
}

self.0.forward 是 LayerNorm::forward：它在输入连续、要减均值、有偏置时走融合的 ops::layer_norm，否则用张量运算一步一步拼（F16 / BF16 先转 F32 再算）。拼出来的那条路每一步都记 Op，可以求导；融合算子用 apply_op2_no_bwd / apply_op3_no_bwd 调用（见 apply_op1_no_bwd），结果的 op 是 BackpropOp::none()，梯度到这里断掉。实验（layers grads，x 是 Var）：

RmsNorm::forward: grad of x = None
RmsNorm::forward_diff: grad of x = Some([[1.6391277e-7, 3.2782555e-7, 5.066395e-7, 6.556511e-7]])
LayerNorm::forward (with bias): grad of x present? false
softmax_last_dim: grad of x present? false
softmax: grad of x present? true
RmsNorm::forward on a transposed (non-contiguous) input: grad present? false
RmsNorm::forward on (4,2).t(): contiguous false, grad present? true

• RmsNorm::forward 在连续输入上没有梯度，backward 不报错，只是 GradStore 里没有 x。用 RmsNorm 训练时要调 forward_diff，它直接走 LayerNorm::forward 的拼接路径。forward_diff 的文档注释写的是「Faster variant of the forward kernel, this can only be used on contiguous tensors though」，描述的其实是 forward 里的融合路径，和它自己的行为正好相反。
• 带偏置的 LayerNorm::forward 同样没有梯度，而且没有对应的 forward_diff。
• 倒数第二行本想测非连续输入，结果也没有梯度：xt 是 (4, 1) 的 Var，t() 之后是 (1, 4)、步长 [1, 1]，Shape::is_contiguous 对长度为 1 的维度不看步长（dim > 1 && stride != acc 才算不连续），判成连续，走的仍是融合路径。换成 (4, 2) 转置（最后一行），真正不连续，走拼接路径，梯度就有了。所以同一个 RmsNorm::forward 有没有梯度取决于输入的布局。
• 梯度的值本身也说明了 RMSNorm 的性质：损失是 0.5 * Σ y²，y 的平方和恒等于元素个数（减去 eps 的影响），所以梯度几乎为 0。

ops::rms_norm 本身在入口检查 alpha 的长度等于 xs 的最后一维，不等就是 shape mismatch in rms-norm [..] [..]；quantized_nn::RmsNorm 不经过 RmsNorm::forward，直接调 ops::rms_norm，非连续输入会由内核报「input has to be contiguous」（CPU）或「Non contiguous rmsnorm is not implemented」（Metal）。

激活

Activation 是一个带 serde::Deserialize 的枚举，config.json 里的 "hidden_act": "silu" 可以直接反序列化成 Activation::Silu。名字按 Hugging Face 的约定对应：gelu 是精确的 erf 版本（gelu_erf），gelu_new、gelu_pytorch_tanh 是 tanh 近似（gelu）。silu、relu、gelu 是 candle-core 里的一元算子；sigmoid 是 candle-nn 里的 CustomOp1，和 softmax 不同，它用 apply_op1 调用，实现了 bwd，可以求导。

融合算子：每个设备一个内核

ops::softmax 是拼出来的：

pub fn softmax<D: candle::shape::Dim>(xs: &Tensor, dim: D) -> Result<Tensor> {
    let dim = dim.to_index(xs.shape(), "softmax")?;
    let max = xs.max_keepdim(dim)?;
    let diff = xs.broadcast_sub(&max)?;
    let num = diff.exp()?;
    let den = num.sum_keepdim(dim)?;
    num.broadcast_div(&den)
}

五个算子，每个是一次独立的遍历（Metal 上是一次内核调用），中间结果各占一块内存。softmax_last_dim 是 CustomOp1（SoftmaxLastDim），只做最后一维：

• CPU：par_chunks(dim_m1) 按行并行，每一行依次 vec_reduce_max、逐个 exp、vec_reduce_sum、逐个除，一个函数里完成，只写一个输出张量。要求输入连续。
• Metal：一次 call_last_softmax，每行一个线程组。内核用在线 softmax：每个线程边读边维护 MD { m, d }（当前最大值、按当前最大值缩放过的指数和），线程组里归约合并，再写一遍结果。d 是 float，BF16 / F16 输入也用 F32 累加。

rms_norm、layer_norm、rotary_emb 的 rope / rope_i / rope_thd 都是同一个模式：一个 CustomOp2 / CustomOp3，CPU 一个函数、Metal 一个内核、CUDA 一个内核，旁边留一个 *_slow 拼接版本。sdpa 只有 Metal 实现，CPU 上直接报「SDPA has no cpu impl」。这些都不记反向。

融合和拼接比一比（f32，形状取 SmolLM2 prefill 512 个 token 时的大小）：

• 数值上两者一致到 F32 的舍入误差。CPU 的 rope 和 rope_slow、rope_i 和 rope_i_slow 逐位相同。
• Metal 上融合版本快 9 到 13 倍：拼接版本的每一步都是一次内核调用，加上中间张量的分配。Metal 上融合 softmax 第一次调用用了 23.2 ms，是编译管线（见 Metal 后端）。
• CPU 上 softmax 快 5.5 倍，layer_norm 快 1.5 倍；rms_norm 融合版反而略慢，rope 打平。

rms_norm 在 CPU 上为什么不快，换行数扫一遍（layers rms-sizes，隐藏维 576）：

[Cpu] rms_norm rows 1: fused 0.9 us, composed 1.6 us
[Cpu] rms_norm rows 32: fused 19.5 us, composed 7.4 us
[Cpu] rms_norm rows 33: fused 50.2 us, composed 7.9 us
[Cpu] rms_norm rows 512: fused 118.4 us, composed 78.7 us
[Cpu] rms_norm rows 4096: fused 440.4 us, composed 731.9 us
[Metal { gpu_id: 4294968354 }] rms_norm rows 1: fused 10.6 us, composed 45.9 us
[Metal { gpu_id: 4294968354 }] rms_norm rows 32: fused 6.6 us, composed 79.1 us
[Metal { gpu_id: 4294968354 }] rms_norm rows 33: fused 5.8 us, composed 58.4 us
[Metal { gpu_id: 4294968354 }] rms_norm rows 512: fused 15.2 us, composed 162.6 us
[Metal { gpu_id: 4294968354 }] rms_norm rows 4096: fused 83.5 us, composed 1350.1 us

CPU 的融合实现（RmsNorm 的 cpu_fwd）在 n_rows <= 32 时单线程逐行算，超过 32 行改用 rayon 的 par_chunks。32 行时它比拼接版本慢一倍多（rms_row 里每个元素一次除法一次乘法，逐个写回）；33 行切到 rayon 后反而从 19.5 μs 涨到 50.2 μs，看起来是并行调度的开销比这点计算还大；到 4096 行才比拼接版本快。生成时每步只有 1 行，融合版本 0.9 μs 对 1.6 μs，赢在少分配几个中间张量。Metal 上融合版本在所有规模上都快 4 到 16 倍。

CPU 上的 BF16 softmax

CPU 的融合 softmax 对 BF16 用的是 half::bf16 自己的 vec_reduce_sum。VecOps 的 vec_reduce_sum 默认实现是：

    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)
        }
    }

f32 覆盖了这个方法，bf16 和 f16 没有覆盖，于是累加器本身就是 BF16。BF16 只有 8 位有效位，到 256 之后相邻两个可表示的数相差 2，256 + 1 舍入回 256。拼接的 ops::softmax 里 sum_keepdim 在 CPU 上也走 T::vec_reduce_sum（ReduceSum），同样受影响。

rotary_emb

旋转位置编码（RoPE）有两种配对方式，candle 各有一个融合算子：

• rope_i：交错，第 2i 和 2i+1 维成对旋转。GGUF 里的 Llama、Mistral 是这种。
• rope：不交错，第 i 和 i + d/2 维成对（rope_slow 里的 rotate_half）。Hugging Face 格式的 Llama（SmolLM2 的 config.json 写着 "rope_interleaved": false）、Qwen 是这种。
• rope_thd：和 rope 一样配对，输入布局是 (B, T, H, D) 而不是 (B, H, T, D)，省一次转置。

三者的 cos、sin 都是 (T, D/2)（或带批维的 (B, T, D/2)），内核按行号 0 到 T−1 取，不知道当前位置；调用方要先按位置 narrow。quantized_llama 的 apply_rotary_emb 就是 self.cos.narrow(0, index_pos, seq_len)，再 x.contiguous()，再按 rope_is_neox 选 rope 或 rope_i；rope_is_neox 按 GGUF 的 general.architecture 决定，llama 是 false。三个函数都要求输入连续，否则报「xs has to be contiguous in rope」。

动脑筋答案

  for n in [256usize, 257, 4096] {
    let x = Tensor::zeros((1, n), DType::BF16, dev)?;
    let fused = ops::softmax_last_dim(&x)?.to_dtype(DType::F32)?;
    let composed = ops::softmax(&x, D::Minus1)?.to_dtype(DType::F32)?;
    // print fused[0], fused.sum(), composed[0], composed.sum(), 1/n
  }

[cpu] bf16 zeros len 256: fused elem 0.00390625 sum 1; composed elem 0.00390625 sum 1; exact 0.00390625
[cpu] bf16 zeros len 257: fused elem 0.00390625 sum 1.0039063; composed elem 0.00390625 sum 1.0039063; exact 0.0038910506
[cpu] bf16 zeros len 4096: fused elem 0.00390625 sum 16; composed elem 0.00390625 sum 16; exact 0.00024414063
[metal] bf16 zeros len 256: fused elem 0.00390625 sum 1; composed elem 0.00390625 sum 1; exact 0.00390625
[metal] bf16 zeros len 257: fused elem 0.0038909912 sum 0.99998474; composed elem 0.00390625 sum 1.0039063; exact 0.0038910506
[metal] bf16 zeros len 4096: fused elem 0.00024414063 sum 1; composed elem 0.00024414063 sum 1; exact 0.00024414063

• CPU，4096 个 0：每个 exp(0) = 1，BF16 累加到 256 就加不上去了，分母是 256，每个元素 1/256 = 0.0039，整行加起来是 16，不是 1。融合和拼接两个版本一样错。
• 257 个：分母同样停在 256，每个元素差了 0.4%，行和 1.0039。
• Metal 融合版本用 F32 累加，三种长度都对（257 时差在最后一位，是把结果存回 BF16 的舍入）。Metal 拼接版本 4096 时是对的，说明 Metal 求和内核的累加器比 BF16 宽；257 时却是 1/256，因为 sum_keepdim 的结果要存回 BF16 张量，257 舍入成 256。4096 是 2 的幂，存回去没有误差。
• 全零是最坏的情况（每一项一样大）。真实的注意力分数不会全相等，但只要一行里有几百个量级相近的概率，CPU 上 BF16 的分母就会偏小，每行和大于 1。CPU 上跑 BF16 模型、上下文超过 256 时要当心；本卷的 CPU 实验都用 F32。

小结

• Linear 的权重和 PyTorch 一样是 [out, in]，前向用 w.t() 的视图直接乘；连续的三维输入压成二维再乘，非连续输入走广播 matmul，只慢一点。
• Embedding 就是 index_select；RmsNorm 是不减均值、无偏置的 LayerNorm。
• 连续输入上 RmsNorm::forward、带偏置的 LayerNorm::forward 走融合算子，不记反向，训练要用 forward_diff（只有 RmsNorm 有），它的文档注释和行为相反。
• softmax_last_dim、rms_norm、layer_norm、rope 每个设备一个内核：Metal 上比拼接版本快 4 到 16 倍；CPU 上 softmax 快 5.5 倍，rms_norm 在 32 到 512 行时反而更慢。
• CPU 上 BF16 / F16 的求和用半精度累加，BF16 softmax 在一行超过 256 个相近值时分母偏小，融合和拼接版本都受影响；Metal 的融合 softmax 用 F32 累加，没有这个问题。