取样

主线位置：推理主线的「带 KV 缓存生成」一步里，从 logits 到下一个 token 的那一下。生成循环用的是 argmax，本页换成 do_sample=True 的那几个参数。

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

flowchart TB
  lg["同一组 logits"] --> hf["generate 一步"]
  lg --> cd["LogitsProcessor"]
  hf --> hs["处理后的 scores"]
  hs --> hp["softmax 得 p_hf"]
  cd --> mir["照抄的截断逻辑"]
  mir --> cp["p_candle"]
  cd --> draw["抽 10 万次"]
  draw --> freq["频率"]
  hp --> cmp["逐 token 比"]
  cp --> cmp
  freq --> cmp

实验代码：Python 侧 sampling_ref.py、edge_py.py、seed_py.py、temp0.py；candle 侧 crate gen-exp 的 bin samp、hfsem、edge、seed。release 构建，平台 M2 Max、macOS 27、rustc 1.99.0；samp 每种设置抽 10 万次，跑的时候机器负载 3–7（uptime），只影响耗时。

candle 一侧每种取法在源码里怎么走（哪些在设备上、哪些读回主机、每次多少毫秒）已经在 candle 卷的取样里讲过，本页只对 transformers 的语义。

怎么让两边拿到同一组 logits

取样是随机的，比较抽出来的 token 没有意义（最后一节解释为什么）。能比的是「从同一组 logits 出发，每个 token 被抽中的概率」。

transformers 一侧，让 generate 只走一步，output_logits=True 给出模型原始的 logits，output_scores=True 给出经过温度、top-k、top-p、重复惩罚处理之后、真正拿去 softmax 的 scores。这样拿到的是 generate 自己的处理器链，不是我手搭的：

ctx = d["prompt_ids"] + d["ids"][:k]  # 对话模板 + 已生成的前 k 个 token
ids = torch.tensor([ctx])
SETTINGS = {
  "t07_default": dict(do_sample=True, temperature=0.7),
  "t07": dict(do_sample=True, temperature=0.7, top_k=0),
  "k40": dict(do_sample=True, temperature=0.7, top_k=40),
  "p09": dict(do_sample=True, temperature=0.7, top_k=0, top_p=0.9),
  "k40p09": dict(do_sample=True, temperature=0.7, top_k=40, top_p=0.9),
  "rep13": dict(do_sample=False, repetition_penalty=1.3),
}
tensors = {}
with torch.inference_mode():
  for name, kw in SETTINGS.items():
    g = model.generate(ids, attention_mask=torch.ones_like(ids), max_new_tokens=1,
                       output_scores=True, output_logits=True, return_dict_in_generate=True, **kw)
    tensors.setdefault("logits", g.logits[0][0].contiguous())
    tensors[name] = g.scores[0][0].contiguous()
save_file(tensors, out)

上下文是生成循环里第二个 prompt 的对话模板（39 个 token），加上回答的前 7 个 token In Chinese, you can say "，下一个 token 该是引号里的第一个字，分布比较平。被过滤掉的 token 在 scores 里是 -inf，softmax 之后是 0。

candle 一侧读进同一份 logits。LogitsProcessor 只吐出一个 token，不暴露截断后的分布，sample_topk、sample_topp、sample_topk_topp 又都是私有函数，所以 samp 把它们的逻辑照抄一份，把最后的 sample_multinomial 换成「返回权重」，再用真的 LogitsProcessor 抽 10 万次，确认照抄的分布就是它实际抽样的分布：

/// candle 0.11.0 `sample_f` 里的 `prs(temperature)`
fn prs(logits: &Tensor, t: f64) -> Result<Vec<f32>> {
  Ok(candle_nn::ops::softmax_last_dim(&(logits / t)?)?.to_vec1()?)
}

