Speculative Decoding Verification

原始题目：LeetGPU - Speculative Decoding Verification

题目描述

实现推测解码（Speculative Decoding）的 token 验证步骤。草稿模型提议 $T$ 个 token；目标模型在一次前向传播中评估它们，并逐个接受或拒绝。给定 $B$ 个序列，生成验证后的输出 token。

对每个序列 $b$，按从左到右顺序处理位置 $i = 0, \dots, T-1$：

1. 计算接受概率：$\alpha_i = \min\!\left(1, \frac{q_i(t_i)}{p_i(t_i)}\right)$
2. 若 $u_i < \alpha_i$：接受 $t_i$，继续位置 $i+1$
3. 若 $u_i \ge \alpha_i$：拒绝并停止。从调整分布重采样：

$$adj(v) = \frac{\max(0, q_i(v) - p_i(v))}{\sum_{v'} \max(0, q_i(v') - p_i(v'))}$$

使用逆 CDF 方法以 $r = uniform\_samples[b, T]$ 采样。若 $adj$ 全为零，回退到均匀分布。

若所有 $T$ 个 token 均被接受：从 $q_{T-1}$ 中采样一个 bonus token。

输出写入 output_tokens（形状 $[B, T+1]$），填充位置 0 到 accepted count，剩余位置为零。

实现要求

• 实现 solve(draft_tokens, draft_probs, target_probs, uniform_samples, output_tokens, B, T, V)。
• 不允许使用外部库。
• 内存布局：行优先，draft_probs[b, i, v] 偏移为 b*T*V + i*V + v。

示例

Input:  B=1, T=3, V=4, draft_tokens=[1,2,0]
        p0=[0.1,0.6,0.2,0.1], q0=[0.1,0.5,0.2,0.2]
        p1=[0.1,0.2,0.5,0.2], q1=[0.3,0.2,0.2,0.3]
        uniform_samples=[0.5, 0.7, 0.3, 0.9]
Pos 0: α₀=min(1,0.5/0.6)≈0.833, u₀=0.5<0.833 → 接受 token 1
Pos 1: α₁=min(1,0.2/0.5)=0.4, u₁=0.7≥0.4 → 拒绝, 从 adj=[0.2,0,0,0.1]→norm=[2/3,0,0,1/3] 重采样得 token 3, u=0.9
Output: [1, 3, 0, 0]

约束条件

• $1 \le B \le 256$，$1 \le T \le 16$，$2 \le V \le 131{,}072$。
• 概率分布有效（非负、和为 1），$draft\_probs[b,i,t_i] > 0$。
• 性能测试在 $B = 64,\ T = 8,\ V = 32{,}768$ 下进行。

解题思路

推测解码验证的核心是逐位置的条件接受/拒绝逻辑。$B$ 个序列可以完全并行处理。每个位置的接受概率计算和 CDF 采样是独立的小计算，但由于拒绝会提前终止（后续位置跳过），需要处理变长输出。可以将所有 $B$ 个序列映射到独立的 warp/block 处理。欢迎在 GitHub Discussions 分享你的解法。