Speculative Decoding · AI infra interview lesson

Why It Matters

A draft proposes three tokens and the target accepts the first two. What is committed and which cache positions are discarded? Start with greedy verification, then derive sampling acceptance and full-round cost.

A concrete ledger: LLaMA-2-7B has 6,738,415,616 parameters, which is 13.48 GB in FP16. An H100 SXM has about 3.35 TB/s of HBM bandwidth, so at batch 1 every decode step takes at least 13.48e9 / 3.35e12 ≈ 4.0 ms, capping you at about 250 tokens/s. Verifying 5 positions in one forward still reads the weights only once. With per-position acceptance a=0.8 and γ=4, each round commits (1 - 0.8^5) / (1 - 0.8) = 3.36 tokens on average. If a draft step costs 5% of a target step, the ideal speedup is about 3.36 / (1 + 4 * 0.05) = 2.8x.

Walk the Shapes First

Setup: committed history h = [5, 9] (P=2), vocab V=10, draft length γ=3, sampling mode (not greedy).

If everything is accepted (k=3), commit y + [bonus] with the bonus sampled from p[3], γ+1=4 tokens in total.

Rules read off the table: the target input has length γ+1 and produces γ+1 distributions; committed length = accepted prefix + 1; candidates after the first rejection do not count, accepted or not; the KV cache rolls back to the "computed" length.

Running the examples: execute this page in order in one Python environment. Later snippets may reuse earlier definitions. CPU examples validate calculations and state changes, not GPU timing. Download the notebook for the complete exercise workspace.

Core Implementations

Step 1: Greedy Verification

What is new: when the target decodes greedily, verification is just "does the draft token equal the target argmax".

Rule: target = argmax(logits, -1) (γ+1 of them), k = cumprod(draft == target[:-1]).sum(), commit draft[:k] + [target[k]]. When all are accepted, target[γ] is the bonus.

Check: a toy target (logits depend only on the last two tokens) generates 40 tokens greedily one at a time; a noisy draft then runs speculative decoding, and the two sequences must match token by token.

import torch
torch.manual_seed(0)

V = 8
A_t, B_t = torch.randn(V, V), torch.randn(V, V)   # target：看最后两个 token
A_d = A_t + 0.3 * torch.randn(V, V)                # draft：target 的有噪声近似

def target_logits(ctx):            # ctx: list[int] -> (V,)
    return A_t[ctx[-1]] + B_t[ctx[-2]]

def draft_logits(ctx):
    return A_d[ctx[-1]] + B_t[ctx[-2]]

def spec_verify_greedy(draft, logits):
    """draft: (γ,) long；logits: (γ+1, V)，第 i 行是"看完 prefix + draft[:i]"后的 target logits。"""
    target = logits.argmax(-1)                      # (γ+1,)
    match = (draft == target[:-1]).long()           # (γ,)
    k = int(match.cumprod(0).sum())                 # 连续接受前缀长度
    return draft[:k].tolist() + [int(target[k])]    # 拒绝处换成 target 的 token；全接受时 k=γ，是 bonus

def generate_greedy(prompt, n):                     # 参考：逐 token 跑 target
    seq = list(prompt)
    while len(seq) < len(prompt) + n:
        seq.append(int(target_logits(seq).argmax()))
    return seq

def generate_spec(prompt, n, gamma):
    seq, rounds = list(prompt), 0
    while len(seq) < len(prompt) + n:
        ctx = list(seq)
        for _ in range(gamma):                      # draft 串行猜 γ 个
            ctx.append(int(draft_logits(ctx).argmax()))
        draft = torch.tensor(ctx[len(seq):])
        # 真实系统里这是一次 target forward；这里逐个前缀算出 γ+1 行 logits
        logits = torch.stack([target_logits(ctx[:len(seq) + i]) for i in range(gamma + 1)])
        seq += spec_verify_greedy(draft, logits)
        rounds += 1
    return seq[:len(prompt) + n], rounds

ref = generate_greedy([1, 2], 40)
for gamma in [1, 3, 5]:
    out, rounds = generate_spec([1, 2], 40, gamma)
    assert out == ref
    print(f"gamma={gamma}: rounds={rounds}, tokens/round={40 / rounds:.2f}")

# 边界：全接受 → γ+1 个；第一个就拒 → 1 个
logits = torch.zeros(4, V); logits[torch.arange(4), torch.tensor([3, 1, 4, 6])] = 1.0
assert spec_verify_greedy(torch.tensor([3, 1, 4]), logits) == [3, 1, 4, 6]
assert spec_verify_greedy(torch.tensor([5, 1, 4]), logits) == [3]
assert spec_verify_greedy(torch.tensor([3, 0, 4]), logits) == [3, 1]