/// 照抄 `sample_topp`：交给 WeightedIndex 的权重
fn topp_weights(mut prs: Vec<f32>, top_p: f32) -> Vec<f32> {
  let mut idx = (0..prs.len()).collect::<Vec<_>>();
  idx.sort_by(|&i, &j| prs[j].total_cmp(&prs[i]));
  let mut cumsum = 0.;
  for &i in &idx {
    if cumsum >= top_p {
      prs[i] = 0.0;
    } else {
      cumsum += prs[i];
    }
  }
  prs
}
// topk_weights 同样照抄 sample_topk / sample_topk_topp，再把位置映射回词表

let mut lp = LogitsProcessor::from_sampling(42, sampling);
let mut count = vec![0usize; vocab];
for _ in 0..draws {
  count[lp.sample(logits)? as usize] += 1;
}

结果（samp，p_hf 是 transformers scores 的 softmax，p_candle 是照抄逻辑得到的分布，TV 是总变差距离 ½Σ|p_hf − p_candle|）：

t07      kept hf=49152 candle=49152  max|p_hf-p_candle|=6.98e-9  TV=2.14e-8
    LogitsProcessor x100000: max|freq-p_candle|=0.0017  freq outside candle set=0.0000  outside hf set=0.0000
t07_def  kept hf=   50 candle=49152  max|p_hf-p_candle|=1.34e-2  TV=6.06e-2
k40      kept hf=   40 candle=   40  max|p_hf-p_candle|=1.49e-8  TV=5.32e-8
    LogitsProcessor x100000: max|freq-p_candle|=0.0016  freq outside candle set=0.0000  outside hf set=0.0000
p09      kept hf=   31 candle=   31  max|p_hf-p_candle|=1.49e-8  TV=2.68e-8
    LogitsProcessor x100000: max|freq-p_candle|=0.0016  freq outside candle set=0.0000  outside hf set=0.0000
    top-40 original mass = 0.9208, p = 0.9
k40p09   kept hf=   14 candle=   31  max|p_hf-p_candle|=1.85e-2  TV=7.45e-2
    LogitsProcessor x100000: max|freq-p_candle|=0.0016  freq outside candle set=0.0000  outside hf set=0.0748
rep13    context=46 changed candle=34 hf=34 max|diff|=0e0

怎么读：

• max|p_hf-p_candle| 在 1e-8 量级就是同一个分布。两边的 logits 逐位相同，差别只来自两个 softmax 实现（exp 和求和顺序）。
• freq 那一行验证照抄的逻辑：10 万次抽样，每个 token 的频率和 p_candle 最多差 0.0017。概率 0.3 左右的 token，10 万次的标准差是 √(0.3×0.7/10⁵) ≈ 0.0014，这个偏差在一两个标准差之内；抽中的 token 没有一次落在 p_candle 为 0 的地方。
• 只有两行对不上：t07_def 和 k40p09。下面逐个说。

prompt 后面不接回答（k = 0，分布更尖）再跑一遍，结论一样：t07、k40、p09 的最大差 ≤ 5.96e-8，p09 两边都留 16 个；t07_def 留 50 个；k40p09 transformers 留 11 个、candle 留 16 个。

温度

两边做的是同一件事：softmax 之前把 logits 除以 temperature。transformers 的 TemperatureLogitsWarper 是 scores / self.temperature，candle 的 sample_f 是 &logits / temperature 再 softmax_last_dim。上面 t07 一行最大差 6.98e-9。

温度为 0 时两边的反应不同（temp0.py）：

{'do_sample': True, 'temperature': 0.0} -> ValueError `temperature` (=0.0) has to be a strictly positive float, otherwise your next token scores will be invalid. If you're looking for greedy decoding strategies, set `do_sample=False`.
{'do_sample': False, 'temperature': 0.0} -> [7042, 30, 7042]

transformers 在构造 warper 时就拒绝，并提示贪心要用 do_sample=False；do_sample=False 时温度根本不参与（只打印一行「not valid and may be ignored」的提示）。candle 的 LogitsProcessor::new 把小于 1e-7 的温度换成 ArgMax，from_sampling 不检查，温度 0 一路算成 NaN，最后 rand 报「A weight is negative, too large or not a valid number」，见 candle 卷的坑页 A weight is negative。从 transformers 搬过来时，把 do_sample=False 对到 Sampling::ArgMax，不要对到「温度 0」。

