Rust in Everything
基于 candle-core 0.11.0

取样

主线位置：推理主线的最后一步。ModelWeights::forward 返回最后一个位置的 logits（形状 (1, vocab)，squeeze(0) 后是 (vocab,)），LogitsProcessor::sample 把它变成一个 u32，作为下一轮 forward 的输入。生成一个 token 用的是贪心的 argmax，本页讲其余几种取法。代码在 candle-transformers/src/generation/mod.rs（158 行）和 utils.rs 的 apply_repeat_penalty。

flowchart TB
  lg["logits vocab"] --> f32["to_dtype F32"]
  f32 -->|ArgMax| am["argmax 在设备上"]
  f32 -->|Gumbel| gs["加噪声后 argmax"]
  f32 -->|其余四种| sm["除以温度 softmax"]
  sm --> rb["to_vec1 读回主机"]
  rb --> cut["top-k 或 top-p 截断"]
  cut --> wi["WeightedIndex 加 StdRng"]
  am --> tok["下一个 token"]
  gs --> tok
  wi --> tok

本页的实验在 models-exp 的 sample 里（release），模型是 SmolLM2-135M-Instruct 的 Q8_0 GGUF，平台 M2 Max、macOS 27、rustc 1.99.0。

六种 Sampling

pub enum Sampling {
    ArgMax,
    All { temperature: f64 },
    TopK { k: usize, temperature: f64 },
    TopP { p: f64, temperature: f64 },
    TopKThenTopP { k: usize, p: f64, temperature: f64 },
    // Note that the rng is not used for the Gumbel-Softmax sampling.
    GumbelSoftmax { temperature: f64 },
}

pub struct LogitsProcessor {
    rng: rand::rngs::StdRng,
    sampling: Sampling,
}

LogitsProcessor 只有两个字段：取法和一个随机数发生器。构造有两条路：

  • from_sampling(seed, sampling) 原样收下 Sampling，用 StdRng::seed_from_u64(seed) 建发生器。
  • new(seed, temperature, top_p) 是老接口：温度是 None 或小于 1e-7 时变成 ArgMax，否则有 top_p 就是 TopP，没有就是 All。它造不出 TopK、TopKThenTopP 和 GumbelSoftmax。

真正干活的是 sample_f：

    pub fn sample_f(&mut self, logits: &Tensor, f: impl FnOnce(&mut [f32])) -> Result<u32> {
        let logits = logits.to_dtype(DType::F32)?;
        let prs = |temperature: f64| -> Result<Vec<f32>> {
            let logits = (&logits / temperature)?;
            let prs = candle_nn::ops::softmax_last_dim(&logits)?;
            let mut prs = prs.to_vec1()?;
            f(&mut prs);
            Ok(prs)
        };

        let next_token = match &self.sampling {
            Sampling::ArgMax => self.sample_argmax(logits)?,
            Sampling::GumbelSoftmax { temperature } => {
                self.sample_gumbel_softmax(&logits, *temperature)?
            }
            Sampling::All { temperature } => {
                let prs = prs(*temperature)?;
                self.sample_multinomial(&prs)?
            }
            // ... TopP、TopK、TopKThenTopP 同样先 prs(*temperature)
        };
        Ok(next_token)
    }

分成两类：

取法在哪算用不用 self.rng
ArgMax设备上 argmax，to_scalar 读回一个 u32不用
GumbelSoftmax设备上加 Gumbel 噪声再 argmax（candle_nn::sampling::gumbel_softmax）不用，噪声来自设备的随机数
All / TopK / TopP / TopKThenTopP设备上除以温度、softmax_last_dim，to_vec1 把整行概率读回主机，截断和抽样都在主机上用，WeightedIndex 抽一次

sample_f 的第二个参数 f 拿到的是读回主机后的概率向量，可以在截断前改它（例如屏蔽某些 token）；sample 就是传一个空闭包。

logits 必须是一维的。forward 返回的是 (1, vocab)，忘了 squeeze(0) 时（rank）：

