执行摘要

• 一句话：回滚 #34854 变更，使用 FP64 提高 Gumbel 噪声数值稳定性，牺牲大 batch 随机采样性能。
• 推荐动作：建议精读，特别是 tl_rand64 的实现和性能数据，以理解 Triton 内核中精度与性能的权衡，以及如何通过代码简化提升可读性。

功能与动机

根据 PR 描述，此变更旨在提高数值稳定性（'for numerical stability'），回滚了之前的 #34854 PR。性能测试显示，对于贪婪采样无影响，对于随机采样在大 batch 下有性能损失。

实现拆解

主要修改两个文件：在 vllm/v1/worker/gpu/sample/gumbel.py 中新增 tl_rand64 函数，使用 Triton 的 tl.randint4x 生成 64 位随机数并转换为 float64；修改 _gumbel_sample_kernel 将 logits 转换为 float64 并使用 tl_rand64 生成 Gumbel 噪声。在 vllm/v1/worker/gpu/spec_decode/rejection_sampler.py 中，将 target_prob 和 draft_prob 转换为 float64，并使用 tl_rand64 替代原有的 tl.rand 调用。

关键文件：

• vllm/v1/worker/gpu/sample/gumbel.py（模块 sampling）: 核心 Gumbel 采样实现，引入 FP64 随机数生成函数 tl_rand64 并修改内核使用 float64
• vllm/v1/worker/gpu/spec_decode/rejection_sampler.py（模块 speculative decoding）: 推测解码中的概率拒绝采样，改用 FP64 精度以保持一致性

关键符号：tl_rand64, _gumbel_sample_kernel, gumbel_sample, _probabilistic_rejection_kernel

评论区精华

review 中 gemini-code-assist[bot] 指出 rejection_sampler.py 中随机数可能为零导致偏差，建议 clamp 以确保正数；WoosukKwon 采纳并修复。另一评论建议简化随机数生成逻辑，直接传递标量 offset 到 tl_rand64，也被接受。

• 避免随机数零值导致的偏差 (correctness): WoosukKwon 修复了代码，使用 tl_rand64 并设置 includes_zero=False 来避免零值。
• 简化随机数生成逻辑 (design): WoosukKwon 接受了建议，修改了代码。

风险与影响

• 风险：技术风险包括：性能风险，随机采样在大 batch（如 1024）下最多有 1.81 倍 slowdown；兼容性风险，改变数据类型可能影响其他依赖模块；正确性风险，已通过修复 bias 问题缓解。核心路径变更需关注回归测试。
• 影响：对用户影响：随机采样性能下降，尤其大 batch 场景，可能影响吞吐量；对系统影响：提高数值稳定性，减少因浮点误差导致的不确定性；对团队影响：需要权衡精度与性能，可能引发后续优化。
• 风险标记：性能下降大 batch 随机采样, 核心采样路径变更

关联脉络

• PR #34854 未知，从上下文推断为使用 FP32 的优化: PR #37798 回滚了 #34854 的变更，以恢复使用 FP64 提高数值稳定性。