do_sample=True 默认带着 top_k=50

t07_def 只写了 do_sample=True, temperature=0.7，transformers 却只留了 50 个 token。5.22.0 的 GenerationConfig._get_default_generation_params() 给没设的字段填默认值，其中：

            "temperature": 1.0,
            "top_k": 50,
            "top_p": 1.0,

_get_logits_processor 只要 top_k 不是 None 也不是 0 就加一个 TopKLogitsWarper。所以 generate(do_sample=True, temperature=0.7) 的语义是「温度 0.7 加 top-50」，对应 candle 的 Sampling::TopK { k: 50, temperature: 0.7 }；对成 Sampling::All { temperature: 0.7 } 时，这一步的分布差 TV = 6.06e-2，candle 会抽到 transformers 永远抽不到的那 49102 个 token。想在 transformers 里关掉它要显式写 top_k=0（本页的 t07、p09 都是这样写的）。模型目录的 generation_config.json 也可以改这些默认值，SmolLM2 的这个文件里没有取样参数。

top-k

transformers 的 TopKLogitsWarper：

        top_k = min(self.top_k, scores.size(-1))  # Safety check
        # Remove all tokens with a probability less than the last token of the top-k
        indices_to_remove = scores < torch.topk(scores, top_k)[0][..., -1, None]
        scores_processed = scores.masked_fill(indices_to_remove, self.filter_value)

candle 的 sample_topk 用 select_nth_unstable_by(top_k, ..) 取出 k 个位置，只在这 k 个里抽。真实 logits 上两边留下同样的 40 个，k40 一行最大差 1.49e-8。

两者只在并列时不同：transformers 删的是「严格小于第 k 名」的，和第 k 名相等的全部留下；candle 正好留 k 个。

top-p

transformers 的 TopPLogitsWarper 从小到大排，删掉「从最小的累加起来不超过 1 − p」的那些：

        sorted_logits, sorted_indices = torch.sort(scores, descending=False)
        cumulative_probs = sorted_logits.softmax(dim=-1).cumsum(dim=-1)

        # Remove tokens with cumulative top_p above the threshold (token with 0 are kept)
        sorted_indices_to_remove = cumulative_probs <= (1 - self.top_p)

candle 的 sample_topp 从大到小排，累加和还没到 p 就收下一个（上面照抄的 topp_weights）。两种写法说的是同一条规则：一个 token 留下，当且仅当比它大的那些加起来还不到 p。真实 logits 上两边都留 31 个（k = 0 时都留 16 个），最大差 1.49e-8。规则一样，只是一个从小端减、一个从大端加，累加和恰好落在 p 上时 f32 舍入可能让两边差一个 token；[0.5, 0.3, 0.2] 取 p = 0.8 这种边界上，两边都留了 2 个。

top-k 加 top-p：要不要重新归一化

这是 k40p09 那一行：transformers 留 14 个，candle 留 31 个，10 万次抽样里有 7.48% 落在 transformers 删掉的 token 上。

transformers 的处理器是一串依次作用在 scores 上的函数，TopKLogitsWarper 把第 40 名之后填成 -inf，TopPLogitsWarper 接着在这个结果上做 softmax、累加。被填成 -inf 的 token 概率是 0，前 40 个自动重新归一化，p = 0.9 量的是「前 40 个里的 90%」。

candle 的 sample_topk_topp 取出前 k 个之后不归一化：

            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 量的是原分布上的概率质量。这一步前 40 个的原始质量是 0.9208，刚过 0.9，于是 candle 的截断点就是「整个词表上的 90%」，和单独 TopP { p: 0.9 } 一样留 31 个；transformers 在 40 个里算 90%，要求更高的集中度，只留 14 个。前 40 个的原始质量不到 0.9 时，candle 走 top_p >= sum_p 那一支，p 完全不起作用。

