Speculative Decoding

The decoding bottleneck

Autoregressive LLM generation is sequential: you generate one token at a time, each step requires a full forward pass through the model. For a 70B parameter model, each token costs ~140 GFLOPS. The throughput bottleneck is memory bandwidth — loading all 140B parameters from HBM to compute units to generate a single token. On an A100 80GB, this limits you to roughly 50 tokens/second for a 70B model.

Speculative decoding exploits the observation that a small model (e.g. 7B) can correctly predict the next 3–5 tokens most of the time. If we can verify all those predictions in a single large-model forward pass, we get multiple tokens per expensive call.

Speculative decoding algorithm

At each step of speculative decoding:

1. Draft: Run the small draft model autoregressively for γ steps, generating γ draft tokens x̃₁, ..., x̃_γ with draft probabilities q(x̃_i).
2. Verify: Run the large target model on the original context + all γ draft tokens in a single forward pass (possible because it's not causal during verification). Get target probabilities p(x̃_i) for each draft position plus p(x_{γ+1}) for the next new token.
3. Accept/reject: For each draft token in order, accept it with probability min(1, p(x̃_i)/q(x̃_i)). If rejected, sample a correction token from a normalised residual distribution. Once any token is rejected, discard all subsequent draft tokens.
4. On average α × γ tokens are accepted per verification step, where α is the acceptance rate (typically 0.7–0.9 for a good draft model).

Acceptance sampling

The acceptance criterion min(1, p/q) is the key insight. It guarantees that the output distribution of speculative decoding is identical to that of the target model — no quality degradation. The math: if we accept x̃ with probability p(x̃)/q(x̃) when p ≤ q, and always accept when p ≥ q, the resulting marginal distribution equals p.

When a token is rejected at position i, we sample from the residual distribution max(0, p - q) / Z rather than simply resampling from p. This ensures no token is wasted — the rejection produces a valid target sample that's different from the draft token.

Python implementation sketch

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

def speculative_decode(draft_model, target_model, tokenizer, prompt, gamma=4, max_new=100):
    input_ids = tokenizer(prompt, return_tensors="pt").input_ids.cuda()
    generated = input_ids.clone()

    for _ in range(max_new // gamma + 1):
        # Step 1: draft gamma tokens
        draft_ids = generated.clone()
        draft_probs = []
        for _ in range(gamma):
            with torch.no_grad():
                logits = draft_model(draft_ids).logits[:, -1, :]
            probs = torch.softmax(logits, dim=-1)
            next_tok = torch.multinomial(probs, 1)
            draft_probs.append(probs[0, next_tok[0, 0]].item())
            draft_ids = torch.cat([draft_ids, next_tok], dim=-1)

        # Step 2: verify with target model (single forward pass)
        with torch.no_grad():
            target_logits = target_model(draft_ids).logits[0, generated.shape[1]-1:, :]
        target_probs = torch.softmax(target_logits, dim=-1)

        # Step 3: accept/reject
        accepted = 0
        for i in range(gamma):
            draft_tok = draft_ids[0, generated.shape[1] + i].item()
            p = target_probs[i, draft_tok].item()
            q = draft_probs[i]
            if torch.rand(1).item() <= p / q:
                accepted += 1
            else:
                break

        new_tokens = draft_ids[:, generated.shape[1]:generated.shape[1] + accepted]
        generated = torch.cat([generated, new_tokens], dim=-1)

        # Add one target-sampled token
        bonus_tok = torch.multinomial(target_probs[accepted], 1).unsqueeze(0)
        generated = torch.cat([generated, bonus_tok], dim=-1)

    return tokenizer.decode(generated[0], skip_special_tokens=True)

Draft model selection

The draft model is the most important hyperparameter. Requirements: (1) same vocabulary as the target model, (2) fast enough that γ draft steps cost less than 1 target step, (3) high acceptance rate ≥ 0.7. Good pairs:

• Llama 3 70B (target) + Llama 3 8B (draft): acceptance rate ~0.75–0.85 on typical prompts
• GPT-4 class target + GPT-3.5-turbo draft (used internally by some providers)
• Self-speculative decoding: use early exit layers of the same model as the draft

For a 70B/7B pair, if acceptance rate α=0.8 and γ=4, the expected tokens per target call is α × γ + 1 ≈ 4.2, giving ~4× fewer target calls and ~2–3× wall-clock speedup (draft overhead reduces the theoretical gain).

Throughput vs latency impact

Speculative decoding primarily reduces latency (time for a single request), not throughput (total tokens/second across many requests). When the GPU is already fully utilised with batched requests, speculative decoding may not help — the verification step has the same batch dimension as the target-only baseline. It shines in:

• Interactive chat (single-user, latency-critical)
• Edge deployment (single GPU, low concurrency)
• Long generation tasks where p99 latency matters

Gotchas

• Memory: Running two models simultaneously doubles GPU memory requirements. On memory-constrained hardware, consider self-speculative decoding (no second model) or medusa heads instead.
• Temperature 0: At temperature=0 (greedy decoding), acceptance becomes deterministic (always accept if draft token = argmax). Acceptance rate is essentially 1 for easy tokens.
• Batching: Speculative decoding with batched requests is harder because different sequences in the batch may reject at different steps. vLLM and TGI implement efficient batched speculative decoding.
• Medusa alternative: Medusa adds multiple prediction heads to the target model to generate draft tokens without a separate model — lower memory cost but requires fine-tuning.