vLLM speculative decoding이 기각 토큰을 다시 뽑는 이유

작은 모델이 쓴 초안을 큰 모델이 검사하면, 결과물도 작은 모델을 조금은 닮게 되지 않을까요? Leviathan 등이 2022년에 낸 논문 「Fast Inference from Transformers via Speculative Decoding」(arXiv 2211.17192)은 초록에서 T5-XXL을 "identical outputs"로 2~3배 빠르게 만들었다고 말합니다. 출력이 같다는 말은 확률론적 보장입니다. 그 보장을 만드는 것은 기각된 토큰을 아무렇게나 다시 뽑지 않는 규칙 한 줄입니다.

드래프트가 제안하고 타깃이 판정한다

설정은 이렇습니다. 빠른 드래프트 모델 Mq가 다음 토큰 분포 q(x)를, 느린 타깃 모델 Mp가 p(x)를 냅니다. 논문 Algorithm 1은 Mq로 γ개 토큰을 먼저 뽑고, Mp를 γ+1개 위치에 대해 병렬로 한 번에 돌립니다. 메모리 대역폭이 병목인 디코딩의 빈 연산 자원을 활용하는 구조입니다.

판정은 위치마다 독립 균등 난수 r을 뽑아 앞에서부터 진행합니다.

위치 i의 드래프트 토큰 x_i (x_i ~ q_i)
  r_i ~ U(0,1)
  r_i <= p_i(x_i) / q_i(x_i)  → 수락, 다음 위치로
  그렇지 않으면              → 기각, x_i 자리를 norm(max(0, p_i − q_i))에서 새로 뽑고 종료
γ개 전부 수락되면           → p_{γ+1}에서 한 토큰을 더 뽑는다 (bonus)

q(x) ≤ p(x)인 토큰은 항상 통과하고, 드래프트가 과하게 확신한 토큰만 p/q 확률로 살아남습니다. 한 반복은 최소 1개, 최대 γ+1개 토큰을 만듭니다.

왜 하필 max(p − q, 0)인가

기각되면 타깃 p에서 그냥 다시 뽑으면 왜 안 될까요? 그렇게 하면 분포가 틀어집니다. 논문 Appendix A.1의 증명을 따라가면 이유가 보입니다.