要在 candle 里得到 transformers 的语义，不用改库：sample_f 的闭包拿到的是 softmax 之后、截断之前的概率，在闭包里做 top-k 并重新归一化，再交给 TopP（hfsem）：

fn hf_topk_renorm(prs: &mut [f32], k: usize) {
  let mut idx: Vec<usize> = (0..prs.len()).collect();
  idx.select_nth_unstable_by(k, |&i, &j| prs[j].total_cmp(&prs[i]));
  for &i in &idx[k..] {
    prs[i] = 0.0;
  }
  let s: f32 = prs.iter().sum();
  prs.iter_mut().for_each(|p| *p /= s);
}

let mut lp = LogitsProcessor::from_sampling(42, Sampling::TopP { p: 0.9, temperature: 0.7 });
let next = lp.sample_f(&logits, |prs| hf_topk_renorm(prs, 40))?;

TopP + top-k renorm in sample_f x100000: tokens drawn=14 max|freq-p_hf|=0.0015 outside hf set=0.0000

10 万次只抽到那 14 个，频率和 transformers 的分布最多差 0.0015。这个闭包没有处理并列（上一节），真实 logits 上碰不到。

重复惩罚

公式两边一样。transformers 的 RepetitionPenaltyLogitsProcessor：

        score = torch.gather(scores, 1, input_ids)
        # if score < 0 then repetition penalty has to be multiplied to reduce the token probabilities
        score = torch.where(score < 0, score * self.penalty, score / self.penalty)
        scores_processed = scores.scatter(1, input_ids, score)

candle 的 apply_repeat_penalty 是 if *logit >= 0. { *logit /= penalty } else { *logit *= penalty }，出现几次都只罚一次。transformers 的 gather / scatter 对重复的 id 也只算一次：重复位置读到的都是原值，写回的是同一个结果。logit 恰好是 0 时两边都走除法，结果还是 0。上面 rep13 一行：同样 46 个 token 的上下文，两边改动的都是 34 个 logit，逐位相同（max|diff|=0e0）。

不同在罚谁：

• transformers 罚的是传进 generate 的整个 input_ids，包括 prompt，也包括对话模板里的系统提示词。处理器有一个 prompt_ignore_length 参数可以跳过 prompt，但 generate 构造它时不传，要自己建处理器传进 logits_processor。
• candle 罚的是你传进去的 context。candle 的 llama 例子传最近 --repeat-last-n 个 token（默认 128），quantized 例子默认 64。

上下文变长之后差别就出来了。同一个对话，取回答的前 60 个 token，上下文共 99 个（samp，k = 60）：

rep13    context=99 changed candle=60 hf=60 max|diff|=0e0
rep13    last 64 of context: changed=39 max|diff vs hf|=4.663518e0

传整个上下文时两边还是逐位相同；照 quantized 例子只传最近 64 个，被罚的 token 从 60 个变成 39 个，那 21 个只在前面出现过的 token 不再被罚（其中 16 个来自系统提示词，ĠAI、Ġhelpful、ĠFace 之类），单个 logit 最多差 4.66。从 transformers 搬过来要对齐结果，就把整个 input_ids 传给 apply_repeat_penalty；要对齐 candle 例子的行为，就知道它们和 transformers 的默认不是一回事。

两边的顺序也一样：transformers 的重复惩罚在处理器链里排在温度和 top-k / top-p 之前，candle 的例子先调 apply_repeat_penalty 再交给 LogitsProcessor。

同一个种子，两边抽到的不一样

每一步的分布对上了，抽出来的 token 照样不同（seed_py.py、seed，同一个 prompt，温度 0.7、不带 top-k，20 个 token）：

torch seed 42: '您好！🙏🎉🎉🎉�'
torch seed 42: '您好！🙏🎉🎉🎉�'
torch seed 43: '"In Chinese, it is said to say "您好？你�'

