复盘 Rejection Sampling 在自研 AI 加速器上的数据流与软件流水优化，以及快速数学指令引入的采样偏差诊断。
国产芯片算子优化：Rejection Sampling
面向读者： 了解概率采样与基本 kernel 优化、希望将算法映射到加速器的开发者。本文讨论如何选择片上数据流、隐藏指令依赖，以及定位快速数学指令引入的采样偏差。
核心问题与结论： 数据搬运、依赖等待与数值近似共同影响这个 kernel。本文记录了数据流选择、软件流水和 log 修复过程；固定全接受测试点得到约 11.9× 的 kernel 加速，扩大样本后仍观察到快速采样路径的统计偏差。
范围与依据： 本文为公司内部源码相关的优化复盘，保留原稿中的测试结果和分析，不提供源码与配套公开复现材料，也不列内部硬件型号、实现与测试版本标识。性能表限定为 BF16、词表大小 154,880、Sequential 模式，采用各场景自己的基线；数值实验的样本规模分别在对应章节说明。
在自研 AI 加速器上实现 Rejection Sampling 时，框架组已经提供了采用 Gumbel-Max 的 PyTorch 参考实现。算法具备逐 token 并行的结构，但直接翻译成 kernel 后，数据搬运、寄存器依赖和特殊函数延迟仍会让硬件等待。换用快速数学指令，又带来了参考实现中没有的数值问题。
本文分上下两篇，先交代两条线共用的算法与语义：
- 上篇：性能优化。 比较 BF16 数据流，解释如何交叠独立工作，最后检查局部调度与完整 kernel 周期。
- 下篇：数值诊断与修复。 从快速 log 的异常定位出发，解释 Div8 与 Div14 的选择，并检查不同样本规模下的分布质量。
实现还区分 Exact 与 Fast 两种路径：前者需要对齐框架的随机数序列要求；后者不要求同 seed 逐位复现，可以使用硬件 RNG，但仍需验证输出分布的质量。本文的性能与数值优化主要围绕 Fast 路径展开。
共同前提：算法与语义
Speculative Decoding 先由 draft 模型提出候选 token，再由 target 模型验证。记 target、draft 的概率分布分别为
对于存在正权重的情况，Gumbel-Max 可以把这次采样写成：
零权重对应的 score 为负无穷；归一化只会给所有 score 增加同一个常数，因此可以省略。原理与证明可参考 A Review of the Gumbel-max Trick，这里关注它对实现的意义：每个 token 独立打分，最后做一次 argmax，无需构造 CDF。
sum-exp 在采样流程中的位置
输入是 target 与 draft 的 logits，还需要先得到各自的 log probability。对于一行 logits
这里的 max 与 sum-exp 就是上篇优化的两遍扫描。当前 kernel 分别计算 target 与 draft 的
拒绝后，令 target、draft logits 分别为
对于
再加 Gumbel 噪声并做 argmax。Fast kernel 将相应指数与对数换到底数 2 实现；无正 residual weight 的项不应参与采样。
因此，省掉的是 residual 权重的再次归一化与 CDF 构造，target/draft 从 logits 得到概率所需的 logsumexp 仍然存在。全接受后的 bonus token 直接走 target logits 采样分支；Block、无 draft logits 与 greedy 路径另有处理，这里只串起本文性能表所采用的普通 Sequential 主线。
这段衔接依据现有 kernel 的 LSE、接受判断及 residual 构造代码，并对照本地 PyTorch 测试中的 reference 整理；不将该 reference 认定为当时交付的原始 PyTorch 文件。
上篇：性能优化
1. 瓶颈与数据流：为什么读两遍 BF16？
对于sum-exp，当时提出了3种方案：
| 方案 | 做法 | 主要代价 |
|---|---|---|
| FP32 中间结果落地 | 第一遍 max unpack 后保存 FP32，第二遍直接读取 | 增加 FP32 片上写入与读取 |
| 两遍 packed BF16 | 两遍分别读取、unpack，不保存整行 FP32 | 重复读取 packed 数据与 unpack |
| Online sum-exp | 一遍维护动态最大值和指数和 | 增加 rescale 计算及循环依赖 |
| 这里的读取发生在片上 VMEM 与寄存器之间。输入已经搬到片上，两遍扫描并不意味着从 HBM 重复搬运两遍。 |
FP32 落地能省掉第二次 unpack，却会增加片上 Load/Store 的负担。Online sum-exp 能省掉第二遍读取，但每次最大值变化都需要调整已有的指数和，增加计算和循环携带的依赖。究竟哪一种更好，需要看硬件上哪类资源先耗尽。最后通过结合硬件后端资源，我选择了第二种。
后续复盘发现，原先对 FP32 落地成本的估计包含了“恢复通用连续布局”的开销。若 sum-exp 的中间缓冲区改用适合顺序扫描的布局，这部分成本可以减少；下面给出从转换封装确认的布局与 store 数量。此前“约省 1 cycle / packed BF16”的估计尚未由新的调度与实测确认，不作为本文性能结果。
Layout 补充：解包与恢复连续顺序是两件事
转换封装先从同一个 packed word 解出低、高两个 BF16，再分别转成 FP32。结合其写回位置与 kernel 的索引构造，可以还原每组 256 个逻辑元素的布局。记
| packed word | 低 16 bit | 高 16 bit |
|---|---|---|
| 第 |
||
| 第一组第 0 个 word | ||
| 第一组第 127 个 word | ||
| 第二组第 0 个 word |
一次完整的宽向量加载包含 8 组、每组 128 个 packed word，共 2048 个 BF16 值。解包后，低半向量包含 [0…127, 256…383, …]，高半向量包含 [128…255, 384…511, …]。
通用转换要恢复 [0…2047] 的 FP32 顺序，因此两路结果都按 stride=2 写回，高半向量的起点比低半向量后移 128 个 FP32 元素。该访问在 8 个 bank 上重复命中同一组 bank；store 封装将每路拆成高、低各 4 个子核的掩码写入。因此，对完整块而言，两个逻辑写回实际展开为 4 次 store。这里说的是封装展开的操作数量，最终发射周期仍取决于调度。
如果中间结果只用于后续 max/sum-exp 扫描，可以考虑先连续保存整个低半向量，再连续保存整个高半向量。这样每路可用 stride=1 写回，合计 2 次 store，但缓冲区不再是原来的逻辑顺序。消费者必须使用同一布局；涉及索引的操作要恢复对应关系，浮点归约顺序变化也需要检查误差。
因此，方案比较应分别计算“转成 FP32”和“恢复通用布局”的成本。上述代码支持少两次 store 的设计机会，不能单凭这一点宣布 FP32 路径最终更快。
2. 软件流水：如何把独立工作交叠起来？
数据流确定之后，还需要解决依赖等待。以 sum-exp 为例，每组数据要经过加载、unpack、减最大值、乘换底常数、发起指数计算、取回结果和累加。若一组做完再开始下一组，即使每条指令本身很快，依赖之间也会留下大量空闲时间。
这里要区分两种等待：普通算术指令的结果写回带来较短的 RAW 依赖；异步特殊函数单元则需要更长时间才能产生可取回的结果。发射吞吐和结果可用时间是两个不同约束。 单元能接受下一次请求，不代表上一次结果已经能被消费。
我用四组独立寄存器状态拉开短依赖的复用距离，同时让不同迭代处在不同计算阶段：
| 逻辑阶段 | 当前推进的工作 |
|---|---|
| Unpack | 已 unpack 的数据减去行最大值 |
| Multipy | 乘以换底常数，为指数计算做准备 |
| Pow | 向特殊函数单元发起指数计算 |
| Pop | 弹出上一轮单元的指数运算，并且完成累加 |
之所以这么切，一方面ResMII只有10cycles，单个body内无法完全隐藏Pow的13cycle延迟。而对于Unpack，因为Pow操作有1 cycle的hold-issue，如果不拆成2条独立链，没有足够的独立指令来避免连续Pow产生的Pipeline-stall。
每个阶段有四组 worker。一轮同时推进不同迭代的阶段，之后轮换寄存器的逻辑角色；更早发起的异步结果通过 FIFO 取回并累加。
角色轮换可以通过硬件旋转，或通过代码展开实现即，下一轮改变同一批物理寄存器的用途。四组 worker 则是在当前依赖距离和寄存器预算下的选择（具体地对于该硬件的FP操作延迟，即从发射到写回，小于4 cycles）；继续增加并行状态会延长值的生命周期，可能引入寄存器压力，不能无限加深。
最终，sum-exp 形成了 10-cycle 的静态循环调度，同时交叠指数请求、算术运算、结果累加、unpack 和加载补充，实测 trace 中该稳态段没有 stall。Residual-Gumbel 的依赖更复杂，28-cycle 的调度体仍约有 1 cycle stall/pass。分别检查这些局部行为，才能知道总周期下降来自哪里，以及哪里仍有空间。
调度前还可以先缩短计算图。例如 argmax 的 value 更新可以由 greater → select 改成独立的 max，比较结果只用于更新 index。这样能缩短 value 的依赖链，但必须同时检查相同最大值和 NaN 的处理语义。
3. 性能结果与验证范围
以下是 BF16、词表大小 154,880、Sequential 模式下的完整 kernel cycles。其中 S 为每个请求的 speculative 长度（draft 步数），R 为 request 数量；两条路径采用各自基线，不能把加速比直接互相比作优化优劣。
| 场景 | 测试参数 | 基线 cycles | 优化后 cycles | 加速比 |
|---|---|---|---|---|
| 全接受 | S=4，R=8 | 1,530,302 | 128,453 | 11.91× |
| 首步拒绝 | S=4，R=2 | 55,896 | 38,032 | 1.47× |
全接受对照的是初始语义正确实现，首步拒绝对照的是已经包含部分优化的专项基线。这些数字描述固定测试点的 kernel 收益，不代表整套推理服务的加速比。
性能侧用局部调度、trace 和总周期互相验证。完整 kernel 的周期用于衡量最终收益，局部调度用于检查等待是否被隐藏。
下篇：数值诊断与修复
性能收益需要与采样语义一起检查。下面分别观察同输入误差、输出分布以及模拟器与实卡的一致性，并通过扩大样本复查修复结论。
1. 数值修复：从异常定位到 Div8
输出没有 NaN 或 Inf，并不意味着采样质量正常。 Fast 路径使用快速硬件 log 生成 Gumbel 噪声，这里遇到的问题是：log 在输入接近 1 的一个区间内，直接返回 0。
实现采用 base-2 形式，概率项也统一换底，保持 argmax 的数学等价关系。噪声的计算形式为：
先保护第一层，再定位第二层
第一层 log 的输入可能接近 1。我用
缩放后的输入上界低于已观察到的坏区间。这与直接 clamp 不同：clamp 会把一段概率质量压到同一个点，而缩放在精确数学中可以通过常数补偿恢复。
但第一层修复之后，1,048,576 个样本仍出现 8627 个精确零。为了定位原因，我把 Gumbel 计算从完整 kernel 中隔离出来，对同一随机输入同时观察中间值与高精度 reference。
当时有两个假设：第一次 log 的普通量化把结果集中到某个值，或第二次 log 自己命中了零平台。根据第一次 log 的量化区间估算，百万样本预计只会命中约 0.5 次，无法解释实际的 8627 次。
逐值 trace 随后显示：某个样本经过第一层 log 和补偿后，第二层输入约为 1.0000002384，恰好落入零平台。按照该平台区间估计的命中数约为 8609 次，与实测量级吻合。概率估算和逐值追踪共同指向了第二层 log。
Div8：把主要概率质量移出坏区间
令
精确数学中，这仍等于
一句话地说，由于第一次log后的分布是无界分布，所以不能通过缩放来避开0平台，但是可以通过调整，不伤害精度的情况下，使得0平台质量分布较小。
| 验证项 | Div8 结果 |
|---|---|
| 样本数 | 1,048,576 |
| 精确零 | 8627 → 0 |
| KS | |
| 64-bin χ² | |
| 同输入绝对误差 P99.9 | 约 |
在这个规模与检验设置下，修复后的结果通过了统计检验。模拟器与实卡的同位置对照中，只有 57 个 FP32 输出不逐位一致，最大绝对差约为
2. 扩大验证：从 Div8 到 Div14
完成 Div8 修复后，我想继续确认更大样本下的随机质量。将规模扩大到十亿样本，64-bin χ² 达到约 851.50，检出了此前百万样本检验未能识别的系统偏差。
进一步检查快速 log 的实现，可以看到分段拟合、中间结果截断和最终舍入都会影响输出。避开显式零平台之后，这些近似误差仍然存在。于是我尝试在数学等价的变换中搜索更好的工作区间：
二次幂缩放主要改变指数；非二次幂缩放还会改变归一化尾数，可能切换拟合分段。因此，不同
初筛采用三个独立 seed、每个 seed 一亿样本；随后对候选方案进行十个 seed、各一亿样本的评估。结果表明，Div10 虽然改善了 KS，χ² 却更差；Div14 在实测候选中取得了更好的综合结果。
| 最终评估指标 | Div8 | Div14 | 变化 |
|---|---|---|---|
| 各 seed 的 KS distance 均值 | −65.03% | ||
| 十亿样本聚合 χ² | 851.50 | 752.78 | −11.59% |
KS 是各 seed 分别计算后取均值；χ² 使用十亿样本的聚合计数。Div14 改善了这些指标，但聚合 χ² 的 p 值仍约为
可迁移的方法与未解问题
分析数据流时，先确定瓶颈位于 HBM、片上存储还是寄存器，再比较存储与重算的代价。安排软件流水时，分别检查发射吞吐、结果延迟与寄存器预算。检查数值问题时，将子计算隔离，用相同输入对照中间值，并结合分布检验和更大样本复查结论。
Div14 改善了已列出的指标，但没有消除大样本偏差；极端尾部与实际 token 采样的关系仍待验证。
参考与复现材料
- 理论参考：A Review of the Gumbel-max Trick，对应“共同前提”的算法背景。
- 公开源码与配套复现材料：—。不公开： 本文涉及公司内部源码，不提供公开源码仓库、复现包或内部日志附件；正文保留实验条件、汇总结果与结论边界。