ArgMax on [1, 6]: Err("unexpected rank, expected: 0, got: 1 ([1])")
All { temperature: 0.8 } on [1, 6]: Err("unexpected rank, expected: 1, got: 2 ([1, 6])")
GumbelSoftmax { temperature: 0.8 } on [1, 6]: Err("unexpected rank, expected: 0, got: 1 ([1])")
apply_repeat_penalty on [1, 6]: Err("unexpected rank, expected: 1, got: 2 ([1, 6])")

报错来自 to_scalar / to_vec1，说的是秩，不提 logits。

温度

温度只做一件事：softmax 之前把 logits 除以 temperature。小于 1 时分布变尖，大于 1 时变平。拿一组固定的 logits [2.0, 1.5, 1.0, 0.5, 0.0, -1.0]，每种设置抽 20 万次（dist）：

softmax: ["0.4197", "0.2546", "0.1544", "0.0936", "0.0568", "0.0209"]
All t=1                ["0.4209", "0.2540", "0.1539", "0.0936", "0.0566", "0.0211"]
All t=0.5              ["0.6353", "0.2341", "0.0856", "0.0319", "0.0117", "0.0015"]

t = 1 时频率和 softmax 一致到小数点后两三位；t = 0.5 等于把 logits 乘 2，最大的那个从 42% 涨到 64%。

温度为 0 只有 new 会替你换成 ArgMax。直接用 from_sampling 传 0（zero）：

new(.., Some(0.0), None): Ok(0)
All { temperature: 0.0 }: Err("A weight is negative, too large or not a valid number")
TopP { p: 0.9, temperature: 0.0 }: Err("A weight is negative, too large or not a valid number")
TopK { k: 2, temperature: 0.0 }: Err("A weight is negative, too large or not a valid number")
All { temperature: 0.001 }: Ok(0)

除以 0 得到 ±inf（logit 恰好是 0 时得到 NaN），softmax 里 inf - inf 是 NaN，整行概率都成了 NaN，rand 的 WeightedIndex::new 拒绝它。报错文字来自 rand，不提温度。很小但不是 0 的温度没问题。GumbelSoftmax 自己处理了这一点：temperature <= 0.0 时直接 argmax。

top-k、top-p 和两者连用

sample_topp：

    fn sample_topp(&mut self, prs: &mut Vec<f32>, top_p: f32) -> Result<u32> {
        let mut argsort_indices = (0..prs.len()).collect::<Vec<_>>();

        // Sort by descending probability.
        argsort_indices.sort_by(|&i, &j| prs[j].total_cmp(&prs[i]));

        // Clamp smaller probabilities to zero.
        let mut cumsum = 0.;
        for index in &argsort_indices {
            if cumsum >= top_p {
                prs[*index] = 0.0;
            } else {
                cumsum += prs[*index];
            }
        }
        // Sample with clamped probabilities.
        self.sample_multinomial(prs)
    }

从大到小累加，累加和还没到 p 时就把下一个收进来，到了之后的全部置零。收进来的是「累加和刚好达到 p 的最短前缀」，最后一个收进来的 token 会让累加和越过 p。置零后不归一化，WeightedIndex 按剩下的权重比例抽，效果等于归一化。TopP 在 p <= 0 或 p >= 1 时不截断，退化成 All。

sample_topk 不全排序：select_nth_unstable_by(top_k, ..) 把最大的 k 个换到前面（它们之间无序），只在这 k 个里抽，再把位置映射回 token id。

sample_topk_topp 先取前 k 个，再对这 k 个调 sample_topp，但传进去的是没有重新归一化的概率：

            let mut prs = indices.iter().map(|&i| prs[i]).collect::<Vec<_>>();
            let sum_p = prs.iter().sum::<f32>();
            let index = if top_p <= 0.0 || top_p >= sum_p {
                self.sample_multinomial(&prs)?
            } else {
                self.sample_topp(&mut prs, top_p)?
            };

所以 p 量的是原分布里的概率质量，不是 top-k 之内的份额。同一组 logits（dist）：