candle seed 42: "嗯？？？\n\n您的希望�"
candle seed 42: "嗯？？？\n\n您的希望�"
candle seed 43: "Hello! Thanks for the offer. I'm so glad they're here. A simple way to say"

每一边自己是可复现的：同一个种子两次逐字相同。两边之间不同，原因有两层：

• 发生器不同。 torch.manual_seed 设的是 PyTorch 的 CPU 发生器；candle 的 LogitsProcessor 里是 rand 0.9 的 StdRng（ChaCha12），from_sampling(seed, ..) 用 seed_from_u64 初始化。同一个整数进去，出来的是两串毫无关系的随机数。

• 抽法不同，消耗的随机数个数也不同。 transformers 的 _sample 对处理后的分布调 torch.multinomial(probs, num_samples=1)。抽一个样本时，torch.multinomial 等价于给每个 token 一个指数分布的随机数，取 probs / 噪声 的最大值：

  Copy

  multinomial == argmax(p / Exp(1)) for 1000/1000 seeds
  

  1000 个种子，先 torch.multinomial，重设种子后再用 torch.empty_like(probs).exponential_(1) 手算，结果每次都相同。也就是说 PyTorch 每抽一个 token 要用掉 49152 个随机数。candle 的 WeightedIndex 建一张累加表，抽一个 [0, 总和) 里的均匀随机数，二分查找它落在哪一格，每个 token 只用一个随机数。即使把两边的发生器换成同一个，第二步开始它们读的也不是随机数流里的同一段。

所以取样的对照只能比分布（本页前面几节），不能比 token。要两边逐 token 相同，只能都用 argmax。candle 一侧种子能管什么、管不到什么（GumbelSoftmax 不用这个种子、Metal 的 set_seed），见 candle 卷的取样和坑页设了种子结果还是不一样。

动脑筋答案

edge_py.py 和 edge：

top_k=2 on [3,2,2,1]: [True, True, True, False]

TopK k=2 on [3,2,2,1], 10000 draws: [7349, 2651, 0, 0]

• transformers 留下 3 个：第 2 名的值是 2，删掉的是严格小于 2 的，两个 2 都留下。三个 token 按 e³ : e² : e² 分。
• candle 只留 2 个，抽中位置 0 和 1：select_nth_unstable_by 把第 k 个位置放好之后，左边正好 k 个，两个并列的 2 里哪个进左边由排序算法决定，这里是位置 1。位置 0 的频率 0.7349，和 e³ / (e³ + e²) = 0.731 一致。
• 并列在 f32 logits 上很少见，但量化模型、被 -inf 屏蔽过的 logits 里会出现成片相同的值。

小结

• 比较取样不比 token，比分布：generate(max_new_tokens=1, output_scores=True, output_logits=True) 给出 transformers 实际用的 scores；candle 照抄截断逻辑，再用真的 LogitsProcessor 抽 10 万次验证。
• 温度、top-k、top-p、重复惩罚的公式两边相同：同一组 logits 上概率最大差 ≤ 6e-8，重复惩罚逐位相同。
• do_sample=True 默认带 top_k=50，对应 Sampling::TopK { k: 50, .. } 而不是 All；do_sample=False 对应 ArgMax，不要用温度 0。
• top-k 加 top-p：transformers 在前 k 个里重新归一化再算 p，candle 的 TopKThenTopP 用原分布的质量，这一步留 14 个对 31 个；在 sample_f 的闭包里做 top-k 加归一化，再用 TopP，就是 transformers 的语义。
• top-k 并列时 transformers 全留、candle 正好留 k 个。
• 重复惩罚 transformers 罚整个 input_ids（含 prompt），candle 罚你传的 context，例子默认最近 64 或 128 个。
• 同一个种子两边抽到不同 token：发生器不同，torch.multinomial 每步用 49152 个指数随机数，WeightedIndex 只用一个。