nano-vllm-omni · 提示 ← P06P09 →
#18 · Track B · Stage 进程与调度 · 书 p.80

P08 · 批不变的采样器（per-request seed）

talker 是唯一随机的 stage，vllm-omni 给它固定 seed 保证可复现；但如果同一个 seed 的结果随 batch 里别人变化，那 seed 就毫无意义。

⏱ 45 min 难度 中 实现 mine/p08_sampler.py 验收 ./check.sh p08 提示 A-18 · 参考解 C-18 solutions/p08_sampler.py

题面速览

1. repetition penalty：对出现过的每个 token，正 logit 除以 p，负 logit 乘以 p
2. temperature < 1e-5 → argmax，结束
3. logits /= T
4. top_k：只保留最大的 k 个
5. top_p：降序排序后，保留累计概率达到 top_p 的最小前缀（至少 1 个）
6. 用该请求自己的 torch.Generator 抽一个 token

验收要点：greedy · top_k=1 / 极小 top_p 等价于 greedy · 同 seed 的请求在不同 batch 组成/顺序下结果完全一致 · top_k=2 的经验分布 ≈ 0.7/0.3 · repetition penalty 数值与效果 · free 后重新开始随机流。

直觉与类比

选秀：扣分 → 放大分差 → 入围名单 → 抽签

一行 logits 就是 V 个选手的评委打分。① repetition penalty：刚唱过的歌手扣分；② temperature：T<1 放大分差（更稳），T>1 压平（更有创意也更容易胡说），T→0 直接点第一名；③ top_k：只留分数最高的 k 个；④ top_p：按分数从高到低累计“得票率”，够 p 了就封闭名单（动态的 k）；最后在名单里按概率抽签——抽签这一下才是唯一用到随机数的地方。

公用摇号机 vs 每人自带骰子

公用摇号机（整 batch 一个 generator）：排在你前面的人每摇一次，机器就往前走一格。你今天排第 3、明天排第 1，摇到的号自然不同。自带骰子（per-request generator）：掷了几次只取决于你被调度过几次。骰子要一直带着（跨 step 复用，不是每步换新的），用完要还掉（free，否则下一个同名请求会接着你的随机流）。

框架视角 · shape trace

步骤 张量 / 形状 做什么 易错点
0 logits [B, V]，可能 bf16/fp16 逐行取出 logits[i] 并 .float() 直接在 bf16 上 softmax / cumsum，与参考差 1e-2
1 row [V] fp32 apply_repetition_penalty：对去重后的历史 token，v > 0 ? v/p : v*p ① 重复 token 罚多次；② 负 logit 也做除法（反而变大 = 奖励）；③ 原地改了调用方的张量
2 row [V] temperature < 1e-5 → argmax；否则 row / temperature 先除后判断 → /0 得 inf/NaN；在 penalty 之前判断 greedy → penalty 对 greedy 行失效
3 row [V] 部分 -inf 仅当 0 < top_k < V：第 k 大的值 kth，row < kth 置 -inf top_k = -1 或 ≥ V 没当成不启用；并列第 k 名会多保留几个（可接受）
4 sorted [V]、order [V] 降序 → softmax → cumsum；“之前的累计质量” = cum − probs；已超过 top_p 的位置置 -inf；强制保留第 0 位 把“越过阈值的那个 token”也删掉 → 极小 top_p 全删光 → NaN；忘了用 order 映射回原位置
5 probs [V] → 标量 torch.multinomial(probs, 1, generator=self._gen(req_id, p)) 整批共用一个 generator；每步重建 generator；忘了 free
6 out [B] long 写回第 i 个位置 返回 float / [B,1]；行顺序与 req_ids 不一致
不变量 怎么验
I1 批不变：固定 seed 的请求连续 N 步的 token 序列与它“单独成批”时逐个相同 塞进随机打乱的 batch 跑 N 步，和单独跑 N 步比较
I2 无副作用：sample() 返回后传入的 logits 逐位不变 before = logits.clone()，调用后 torch.equal
I3 temperature < 1e-5、top_k == 1、top_p → 0 都退化成 argmax 三种参数各跑一次
I4 保留集合非空，无 NaN top_p=1e-9、top_k=1 等极端值
I5 随机流连续：第 n 次抽样用 generator 的第 n 个随机数；free 后重置 抽 4 次 → free → 再抽 4 次，两段相同
I6 经验分布 ≈ 过滤后的概率 同一行抽几千次统计频率

易错点

bug 症状
重复 token 惩罚多次 例 1 中 token 3 变成 -0.45 而不是 -0.3
负 logit 也做除法 被惩罚的负分 token 反而更容易被选中
top_p 删掉越过阈值的那个 token top_p=0.7 时只剩 {0}；极小 top_p 时全被删 → NaN
所有行共用一个 generator 例 3 的“混合”结果与“单独”不一致
每步都 manual_seed 新建 generator 同一请求每一步都抽到同一个随机数，分布明显偏
free 后没删 generator 新请求复用旧 id 时接着旧随机流
对整个 [B, V] 调一次 multinomial 各行按顺序消费同一个 generator，批不变性没了

逐步图：top_k=3, top_p=0.6, T=1（原概率 [0.5, 0.25, 0.15, 0.1]）

▶ 一行 logits 的过滤
1. logits = log(p)

-0.693-1.386-1.897-2.303

2. top_k=3

第 3 大的值 = -1.897，小于它的置 -inf：

-0.693-1.386-1.897-inf

3. 降序 softmax（在 top_k 之后的分布上）

0.5560.2780.1670

4. “之前的累计质量” cum − p

00.5560.8331.0

