Speculative Decoding, Early Exit and Quantization: Faster LLM Decode

LLM inference runs in two phases. Prefill processes the whole prompt in parallel and can keep a GPU busy. Decode emits one token at a time, and each token depends on the one before it, so the hardware mostly waits. That sequential dependency is why interactive latency stays expensive even when the FLOPs look plentiful on paper.

Prefill phase:
  Input: [token_1, token_2, ..., token_512]  → all 512 tokens processed in parallel
  Matrix shape: [batch, 512, 4096]
  GPU utilization: high — large matrix, full Tensor Core throughput

Decode phase (one step):
  Input: [token_513]  → one token processed
  Matrix shape: [batch, 1, 4096]
  GPU utilization: low - tiny matrix, most CUDA cores idle

Why decode sits on the critical path

At each decode step the model runs queries against a growing KV cache. Work per step climbs with context length, and for chat you can't batch your way out freely, because latency SLOs cap how long a request may wait for company. The arithmetic shows how little of the chip a single step actually uses:

One attention head, decode step:

Query Q: [8, 1, 128]   →  8 × 128 = 1,024 elements
Keys  K: [8, 2048, 128] →  8 × 2048 × 128 = 2M elements
Values V: [8, 2048, 128] →  same

Matrix multiply (QKᵀ) per head:
  Shape: [8, 1, 128] × [8, 128, 2048] → [8, 1, 2048]
  FLOPs: 8 × 1 × 128 × 2048 ≈ 2.1M FLOPs per head

For 32 heads: ≈ 67M FLOPs per layer
For 32 layers: ≈ 2.1B FLOPs total per decode step

A100 peak at BF16: ~312 TFLOPS = 312 × 10¹² FLOPs/sec

Time to execute 2.1B FLOPs at 100% utilization:
  2.1 × 10⁹ / 312 × 10¹² ≈ 0.0067 ms

Actual observed decode latency per step: ~10–30ms

Effective compute utilization: < 0.1%

Under 0.1% effective utilization means decode is limited by memory traffic and step count, not raw compute. There are two ways to attack it: produce more accepted tokens per expensive step, or do less work per step. Speculative decoding does the first, early exit the second.

Speculative decoding: draft, then verify

A small draft model proposes several future tokens. The target model checks all of them in one parallel forward pass, keeps the longest agreeing prefix and resamples at the first rejection. Each accepted draft token is a token you didn't pay a full target step for.

Without speculative decoding:
  Generate 5 tokens: 5 × 30ms = 150ms

With speculative decoding (K=4):
  Draft 4 tokens: 4 × 1.5ms = 6ms
  1 target verification pass: ~35ms  (slightly longer than decode,
                                       processes K+1=5 positions)
  Expected accepted tokens per round:
    (1 - 0.80^5) / (1 - 0.80) ≈ 3.36 tokens
  Time per round: 6ms + 35ms = 41ms
  Time per token: 41ms / 3.36 ≈ 12.2ms
Speedup: 30ms → 12.2ms ≈ 2.5×

The speedup depends entirely on how often draft and target agree. Easy, low-entropy text accepts long drafts; a surprising token cuts the round short. Here's the loop in conceptual form (real implementations also manage the KV cache across both models):

import torch
import torch.nn.functional as F

def speculative_decode(target_model, draft_model, input_ids,
                       max_new_tokens, K=4, temperature=1.0):
    """
    Conceptual speculative decoding loop.
    Real implementations handle KV cache management across both models.
    """
    generated = input_ids.clone()
    while generated.shape[1] - input_ids.shape[1] < max_new_tokens:
        # --- Draft phase ---
        draft_tokens = []
        draft_probs = []
        draft_input = generated.clone()
        for _ in range(K):
            with torch.no_grad():
                draft_logits = draft_model(draft_input).logits[:, -1, :]
            q = F.softmax(draft_logits / temperature, dim=-1)
            token = torch.multinomial(q, num_samples=1)
            draft_tokens.append(token)
            draft_probs.append(q)
            draft_input = torch.cat([draft_input, token], dim=1)
        # --- Verify phase: one target forward pass over all K+1 positions ---
        verify_input = torch.cat([generated] + draft_tokens, dim=1)
        with torch.no_grad():
            target_logits = target_model(verify_input).logits
        # target_logits[:, -K-1:, :] covers all K draft positions + bonus
        # --- Accept/reject ---
        accepted = 0
        for i in range(K):
            p = F.softmax(target_logits[:, -(K+1)+i, :] / temperature, dim=-1)
            q = draft_probs[i]
            token = draft_tokens[i]
            # Acceptance probability
            accept_prob = torch.min(
                torch.ones_like(p.gather(1, token)),
                p.gather(1, token) / (q.gather(1, token) + 1e-9)
            )
            if torch.rand(1) < accept_prob:
                generated = torch.cat([generated, token], dim=1)
                accepted += 1
            else:
                # Sample corrected token and stop this round
                corrected_dist = F.relu(p - q)
                corrected_dist = corrected_dist / corrected_dist.sum(dim=-1, keepdim=True)
                corrected_token = torch.multinomial(corrected_dist, num_samples=1)
                generated = torch.cat([generated, corrected_token], dim=1)
                break
        else:
            # All K accepted - take bonus token
            bonus_logits = target_logits[:, -1, :]
            p_bonus = F.softmax(bonus_logits / temperature, dim=-1)
            bonus_token = torch.multinomial(p_bonus, num_samples=1)
            generated = torch.cat([generated, bonus_token], dim=1)
    return generated

