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 处理。

代码实现

CUDA

#include <cuda_runtime.h>
__global__ void spec_verify(const int* dt, const float* dp, const float* tp,
    const float* us, int* out, int B, int T, int V) {
    int b=blockIdx.x; if(b>=B)return;
    for(int i=0;i<T;i++){
        int ti=dt[b*T+i]; float pi=dp[b*T*V+i*V+ti], qi=tp[b*T*V+i*V+ti];
        float alpha=(pi>0)?fminf(1.0f,qi/pi):0.0f;
        if(us[b*(T+1)+i] < alpha){ out[b*(T+1)+i]=ti; }
        else {
            // Rejection sampling from adjusted distribution
            float sum=0; for(int v=0;v<V;v++) sum+=fmaxf(0.0f,tp[b*T*V+i*V+v]-dp[b*T*V+i*V+v]);
            float r=us[b*(T+1)+T],cum=0;
            for(int v=0;v<V;v++){cum+=fmaxf(0.0f,tp[b*T*V+i*V+v]-dp[b*T*V+i*V+v])/sum;if(r<cum){out[b*(T+1)+i]=v;break;}}
            for(int j=i+1;j<T;j++)out[b*(T+1)+j]=0;
            return;
        }
    }
    out[b*(T+1)+T]=tp[b*T*V+(T-1)*V + out[b*(T+1)+(T-1)]];
}
extern "C" void solve(const int* dt,const float* dp,const float* tp,const float* us,int* out,int B,int T,int V) {
    spec_verify<<<B,1>>>(dt,dp,tp,us,out,B,T,V); cudaDeviceSynchronize();
}

Triton

import triton, triton.language as tl
@triton.jit
def spec_verify(dt_ptr,dp_ptr,tp_ptr,us_ptr,out_ptr, B,T,V, BLOCK:tl.constexpr):
    b=tl.program_id(0)*BLOCK+tl.arange(0,BLOCK); mask=b<B
    for i in range(T):
        ti=tl.load(dt_ptr+b*T+i,mask=mask)
        pi=tl.load(dp_ptr+(b*T+i)*V+ti,mask=mask)
        qi=tl.load(tp_ptr+(b*T+i)*V+ti,mask=mask)
        alpha=tl.where(pi>0,tl.minimum(1.0,qi/pi),0.0)
        accept=tl.load(us_ptr+b*(T+1)+i,mask=mask) < alpha
        tl.store(out_ptr+b*(T+1)+i,ti,mask=mask&accept)