规则：之前的质量已经 > top_p（0.6）的才删；第 0 个强制保留。

5. > 0.6 → 删？
token1、token2 各自删不删？最终概率？

token1 保留（0.556 未达 0.6），token2 删；最终概率 [0.667, 0.333, 0, 0]。

保留保留删删

三级提示

L1I/O 示例与关键问题I/O 示例 + 关键问题 + 细节（轻度剧透）

例 1：repetition penalty

logits = torch.tensor([3.0, -1.0, 0.5, -0.2])
apply_repetition_penalty(logits, prev_tokens=[1, 3, 3], penalty=1.5)
# -> tensor([ 3.0, -1.5, 0.5, -0.3])
logits   # 仍是 [3.0, -1.0, 0.5, -0.2]（不能原地改）
s = Sampler()
s.sample(torch.tensor([[3.0, -1.0, 2.5, -0.2]]), ["r"],
         [SamplingParams(temperature=0.0, repetition_penalty=1.5)], [[0]])
# -> tensor([2])   token 0: 3.0/1.5 = 2.0 < 2.5，greedy 改选 token 2

token 3 在历史里出现两次，只惩罚一次。

例 2：top_k / top_p 保留集合（T=1，概率 [0.5, 0.25, 0.15, 0.1]）：

参数 降序累计概率 保留 过滤后重新归一化
top_p=0.4 0.5 | 0.75 | 0.9 | 1.0 {0} [1, 0, 0, 0]
top_p=0.7 同上 {0, 1} [0.667, 0.333, 0, 0]
top_p=0.8 同上 {0, 1, 2} -
top_p=1e-9 同上 {0}（至少 1 个）
top_k=2 - {0, 1} [0.667, 0.333, 0, 0]
top_k=3, top_p=0.6 top_k 后：0.556 | 0.833 | 1.0 {0, 1} -

温度：T=0.5 → [0.725, 0.181, 0.065, 0.029]，T=2 → [0.370, 0.262, 0.203, 0.166]。

例 3：批不变性（row = [1.0, 0.5, 0.0, -0.5, 0.2, 0.9]，请求 q 用 seed=7）：单独跑 6 步 → [0, 5, 4, 0, 4, 5]；混在别的请求里（位置/同伴每步都变）→ 完全一致；错误做法（全 batch 共用 seed=7 的 generator）→ [5, 4, 1, 0, 0, 0]；free("q") 后重新以 seed=7 跑 3 步 → [0, 5, 4]。具体数字取决于抽样方式；“单独 == 混合”与“free 后重放一致”才是要满足的性质。

关键问题

  1. prev_tokens=[4, 4, 4]、penalty=2，token 4 的正 logit 2.0 惩罚后是 1.0 还是 0.25？
  2. 为什么负 logit 要乘而不是除？
  3. temperature=0 时先做 logits / temperature 会怎样？repetition penalty 对 greedy 行还生效吗？
  4. top_p 的“smallest prefix whose mass reaches top_p”：越过阈值的 token 留不留？用“之前的累计质量”怎么表达？
  5. top_k=-1、top_k ≥ V 怎么处理？并列第 k 大时你的实现保留几个？
  6. 带 seed 的请求的 generator 什么时候创建、存哪、跨 step 复用吗？不带 seed 的用什么？
  7. 为什么不能对整个 [B, V] 调一次 torch.multinomial(probs, 1, generator=g)？
L2思路与常见坑思路拆分与常见坑（中度剧透）
  • 逐行处理：for i in range(B)，每行只用自己的 logits、params、prev_tokens、generator。先写对，向量化是 follow-up。
  • apply_repetition_penalty：去重得下标张量 → 取值 → torch.where(v > 0, v / p, v * p) → 写回 clone。penalty==1 或 prev 为空直接返回。
  • top_k：kth = torch.topk(row, k).values[-1]，row < kth 设 -inf；仅当 0 < k < V。
  • top_p：降序排序 → softmax → cumsum；cum − probs 已 > top_p 的删；强制保留第 0 位；用排序下标映射回原位置置 -inf。
  • generator 管理：dict[req_id, torch.Generator]；seed 为 None 用一个共享默认 generator。
  • 抽一个 token：torch.multinomial(softmax(row), 1, generator=g)，每行一次调用。
L3详解：答案详解：答案、逐步追踪、骨架提示（重度剧透）
关键问题答案
  1. 1.0：distinct，只除一次。
  2. −2.0/2 = −1.0 更大 → 更可能被选，与“惩罚”相反；乘以 2 得 −4.0 才是降低。
  3. 除以 0 得 inf/NaN。阈值判断必须在除法之前；penalty 在 greedy 之前，所以对 greedy 行仍生效（例 1 第二段）。
  4. 留。token 之前（不含自己）的累计质量已达到 top_p 的才删；第 0 个的“之前质量”是 0，天然保留。
  5. 都视为不启用。用“≥ 第 k 大值”保留时并列会多保留，可以接受。
  6. 第一次 sample 到该 req_id 时创建，存在 Sampler 实例的字典里，跨 step 复用，free 时删除。无 seed 用共享 generator。
  7. 一次调用会按行顺序依次消费同一个 generator，q 拿到的随机数取决于前面有几行；而且一个 g 无法同时代表多个不同 seed 的请求。

自检：① 例 2 每一行用你的过滤函数算出保留集合；② 同一 seed 的请求单独跑 N 步、再随机打乱 batch 跑 N 步，结果列表相等；③ 构造 0.8/0.2 的两个 token，抽几千次看频率。

面试官追问 & 对照源码

对照源码：deploy/qwen3_omni_moe.yaml:49-62 · worker/gpu_model_runner.py:1998-2010