TopK k=2               ["0.6228", "0.3772", "0.0000", "0.0000", "0.0000", "0.0000"]
TopP p=0.7             ["0.5066", "0.3076", "0.1859", "0.0000", "0.0000", "0.0000"]
TopP p=0.8             ["0.5066", "0.3076", "0.1859", "0.0000", "0.0000", "0.0000"]
TopKThenTopP k=3 p=0.6 ["0.6228", "0.3772", "0.0000", "0.0000", "0.0000", "0.0000"]
TopKThenTopP k=3 p=0.75 ["0.5066", "0.3076", "0.1859", "0.0000", "0.0000", "0.0000"]
TopKThenTopP k=3 p=0.9 ["0.5066", "0.3076", "0.1859", "0.0000", "0.0000", "0.0000"]

逐行核对：

  • TopK k=2：0.4197 和 0.2546 按比例分，0.4197 / 0.6743 = 0.622。
  • TopP p=0.7 和 p=0.8 一样：前两个累加 0.6743，还不到 0.7，第三个也收进来，累加到 0.8287，同时越过了 0.7 和 0.8。三个按比例分：0.4197 / 0.8287 = 0.506。
  • TopKThenTopP k=3 p=0.75：前 3 个的原始概率累加到第二个是 0.6743，没到 0.75，第三个也收进来。如果先在前 3 个里归一化（0.506、0.307、0.186），前两个已经是 0.813，第三个就会被截掉。candle 的做法是前一种。
  • k=3 p=0.9：0.9 大于前 3 个的总和 0.8287，走 top_p >= sum_p 那一支，不截断。
📝NOTE
没有蠢问题

问： TopKThenTopP { k: 40, p: 0.9 } 和先 TopK 再 TopP 是一回事吗？

答： 不完全是。p 对照的是整个词表上的概率质量。前 40 个加起来不到 0.9 时（分布很平时常见），top_p >= sum_p 成立，p 不起作用，等于只做了 top-k；加起来超过 0.9 时，截断点按原始概率算，比「在 40 个里归一化后再取 90%」收进来的 token 更多或一样多，不会更少。

每次取样要多久

sample_topp 每一步都把整个词表排一次序。SmolLM2 的词表是 49152，同一组 logits 连续取样 200 次的平均（cost，release；温度 0.8、k = 40、p = 0.9）：

取法logits 在 CPU 上logits 在 Metal 上
ArgMax0.059 ms0.22–0.36 ms
All0.155–0.168 ms0.30–0.47 ms
TopK0.195 ms0.32–0.47 ms
TopP1.23–1.28 ms1.67–1.72 ms
TopKThenTopP0.197–0.201 ms0.38–0.68 ms
GumbelSoftmax0.44–0.46 ms0.45–0.60 ms

两次运行的范围；机器上同时有别的负载，Metal 一列波动较大。

  • 在 CPU 上的 logits：argmax 0.06 ms；All 要做 softmax 再在 49152 个权重上建 WeightedIndex（累加表），约 0.16 ms；TopP 多一次全排序，1.2 ms 以上。TopKThenTopP 的排序只在 40 个里做，又回到 0.2 ms。
  • 在 Metal 上的 logits 每种都多一次提交和读回（argmax 读一个 u32，其余读 49152 个 f32），各多出零点几毫秒。
  • Metal 上热缓存时一个 decode 步约 2.5 ms（见生成一个 token），TopP 一次取样就占它的一半以上。要 top-p 又在意速度，TopKThenTopP 加一个足够大的 k 能把排序的规模先降下来。

种子：管得到什么

All、TopK、TopP、TopKThenTopP 的随机性全来自 self.rng，即 rand 0.9 的 StdRng（ChaCha12）。同一个种子、同一个模型、同一个 prompt，两次生成逐 token 相同（seed，CPU，All { temperature: 0.8 }，25 个 token）：