Pick a draft from the same family when you can, such as a distilled or quantized smaller sibling, and measure acceptance rate on your own traffic rather than on public blog examples.

When drafts diverge

If the draft's distribution drifts away from the target's, acceptance collapses and you pay the draft cost for almost nothing. Monitor acceptance continuously and fall back to plain decode when it drops.

A quick sanity check with round numbers: say a target step costs 10 ms and the draft proposes 5 tokens, of which about 3 are accepted on average (60%). The cost per accepted token falls compared with one-token steps, even after the draft overhead, as long as acceptance stays healthy. If acceptance sinks to about one token, the scheme loses. That sensitivity is why dashboards beat anecdotes.

Early exit: skip deep layers on easy tokens

Some architectures let a token leave the network at an intermediate layer once confidence is high, which saves compute on punctuation, articles and other predictable tokens.

Layer distribution of exits:Exit at layers 1–8  (very easy tokens like punctuation, articles):  15%
Exit at layers 9–16 (medium tokens, common continuations):           35%
Exit at layers 17–24 (harder tokens, named entities, numbers):       30%
Exit at layers 25–32 (full computation required):                    20%Weighted average layers executed:
  0.15 × 6 + 0.35 × 12 + 0.30 × 20 + 0.20 × 32
  = 0.90 + 4.20 + 6.00 + 6.40
  = 17.5 layers averageSpeedup vs always running 32 layers:
  32 / 17.5 ≈ 1.83×

Unlike speculative decoding, early exit does change outputs when it skips layers, by design. It trades quality for speed and depends heavily on how well confidence is calibrated, so budget the accuracy delta and measure it.

Stacking with quantization, batching and the KV cache

Speculative decoding changes the token proposal loop by adding a second model; early exit changes depth inside one model. They hit related bottlenecks in different ways, and both can sometimes stack with quantization, continuous batching and KV cache paging:

Combined stack example:

Target: 7B model, BF16, FlashAttention, PagedAttention
  → Model: ~14 GB, memory-efficient attention, paged KV cache

Draft: 70M model, INT4, FlashAttention
  → Model: ~35 MB, near-zero memory overhead

Speculative decode:
  → with K=4, α=0.80
  → ~2.5× token generation speedup on long outputs

Full stack speedup vs FP32 no-optimization baseline:
  - Quantization: 2–2.3× throughput
  - FlashAttention: 20–40% attention latency reduction
  - Speculative decoding: 2–2.5× decode speedup (on eligible requests)

Combined: 5–8× improvement in end-to-end tokens/second

Stacking takes care. Draft lengths interact with the assumptions of a continuous batching scheduler, and a poor scheduler fragments batches and hurts utilization. KV cache eviction policies need tuning so long chats don't thrash. Prompt caching on the prefill side matters too, otherwise you "fix decode" and leave prefill wasteful. And stacking early exit with aggressive quantization can push hallucinations up quickly. A coherent serving design beats a single paper transplanted into the hot path.

When each technique pays off

Speculative decoding helps when draft agreement is high and target steps dominate latency. It fails when drafts rarely match or the draft overhead eats the gain. Early exit helps when many tokens are easy and the accuracy budget allows it; it fails on hard tokens or when confidence is poorly calibrated.

A concrete picture: a 7B target serving chat with replies of 50 to 150 tokens. Prefilling a 2k-token system prompt is heavy but happens rarely per turn, while decode steps dominate the delay users notice. Cutting the number of decode steps with speculative accepts, or the work per step with early exit, is what they actually feel.

Does it matter for tiny on-device models? Usually less than for large targets. Can the draft run on CPU? Sometimes; measure it. As a priority order, fix batching and caching first, then add speculative decoding, then early exit if your stack supports it.

Putting it into production LLM inference

• Baseline vanilla p95 latency and quality before changing anything.
• Colocate the draft model with the target on the same node. Cross-host drafts add network jitter that can erase the gain, and two models plus the KV cache can run a box out of memory that comfortably held one.
• Log acceptance histograms, mean accepted length, quality eval deltas, GPU utilization and p95 latency.
• Watch the CPU and GPU balance when drafts run on different devices.
• A/B quality on factual suites. Before a global rollout, run fixed prompts across factual QA, coding and refusal suites and compare token-identical rates when speculative decoding is configured for exact distribution match. For early exit, track win rate on graded tasks and human preference where you have it.
• Put each optimization behind a feature flag with a kill switch back to vanilla decode, and rehearse the rollback.
• Revisit everything after a tokenizer or fine-tune change.

A correctly implemented speculative decoder should match the target distribution, but verify that with paired tests instead of trusting it. Using an unrelated draft, ignoring how temperature affects acceptance, or declaring victory from throughput benchmarks without quality gates are the usual ways this goes wrong, and every one of them is measurable.

Platform teams should own these optimizations in the serving layer so each app team doesn't invent its own draft selection. Coordinate with the serving maintainers instead of flipping flags in app code, expose headers or trace attributes showing whether speculative decoding or early exit applied, and tag token bills by optimization so finance can see the win.

The bottleneck that's left

Once decode gets faster, users may still be waiting on tool calls, retrieval or client-side markdown rendering. The remaining limit can be memory bandwidth, the network or the client, not FLOPs, so trace end to end and profile before stacking every trick from the papers. Speculative decoding also doesn't replace good retrieval or good prompts; it only makes the next tokens cheaper when the model already knows what it wants to say.

Sequential decode is the structural tax of autoregressive generation. Speculative decoding and early exit reduce that tax under conditions you can measure, so ship them like any production feature: metrics, flags, rollbacks and a clear owner in the serving platform team.