执行摘要

• 一句话：修复 CPU 推测解码因 kernel 签名缺失崩溃
• 推荐动作：建议精读，这是一个展示如何通过最小签名变更加入兼容层以修复跨模块调用错误的优秀案例。对于维护 CPU 或扩展 spec decode 的开发者尤其值得关注。

功能与动机

CPU speculative decoding 在启动时立即崩溃，报错 TypeError: _rejection_greedy_sample_kernel_impl() got an unexpected keyword argument 'SYNTHETIC_MODE'。原因是 PR #40662 统一引入 synthetic mode 支持后，rejection sampler 总是传递这些参数，但 CPU Triton fallback kernel 没有对应的形参。

实现拆解

本 PR 的修改集中在 vllm/utils/cpu_triton_utils.py，包含两个步骤：

1. 修改 _rejection_greedy_sample_kernel_impl：在签名末尾添加三个带默认值的形参 uniform_probs=None、synthetic_conditional_rates=None、SYNTHETIC_MODE=False，并加入 assert not SYNTHETIC_MODE 断言和注释说明。该函数封装了 C++ CPU kernel 的调用，实际 C++ kernel 仍不处理 synthetic mode，因此参数仅用于接口兼容。

2. 修改 _rejection_random_sample_kernel_impl：类似添加 synthetic_conditional_rates=None 和 SYNTHETIC_MODE=False（注意 uniform_probs 此前已存在），同样添加断言和注释。

两个修改均不改变既有逻辑，不引入新依赖，保持向后兼容。

关键文件：

• vllm/utils/cpu_triton_utils.py（模块 CPU工具；类别 source；类型 core-logic；符号 _rejection_greedy_sample_kernel_impl, _rejection_random_sample_kernel_impl）: 单文件变更，包含两个 CPU fallback kernel 包装函数，是修复的唯一修改目标。

关键符号：_rejection_greedy_sample_kernel_impl, _rejection_random_sample_kernel_impl

关键源码片段

vllm/utils/cpu_triton_utils.py

单文件变更，包含两个 CPU fallback kernel 包装函数，是修复的唯一修改目标。

def _rejection_greedy_sample_kernel_impl(
    output_token_ids,
    cu_num_draft_tokens,
    draft_token_ids,
    target_argmax,
    bonus_token_ids,
    is_greedy,
    max_spec_len,
    uniform_probs=None, # added: 用于合成接受率计算，CPU 端尚未实现
    synthetic_conditional_rates=None, # added: 同上
    SYNTHETIC_MODE=False, # added: 若为 True 则断言失败，防止误用
):
    # C++ kernel 要求所有整数张量为 int64 。
    # 注意：上述三个参数由 rejection sampler 传入以支持 synthetic mode，
    # 但 C++ CPU kernel 尚未实现该功能。此处仅接收参数以保持调用约定兼容。
    assert not SYNTHETIC_MODE, 'Synthetic acceptance not supported with CPU sampling'
    orig_dtype = output_token_ids.dtype
    output_token_ids_i64 = _ensure_int64(output_token_ids)
    torch.ops._C.rejection_greedy_sample_kernel_impl(
        output_token_ids_i64,
        _ensure_int64(cu_num_draft_tokens),
        _ensure_int64(draft_token_ids),
        _ensure_int64(target_argmax),
        _ensure_int64(bonus_token_ids),
        is_greedy,
        max_spec_len,
    )
    if orig_dtype != torch.int64:
        output_token_ids.copy_(output_token_ids_i64.to(orig_dtype))

评论区精华

审核者 benchislett 在两个 kernel 函数中均建议添加 assert not SYNTHETIC_MODE 断言，以防止在 CPU 上误用 synthetic mode（C++ kernel 尚未实现）。开发者接受建议并纳入最终补丁。ganeshr10 评论称赞最小签名变更是正确的方式，保持 CPU/GPU 调用约定一致。

• 添加 SYNTHETIC_MODE 断言防止误用 (design): 开发者采纳，在最终补丁中加入断言。

风险与影响

• 风险：风险极低。所有新参数都有默认值，不会破坏现有调用方。新增的断言在 synthetic mode 被意外启用时会主动报错，提供清晰的错误信息，避免静默错误。未来若在 CPU 上实现 synthetic mode，只需移除断言并更新 kernel 即可。没有回归风险。
• 影响：对用户：CPU 推测解码用户不再遇到启动崩溃，可直接使用 spec decode 获得约 2x 加速。对 GPU 用户无影响。对系统：仅修改 CPU fallback 路径，不影响 GPU Triton kernel 或其他模块。对团队：合并容易，无冲突风险。
• 风险标记：兼容性变更

关联脉络

• PR #32662 Enable CPU speculative decoding: 本 PR 修复了 #32662 合并后引入的回归，该 PR 首次启用 CPU spec decode。
• PR #40662 Add unified synthetic acceptance rate support: 该 PR 引入 synthetic mode 参数导致 rejection sampler 传递新参数，从而引发 CPU kernel 签名不匹配。