seed 42: "The misty dawn barely penetrates the sleepy way of the sea, beckoning the speaker, creaming with the salty"
seed 42: "The misty dawn barely penetrates the sleepy way of the sea, beckoning the speaker, creaming with the salty"
seed 43: "In the vast expanse of the sea, where the spray danced on the waves, I find myself lost in a world of wonder"

每取样一次，WeightedIndex 都从发生器里取随机数，所以「同一个种子」复现的是整串抽取，前提是每一步的概率向量都一样：换了 dtype、换了量化格式，或者中途多调了一次 sample，后面的序列就会分岔。rand 的文档还写明 StdRng 不保证跨版本可移植，升级 rand 之后同一个种子可以给出不同的序列。

同一个 prompt、同一个种子 299792458，几种取法的前 25 个 token（text，CPU）：

ArgMax            "The sea, a vast and mysterious realm, beckons sailors to explore its depths and uncover its secrets.\n"
All t=0.8         "The deep blue of the ocean was an otherworldly sight, like nothing I'd ever seen before. Everywhere I looked"
All t=1.5         "##  Honpoaatio AT Mallorca Sarea Mayfelidas Near Bella afflictatica erab Repeter Foredo"
TopK k=40 t=0.8   "The deep blue of the ocean, its surface painted with hues of sapphire and turquoise, stretches out as far as the eye"
TopP p=0.9 t=0.8  "The deep blue of the ocean is a sea of wonder and mystery. It beckons sailors to set sail, their hearts"

后三行开头几个词相同：三者的发生器给出同一串随机数，前几步的概率集中在头部，截掉的只是尾部的小概率 token，同一个随机数抽中的还是同一个 token；等到某一步截断改变了被抽中的那一段，三条路就分开了。温度 1.5 时，135M 的小模型几步之后就离开了正常的词。

❗IMPORTANT
动脑筋：同样是 LogitsProcessor::from_sampling(42, ..)，在同一组 logits 上连抽 16 次，Sampling::All { temperature: 1.0 } 两次结果相同。换成 Sampling::GumbelSoftmax { temperature: 1.0 }，在 CPU 上两次结果相同吗？在 Metal 上呢？

gumbel_softmax 的噪声来自 logits.rand_like(1e-7, 0.999)，也就是设备的随机数发生器，和 LogitsProcessor 的种子无关：

        let logits = logits.to_dtype(candle::DType::F32)?;
        let minus_g = logits.rand_like(1e-7, 0.999)?.log()?.neg()?.log()?;
        if temperature == 1.0 {
            let sampled = (logits - minus_g)?.argmax(dim)?;
            Ok(sampled)
        } else {
            let sampled = (logits + minus_g * (-temperature))?.argmax(dim)?;
            Ok(sampled)
        }

argmax(logits + Gumbel 噪声 × T) 和按 softmax(logits / T) 抽样是同一个分布，但整个过程留在设备上，不用把整行概率读回主机。

📝NOTE
冷知识：Gumbel-Softmax 一进来就不用种子

GumbelSoftmax 是 PR #2894（「Gumbel-Softmax sampling.」，首次进入 0.9.0-alpha.3，正式版 0.9.0）加的，提交说明只有三条：「Gumbel-Softmax sampling.」「Add a sampling test.」「Share the gumbel-softmax bits.」。之前 Sampling 只有五种，全靠 self.rng；这次在 candle-nn 新建了 sampling.rs，gumbel_softmax 写在那里，Sampling 多了一个变体，旁边那行注释「Note that the rng is not used for the Gumbel-Softmax sampling.」也是这次加的。注释说的就是上面的动脑筋：LogitsProcessor 的种子对它不起作用。

Metal 的 set_seed：种子的值被忽略

想让 Gumbel 可复现，只能给设备设种子。CPU 不行（Device::Cpu.set_seed 返回「cannot seed the CPU rng with set_seed」，见优化器）。Metal 可以设，但设什么值都一样（seedprobe，release）：

