浮点误差如何导致手写采样逻辑偏离 torch.multinomial

在复现大语言模型（LLM）token 采样行为时，我最近掉进了一个非常隐蔽的陷阱。许多开发者在处理概率采样时，习惯于将数学定义直接转为代码，认为只要逻辑正确，使用框架内置函数还是手写逆变换采样，最终结果应该完全一致。然而，在 float32 精度下，这两种方法却出现了偏差。

常规流程是先用 softmax 将模型输出的 logits 转为概率分布 probs，然后尝试两种不同的采样方式。

一种是标准的 PyTorch 官方做法，直接调用内置采样函数：

next_token = torch.multinomial(probs, num_samples=1)

另一种是我尝试手写的 CDF 累加法。其核心是先计算概率的前缀和，生成累积分布区间，然后随机生成一个 0 到 1 之间的均匀数 u，最后用二分查找确定 u 落在哪个区间内：

手写 CDF 累加法是否会导致区间宽度为 0 的问题？

cdf = torch.cumsum(probs, dim=-1)
u = torch.rand(1, device=probs.device)
next_token = torch.searchsorted(cdf, u)

从纯数学定义看，这两种方法是完全等价的，都执行了 Categorical Distribution 的采样。但在 float32 下，手写的 CDF 采样器无法完全复现 torch.multinomial 的行为。

最让我困惑的细节是：当词表规模极大时，float32 的精度非常危险。实验中观察到在执行 torch.cumsum 后，会出现 cdf[i] == cdf[i-1] 的情况。理论上，只要某个 token 的概率 probs[i] > 0，那么累加值 cdf[i] 必然应该比 cdf[i-1] 大。但在有限精度的浮点数运算中，累积值会随着索引增加而迅速增大。如果某个 token 的概率值极低，而累积值已达到较大量级，那么加法运算时极小的概率值会在巨大累积值上被直接“吞掉”，导致结果无变化。

这意味着手写采样逻辑中，该 token 对应的区间宽度实际上变成了 0。随机数 u 的选取过程中，没有任何值能落在该 token 的区间内，该 token 在逻辑上被“抹除”，永远不可能被抽中。

而 torch.multinomial 在实际表现中显然避免了这个问题，这说明它的底层实现做了深度优化，并没有简单地走 cumsum -> rand -> searchsorted 的线性累加流程。尽管 PyTorch 的官方文档没有详细披露具体算法，但我们可以推断，内置函数在 CUDA 和 CPU 上的实现可能完全不同。

为了保证数值稳定性，torch.multinomial 内部可能采用了比简单累加更复杂的算法，例如 Alias Method 将采样复杂度降低到 O(1) 且避免累加误差，或采用了经过优化的累加策略（如 Kahan Summation）来减少精度丢失。

官方实现的优化算法是否能完全复现数学定义的采样行为？

这次经历的最大启示是，在处理大规模词表的概率运算时，不能简单地将数学公式等同于代码实现。float32 的精度在面对数万个元素的累加时，其数值不稳定性会通过 torch.cumsum 等操作被放大。如果你复现模型行为时发现采样结果与官方实现不一致，请务必检查累积分布函数是否在某些位置出现了数值停滞，因为那些被“抹除”的微小概率，正是导致结果偏差的根源。