어떤 토큰 x'가 최종 출력될 확률은 두 경로의 합입니다. 수락 경로의 확률은 q(x')·min(1, p(x')/q(x')) = min(q(x'), p(x'))입니다. 전체 수락 확률을 β라 하면 β = Σ min(p, q)이고(논문 Theorem 3.5), 기각 경로는 (1−β)·p'(x')입니다. 출력 분포가 p가 되려면 기각 경로가 정확히 p(x') − min(q(x'), p(x'))를 채워야 합니다. 그 값을 1−β로 나눈 것이 바로 norm(max(0, p − q))입니다.

잔여 분포는 수락 경로가 채운 몫을 p에서 뺀 나머지입니다. 기각 후 p에서 다시 뽑으면 이미 채운 몫에 확률이 한 번 더 얹힙니다.

논문은 Lemma 3.3에서 1 − Σ min(p, q)가 두 분포의 거리(D_LK)임을 보입니다. 드래프트가 타깃과 닮을수록 수락률 α = E[Σ min(p, q)]가 1에 가까워집니다.

vLLM 커널은 이 규칙을 어떻게 옮겼나

vLLM main@5689c27(2026-11-01 커밋, 작성 시점 2026-11-02 기준)의 vllm/v1/sample/rejection_sampler.py는 독스트링에서 이 논문을 "strictly follows"한다고 밝히고, 출력을 accepted + recovered + bonus 토큰으로 나눕니다.

rejection_random_sample_kernel의 수락 조건은 draft_prob > 0 and target_prob / draft_prob >= uniform_prob입니다. 균등 난수는 generate_uniform_probs가 float64로 만듭니다. 주석은 float32면 정확히 0.0이 나올 확률을 무시할 수 없어서라고 설명합니다. 조건식대로라면 u = 0일 때 타깃 확률 0인 토큰도 통과할 수 있다는 점에서 나온 선택으로 읽힙니다(코드에서 끌어낸 추론입니다).

기각 시 대체 토큰은 sample_recovered_tokens_kernel이 고릅니다. 정규화를 하지 않고 max(target_prob − draft_prob, 0)에 Exp(1) 난수의 역수를 곱한 뒤 argmax를 취합니다. w_i/E_i(E_i는 독립 지수분포)의 argmax는 w에 비례하는 샘플이 되므로, 주석 말대로 어휘 전체 합으로 나누는 정규화가 필요 없습니다.

타깃 logits에는 판정 전에 apply_sampling_constraints가 temperature·top_k·top_p를 적용합니다. 논문 2.2절처럼 샘플링 설정까지 반영한 분포를 p로 보는 셈입니다. temperature와 top_p의 적용 순서는 temperature는 top_p 앞일까 뒤일까에서 다뤘습니다.

기본값 greedy 드래프트에서도 분포가 지켜질까

vllm/config/speculative.py의 SpeculativeConfig.draft_sample_method 기본값은 "greedy"입니다. 독스트링에 따르면 드래프트는 항상 argmax를 고르고, 기각 판정에서 드래프트 확률은 one-hot으로 다뤄집니다("probabilistic"은 드래프트 분포에서 샘플링하는 대신 GPU 메모리를 더 씁니다). 이때 커널에는 draft_probs가 None으로 들어가고, 코드는 draft_prob = 1로 둡니다.

드래프트가 argmax만 낸다면 타깃이 temperature 1로 샘플링할 때 출력이 드래프트 쪽으로 쏠리지 않을까요? one-hot q를 넣으면 수락 확률은 p(x)이고, 잔여 분포 max(p − δ_x, 0)은 x 자리만 0인 p입니다. 커널도 이 경우 드래프트 토큰 위치를 마스킹하고 타깃 확률만 읽습니다. x가 나올 확률은 p(x), 다른 y가 나올 확률은 (1 − p(x))·p(y)/(1 − p(x)) = p(y)이므로 분포는 그대로 보존됩니다. 달라지는 것은 수락률뿐입니다. temperature 0 요청은 rejection_greedy_sample_kernel이 argmax 일치 여부로만 판정합니다.

Java로 규칙을 재현해 본 결과

아래는 vLLM 커널의 수락 조건과 recovered 샘플링 방식을 Java 25로 옮긴 코드의 핵심부 발췌입니다(빈도 집계용 main은 생략). 이번 실행에서 전체 코드로 5토큰 예시 분포를 seed 42, 200만 회 샘플링했습니다(vLLM 실행은 아닙니다).

import java.util.SplittableRandom;

public class SpecSample {
    static final SplittableRandom R = new SplittableRandom(42);

    // vLLM recovered 방식: argmax(w_i / E_i), E_i ~ Exp(1) → 정규화 없이 w에 비례
    static int race(double[] w) {
        int best = 0; double bestScore = -1;
        for (int i = 0; i < w.length; i++) {
            double s = w[i] / -Math.log(1.0 - R.nextDouble());
            if (s > bestScore) { bestScore = s; best = i; }
        }
        return best;
    }

    static int sample(double[] q) {
        double u = R.nextDouble(), acc = 0;
        for (int i = 0; i < q.length; i++) { acc += q[i]; if (u < acc) return i; }
        return q.length - 1;
    }

    // mode: prob / greedy(argmax를 one-hot으로) / naive(기각 시 p에서 재샘플)
    static int step(String mode, double[] p, double[] q, int argmaxQ) {
        boolean greedy = mode.equals("greedy");
        int x = greedy ? argmaxQ : sample(q);
        double qx = greedy ? 1.0 : q[x];
        if (qx > 0 && p[x] / qx >= R.nextDouble()) return x;  // 수락 조건
        if (mode.equals("naive")) return sample(p);           // 잘못된 구현
        double[] w = new double[p.length];
        for (int i = 0; i < p.length; i++)
            w[i] = Math.max(p[i] - (greedy ? (i == x ? 1 : 0) : q[i]), 0);
        return race(w);                                       // max(p - q, 0)에서 재샘플
    }
}

예시 분포는 p = [0.50, 0.25, 0.15, 0.07, 0.03], q = [0.30, 0.40, 0.10, 0.15, 0.05]이고, 결과의 L1 거리는 Σ|빈도 − p|입니다.

방식	수락률	출력 빈도	p와의 L1 거리
prob	0.7498	0.5001 0.2496 0.1500 0.0701 0.0301	0.0008
greedy	0.2496	0.5004 0.2496 0.1497 0.0702 0.0301	0.0013
naive	0.7500	0.4254 0.3124 0.1372 0.0875 0.0375	0.1748

수락률은 각각 Σ min(p, q) = 0.75, p(argmax q) = 0.25와 맞고, 두 방식 모두 출력이 p에 붙어 있습니다. naive는 첫 토큰이 0.425로 내려갔는데, min(0.5, 0.3) + 0.25 × 0.5 = 0.425로 이중 계산과 정확히 맞습니다.

greedy가 0.25로 낮은 것은 이 예시에서 드래프트 argmax가 타깃의 1순위가 아니어서입니다. 두 수락률 사이에 일반적인 대소 관계는 없고, 타깃이 뾰족하고 argmax가 같으면 greedy 쪽이 더 높을 수도 있습니다.

수락률은 속도만 바꾸고 정답은 바꾸지 않는다

수락률이 얼마나 중요한지는 논문 식 (1)이 보여 줍니다. 위치별 수락이 i.i.d.라는 단순화 가정 아래 한 반복의 기대 토큰 수는 (1 − α^(γ+1)) / (1 − α)입니다. 이번 실행에서 값을 넣어 계산하면 α = 0.75일 때 γ = 3이면 2.734개, γ = 7이면 3.600개입니다. α = 0.9면 같은 γ에서 3.439개와 5.695개입니다. α가 낮으면 γ를 늘려도 기대값이 금방 포화합니다. 논문은 Theorem 3.8에서 α와 드래프트 비용 계수 c로 속도 향상을 표현합니다.

vLLM에서 출력 품질이 걱정되신다면 rejection_sample_method를 보면 됩니다. 기본값 "standard"가 이 규칙을 따르고, 독스트링상 lossy인 "fly"와 정해 둔 확률로 수락하는 "synthetic"은 이 보장 밖입니다.

기각된 자리를 수락 경로가 채우지 못한 몫에서만 다시 뽑는 한, 드래프트가 서툴러도 출력 분포는 타깃의 것입니다.

참고 자료

• Yaniv Leviathan, Matan Kalman, Yossi Matias, 「Fast Inference from Transformers via Speculative Decoding」, arXiv:2211.17192 — Algorithm 1, Definition 3.1·3.2, Lemma 3.3, Theorem 3.5, 식 (1), Appendix A.1
• vllm-project/vllm main@5689c27 — vllm/v1/sample/rejection_sampler.py (RejectionSampler, rejection_random_sample_kernel, rejection_greedy_sample_kernel, sample_recovered_tokens_kernel, generate_uniform_probs, apply_sampling_constraints)
• vllm-project/vllm main@5689c27 — vllm/config/speculative.py (SpeculativeConfig.draft_sample_method, rejection_sample_method)