fresh device:            [0.0030706443, 0.3755374, 0.676627, 0.69485384]
second draw:             [0.037836768, 0.99889874, 0.879909, 0.42845085]
set_seed(       42):       [0.0030706443, 0.3755374, 0.676627, 0.69485384]
set_seed(        7):       [0.0030706443, 0.3755374, 0.676627, 0.69485384]
set_seed(299792458):       [0.0030706443, 0.3755374, 0.676627, 0.69485384]
set_seed(       42):       [0.0030706443, 0.3755374, 0.676627, 0.69485384]
set_seed(4294967296):       [0.16422416, 0.6684503, 0.31429586, 0.50046116]
set_seed(21474836480):       [0.13277194, 0.124092355, 0.924061, 0.028893871]

每次 Tensor::rand(0, 1, 4)。新建的设备、set_seed(42)、set_seed(7) 给出同一串数；只有 2³² 以上的种子才不同。原因在主机和内核对这 8 个字节的理解不一致：MetalDevice::set_seed 把 u64 按小端序拷进缓冲区（低 32 位在前），random.metal 的 atomic_load_seed 却把前 4 个字节当高 32 位：

METAL_FUNC ulong atomic_load_seed(device seed_buffer *sb) {
    uint x = atomic_load_explicit(&sb->seed[0], memory_order_relaxed);
    uint y = atomic_load_explicit(&sb->seed[1], memory_order_relaxed);
    return static_cast<ulong>(x) << 32 | y;
}

种子小于 2³² 时，内核读到的是 seed << 32，低 32 位全是 0。seed_per_thread 把它乘上常数再截成 uint，只留低 32 位，结果和种子无关。内核自己写回的下一个种子按同样的顺序存取，所以同一设备上第二次抽取会变，看起来像正常的随机数。设备初始的种子 299792458 也小于 2³²，于是「新设备」「set_seed 任意小种子」落在同一串上。

这同时解释了优化器一页的观察（Metal 上两次训练每一步的 loss 相同）：可复现是因为每个新设备都从同一个状态开始，不是因为种子起了作用。

重复惩罚

apply_repeat_penalty 不在 LogitsProcessor 里，调用方在取样前自己调：

pub fn apply_repeat_penalty(logits: &Tensor, penalty: f32, context: &[u32]) -> Result<Tensor> {
    let device = logits.device();
    let mut logits = logits.to_dtype(candle::DType::F32)?.to_vec1::<f32>()?;
    let mut already_seen = std::collections::HashSet::new();
    for token_id in context {
        if already_seen.contains(token_id) {
            continue;
        }
        already_seen.insert(token_id);
        if let Some(logit) = logits.get_mut(*token_id as usize) {
            if *logit >= 0. {
                *logit /= penalty
            } else {
                *logit *= penalty
            }
        }
    }
    let logits_len = logits.len();
    Tensor::from_vec(logits, logits_len, device)
}
  • 出现过的 token，正的 logit 除以 penalty，负的乘以 penalty，两种情况都让它变小。
  • 出现几次都只罚一次（HashSet 去重）；超出词表的 id 被 get_mut 静默跳过。
  • 整行 logits 读回主机、改完再建一个新张量放回原设备。(logits, 2.0, [0, 1, 1, 1, 7]) 作用在 [2, -2, 0.5] 上得到 [1.0, -4.0, 0.5]。
  • 读回再上传在 Metal 上有代价：49152 个 logits、64 个 token 的上下文，CPU 上 0.007 ms，Metal 上约 0.25 ms（cost）。

context 传多少由调用方决定，candle-examples 的 quantized 例子默认 --repeat-penalty 1.1、--repeat-last-n 64，即最近 64 个 token。贪心生成 80 个 token，罚与不罚（penalty，CPU）：

no penalty: "Here are some common colors:\n\n* Red: A strong, fiery, and vibrant color that is often associated with passion, energy, and excitement.\n* Blue: A deep, soothing, and calming color that is often associated with trust, stability, and calmness.\n* Green: A versatile color that can be used for everything from nature to abstract art, and is often associated with growth"
penalty 1.3, last 64: "Here are some common colors:\n1. Red - This is a popular choice for many people, as it's often associated with fire and passion in the outdoors or when you're excited about something exciting happening (like sports games). \n2. Blue - Used to represent calmness; blue can also be used on walls because of its calming properties that help reduce stress levels by reducing anxiety caused from bright lights"

