执行摘要
- 一句话：为推测解码拒绝采样器添加贪婪采样支持，优化温度为零时的性能。
- 推荐动作：建议工程团队精读此PR，特别关注
_gather_draft_logits_and_target_argmax_kernel和_probabilistic_rejection_kernel的设计，以及review中讨论的正确性问题。设计决策如本地argmax计算和贪婪路径隔离值得学习。
功能与动机
PR body指出此PR是跟进#35461，专门为贪婪采样（temperature=0）提供支持，以高效处理贪婪请求而不影响批次性能。
实现拆解
- 新增目标argmax计算内核：在
vllm/v1/worker/gpu/spec_decode/rejection_sampler.py中新增_gather_draft_logits_and_target_argmax_kernel函数，根据温度是否为0计算目标logits的局部argmax和max值，为贪婪采样准备数据。
- 修改概率拒绝采样内核：将原
_probabilistic_rejection_sample_kernel重命名为_probabilistic_rejection_kernel，并集成贪婪采样逻辑；当温度=0时，只接受与目标argmax匹配的草稿token。
- 调整数据接口：在
vllm/v1/worker/gpu/model_runner.py的sample方法中，简化draft_logits的传递，移除索引映射和空值检查，直接使用self.req_states.draft_logits。
- 变量重命名与位置移动：将
residual_pos重命名为rejected_pos，并将计算从_compute_residual_logits_kernel移动到_probabilistic_rejection_kernel，提高逻辑一致性。
关键文件：
vllm/v1/worker/gpu/spec_decode/rejection_sampler.py（模块 推测解码；类别 source；类型 core-logic；符号 _gather_draft_logits_and_target_argmax_kernel, _probabilistic_rejection_kernel, probabilistic_rejection_sample）: 核心变更文件，新增和修改Triton内核以实现贪婪采样支持。vllm/v1/worker/gpu/model_runner.py（模块 模型运行器；类别 source；类型 data-contract）: 调整数据接口，简化draft_logits传递以支持新采样逻辑。
关键符号：_gather_draft_logits_and_target_argmax_kernel, _probabilistic_rejection_kernel, probabilistic_rejection_sample
评论区精华
TheEpicDolphin解释将residual_pos计算移动到_probabilistic_rejection_kernel并重命名为rejected_pos的原因，以提升代码清晰度。gemini-code-assist[bot]指出新增的_flatten_sampled_kernel中循环可能读取未初始化值，存在正确性风险。WoosukKwon建议未来可以融合更多内核以减少张量物化，但认可当前实现可作为后续优化基础。
- 移动residual_pos计算和重命名 (design): 已实现变更。
- 潜在未初始化读取问题 (correctness): 问题被指出，但PR已合并，可能需后续修复。
- 内核融合建议 (performance): 建议被记录，未来可能跟进。
风险与影响
- 风险：主要风险在于gemini-code-assist[bot]指出的潜在未初始化读取问题，可能导致输出错误；贪婪采样路径增加了内核复杂度，可能引入性能回归；修改涉及核心推测解码逻辑，需确保与现有严格拒绝采样和概率采样模式的兼容性。
- 影响：对用户：使贪婪采样在推测解码中更高效，提升温度为零场景的吞吐量。对系统：优化了拒绝采样器的性能，减少对非贪婪请求的影响。对团队：引入新内核需加强测试覆盖，后续可能需进行内核融合以进一步提升性能。
- 风险标记：潜在未初始化读取, 内核分离性能影响, 兼容性风险
关联脉络
- PR #35461 推测解码拒绝采样器基础功能（从PR body提及推断）: 此PR是#35461的跟进，专门添加贪婪采样支持，属于同一功能线。
- PR #39773 [Model Runner V2] Disable piecewise cudagraph mode fallback for eagle draft decodes: 同属模型运行器V2和推测解码功能线，涉及相关组件。
- PR #38372 [Hybrid] Simplify accepted token counting in spec decode for hybrid models: 涉及推测解码的令牌计数简化，功能相关。
参与讨论