The closer the draft is to the target, the more tokens per round. However bad the draft is, the output still equals plain greedy; only the number of rounds changes.

Step 2: Rejection Sampling at One Position

What is new: in sampling mode, accept with a probability and resample from the residual on rejection.

• Acceptance: a(y) = min(1, p(y) / q(y))
• Residual: r(x) = max(p(x) - q(x), 0) / sum_z max(p(z) - q(z), 0)
• Proof: the accept path contributes q(x) * min(1, p/q) = min(p(x), q(x)); rejection probability beta = 1 - sum_x min(p, q) = sum_x max(p - q, 0); the reject path contributes beta * r(x) = max(p(x) - q(x), 0). Together: min(p, q) + max(p - q, 0) = p.
• Total acceptance rate A = sum_x min(p, q) = 1 - TV(p, q).

With p=[0.6, 0.3, 0.1] and q=[0.2, 0.5, 0.3]: the accept path gives [0.2, 0.3, 0.1], beta=0.4, the residual puts everything on token 0 and adds back 0.4, giving exactly p. The wrong rule "resample from p after rejection" gives [0.44, 0.42, 0.14].

import torch
torch.manual_seed(0)

def residual_dist(p, q):
    """norm(max(p - q, 0))；p == q 时拒绝概率为 0，返回 p 占位，避免 0/0。"""
    r = (p - q).clamp(min=0)
    s = r.sum(-1, keepdim=True)
    return torch.where(s > 1e-12, r / s.clamp(min=1e-12), p)

def exact_output_dist(p, q):
    accept_path = torch.minimum(p, q)               # q(x) * min(1, p/q)
    beta = 1 - accept_path.sum()                    # 拒绝概率 = sum max(p-q, 0)
    return accept_path + beta * residual_dist(p, q), accept_path.sum()

p = torch.tensor([0.6, 0.3, 0.1], dtype=torch.float64)   # target
q = torch.tensor([0.2, 0.5, 0.3], dtype=torch.float64)   # draft
out, A = exact_output_dist(p, q)
print("accept rate:", A.item(), " residual:", residual_dist(p, q).tolist())
assert torch.allclose(out, p)
assert torch.isclose(A, 1 - 0.5 * (p - q).abs().sum())  # A = 1 - TV(p, q)

# 错误写法：拒绝后直接从 p 重采样
wrong = torch.minimum(p, q) + (1 - A) * p
assert torch.allclose(wrong, torch.tensor([0.44, 0.42, 0.14], dtype=torch.float64))

# Monte Carlo：一次采 N 个，向量化
N = 200_000
y = torch.multinomial(q.float(), N, replacement=True)          # draft 提议
u = torch.rand(N)
accept = u < (p[y] / q[y]).clamp(max=1).float()
fix = torch.multinomial(residual_dist(p, q).float(), N, replacement=True)
tok = torch.where(accept, y, fix)
emp = torch.bincount(tok, minlength=3).double() / N
tv = 0.5 * (emp - p).abs().sum().item()
print("empirical:", [round(v, 4) for v in emp.tolist()], " TV:", round(tv, 4))
assert tv < 0.01

# 边界：p == q 时不除零，接受率 1
same, A1 = exact_output_dist(p, p)
assert torch.allclose(same, p) and torch.isclose(A1, torch.tensor(1.0, dtype=torch.float64))
assert not torch.isnan(residual_dist(p, p)).any()

Step 3: Verifying a Whole Chain, and Expected Committed Length

What is new: apply Step 2 left to right over γ positions, stop at the first rejection; if all pass, sample a bonus from p[γ] . (Pro lesson)

Count draft and verification time together

Accepting more tokens helps only when the full round is cheap enough. (Pro lesson)

Step 4: N-gram Proposals (Prompt Lookup)

What is new: no draft model at all, just search the token history. (Pro lesson)

Step 5: Tree Attention Mask

What is new: candidates form a tree instead of a chain (Medusa / EAGLE keep top-k per position). (Pro lesson)

In Real Systems

nano-vllm : no speculative decoding. ModelRunner.prepare_decode puts exactly one token per sequence ( input_ids.append(seq.last_token) ), Bl… (Pro lesson)

Common Pitfalls

Pattern: on rejection, torch.multinomial(p[i], 1) . (Pro lesson)

Say It in the Interview

"Speculative decoding exploits the fact that small-batch decode is memory bound: a draft proposes γ tokens, and the target verifies all of t… (Pro lesson)

Hands-on Notebook

Each exercise in the notebook gives a statement and a function skeleton; write your code, then run the tests below it. (Pro lesson)