不罚时每一项都是「is often associated with」的同一个句式；罚 1.3 之后连 * 和换行这种格式 token 也被压低，列表换成了 1.、2.，句子越拉越长。惩罚不分内容词和格式词，penalty 和窗口都要按模型试。

动脑筋答案

gumbel，同一组 6 个 logits，每次新建 LogitsProcessor::from_sampling(42, ..) 连抽 16 次，CPU：

gumbel seed 42: [1, 1, 2, 0, 0, 0, 2, 0, 0, 1, 2, 0, 1, 4, 0, 0]
gumbel seed 42: [0, 2, 1, 0, 0, 0, 3, 5, 2, 0, 5, 1, 5, 0, 0, 0]
all    seed 42: [0, 1, 0, 1, 3, 1, 5, 0, 4, 0, 1, 0, 0, 2, 0, 3]
all    seed 42: [0, 1, 0, 1, 3, 1, 5, 0, 4, 0, 1, 0, 0, 2, 0, 3]
Device::Cpu.set_seed(42): Some("cannot seed the CPU rng with set_seed")

Metal（新进程，新设备）：

gumbel seed 42: [1, 1, 0, 4, 0, 0, 1, 2, 3, 3, 0, 0, 3, 0, 2, 0]
gumbel seed 42: [0, 2, 2, 1, 2, 0, 0, 0, 1, 1, 2, 4, 1, 3, 0, 3]
all    seed 42: [0, 1, 0, 1, 3, 1, 5, 0, 4, 0, 1, 0, 0, 2, 0, 3]
all    seed 42: [0, 1, 0, 1, 3, 1, 5, 0, 4, 0, 1, 0, 0, 2, 0, 3]
gumbel seed 42 + device.set_seed(42): [1, 1, 0, 4, 0, 0, 1, 2, 3, 3, 0, 0, 3, 0, 2, 0]
gumbel seed 42 + device.set_seed(42): [1, 1, 0, 4, 0, 0, 1, 2, 3, 3, 0, 0, 3, 0, 2, 0]
  • All 两次相同，CPU 和 Metal 也相同：随机数来自 StdRng，概率向量是同一组 logits 的 softmax。
  • GumbelSoftmax 在 CPU 上两次不同：噪声来自 CPU 的 rand_uniform，它每次用 rand::rng()，不能设种子。
  • Metal 上两次也不同：第二次时设备的种子缓冲区已经被第一次推进了。但新设备上的第一次是固定的；每次先 device.set_seed(42) 就能复现，而且正如上一节所说，复现出来的就是新设备那一串，42 这个值本身没起作用。

小结

  • LogitsProcessor = 一种 Sampling + 一个 StdRng。ArgMax 和 GumbelSoftmax 在设备上算、只读回一个 u32；其余四种把整行概率读回主机，截断和抽样在主机上做。
  • 温度是 softmax 前的除数；只有 new 把温度 0 换成 ArgMax，from_sampling 传 0 会得到 NaN 和 rand 的「A weight is negative, too large or not a valid number」。
  • TopP 收下累加和刚好达到 p 的最短前缀；TopKThenTopP 的 p 量的是原分布的概率质量，不在 top-k 里重新归一化。
  • TopP 每步把整个词表排序，49152 个 token 要 1.2 ms 以上，比 Metal 上一个 decode 步的一半还多；TopKThenTopP 先缩到 k 个再排。
  • 同一个种子复现整串抽取，前提是每一步的概率完全相同；Gumbel 不用这个种子，CPU 上不可复现。
  • Metal 的 set_seed 有字节序 bug：小于 2³² 的种子全部等价于默认种子。
  • apply_repeat_penalty 去重后把出现过的 token 压低一次，要在取样前自己调，整行 logits 在主机上走一个来回。

💬 关联讨论

发起讨论
加载中...