A 70B parameter model running on 4xA100 GPUs produces roughly 18-22 tokens per second per sequence. That works fine for batch offline processing - but for a real-time chat product, it translates to first tokens arriving 800ms after the user hits send, and visible letter-by-letter streaming that feels like watching paint dry. Users tolerate 200ms. At 800ms they bounce.

The fundamental constraint is that autoregressive decoding forces the model to execute a full forward pass - touching all 140 GB of weights across all 80 transformer layers - for every single token. A 200-token response requires 200 sequential forward passes. There is no way around this with standard decoding; parallelism at the batch level helps throughput but does nothing for per-sequence latency.

Speculative decoding reframes the problem. Think of it like a senior engineer reviewing a junior engineer’s code. The junior (a cheap 8B model) writes a plausible draft quickly. The senior (the 70B model) reviews the whole draft in a single focused pass - approving most of it, correcting the odd token, and adding one clean next step. The output is mathematically equivalent to what the senior would have produced alone, but the wall-clock time is dominated by the single review pass, not by writing each line from scratch.

At 80% acceptance rate with speculation depth gamma=4, a single 70B forward pass produces 3.5 tokens on average instead of 1. With gamma=6 and a high-quality draft model, we can push that to 4.8 tokens per pass. On Llama-3 70B, this translates from 20 tokens per second to 70-80 tokens per second, cutting TTFT from 800ms to under 150ms and TPOT (time per output token) from 50ms to 14ms.

Requirements and Constraints

Functional requirements:

• Accept a prompt and stream back generated tokens with identical statistical distribution to standard autoregressive decoding from the target model
• Support configurable speculation depth gamma (2-8 tokens)
• Handle rejection of draft tokens gracefully via rollback without visible output artifacts
• Support tree-based multi-candidate drafting (Medusa-style) in addition to linear drafting
• Dynamically adjust gamma per sequence based on observed acceptance rate

Non-functional requirements:

• p50 TTFT under 100ms, p99 TTFT under 150ms for prompts up to 512 tokens
• p50 TPOT under 18ms (target: 14ms with effective gamma around 4.5)
• Draft model memory overhead below 20% of target model HBM footprint
• Mathematically lossless: output distribution identical to sampling from target model alone
• Zero added latency when draft acceptance rate falls below threshold - fall back to standard decoding

Constraints:

• Draft model must share the same vocabulary and tokenizer as the target model (otherwise token IDs do not align for verification)
• Draft and target models must be co-located on the same GPU node to avoid PCIe transfer overhead on every speculation step
• Total GPU HBM budget: 80GB per A100, 4 GPUs per node (320 GB total); target model needs 140 GB, leaving 180 GB for KV caches and the draft model

High-Level Architecture

The system has five moving parts. The Speculation Scheduler manages incoming requests and maintains a live gamma configuration per active sequence. The Draft Model (Llama-3 8B) runs the fast autoregressive generation, producing gamma draft tokens with their probability distributions. The Speculation Engine takes those draft tokens plus the target model’s verification pass and runs the accept/reject loop. The Target Model (Llama-3 70B) runs a single forward pass over the entire input including draft tokens. The KV Cache Manager maintains separate caches for draft and target models, handling rollback atomically when tokens are rejected.

Component Deep Dives

The Accept/Reject Mechanism

The non-obvious property of speculative decoding is that it is mathematically lossless - the output distribution is identical to the target model sampling alone. This is not an approximation. It is guaranteed by the rejection sampling correction.

For each draft token d_i, the draft model produced it with probability q(d_i | x) and the target model assigns it probability p(d_i | x). The acceptance rule is: accept d_i with probability min(1, p(d_i|x) / q(d_i|x)). If rejected, resample from the residual distribution normalize(max(0, p(t|x) - q(t|x))) for all tokens t in the vocabulary.

This rule has two critical properties. First, if the draft model proposed d_i and the target agrees it’s at least as likely as the draft thought, it is always accepted (ratio >= 1). Second, if the target rejects, the correction token is sampled from exactly the distribution the target would have produced had it been sampling unconditionally. The composition of these two cases provably recovers the target’s marginal distribution at every position.

After any rejection, all subsequent draft tokens are discarded and the KV cache is rolled back to the position of the last accepted token. This is why rollback handling must be atomic - if the KV cache for the target gets out of sync with the accepted token position, subsequent forward passes will attend over a stale context and produce incorrect outputs.

import torch
from dataclasses import dataclass
from typing import Optional

@dataclass
class SpeculationResult:
    accepted_tokens: list[int]
    correction_token: Optional[int]
    n_accepted: int
    n_rejected: int
    rollback_to_position: int

def speculative_accept_reject(
    draft_token_ids: list[int],
    draft_logits: torch.Tensor,   # shape: [gamma, vocab_size]
    target_logits: torch.Tensor,  # shape: [gamma+1, vocab_size] - one extra for correction
    temperature: float = 1.0,
) -> SpeculationResult:
    """
    Run the speculative decoding accept/reject loop.
    Returns accepted tokens and an optional correction token.
    Output distribution is identical to sampling from target alone.
    """
    gamma = len(draft_token_ids)
    vocab_size = draft_logits.shape[-1]

    # Convert logits to probabilities
    if temperature != 1.0:
        draft_probs = torch.softmax(draft_logits / temperature, dim=-1)
        target_probs = torch.softmax(target_logits / temperature, dim=-1)
    else:
        draft_probs = torch.softmax(draft_logits, dim=-1)
        target_probs = torch.softmax(target_logits, dim=-1)

    accepted = []
    for i, draft_tok in enumerate(draft_token_ids):
        p = target_probs[i, draft_tok].item()
        q = draft_probs[i, draft_tok].item()
        # Clamp to avoid division-by-zero on zero-probability draft tokens
        q = max(q, 1e-10)
        accept_prob = min(1.0, p / q)

        if torch.rand(1).item() < accept_prob:
            accepted.append(draft_tok)
        else:
            # Rejection: sample correction from residual distribution
            residual = torch.clamp(target_probs[i] - draft_probs[i], min=0.0)
            residual_sum = residual.sum()
            if residual_sum > 1e-8:
                residual = residual / residual_sum
                correction = torch.multinomial(residual, num_samples=1).item()
            else:
                # Fallback: sample from target directly (numerical edge case)
                correction = torch.multinomial(target_probs[i], num_samples=1).item()
            return SpeculationResult(
                accepted_tokens=accepted,
                correction_token=correction,
                n_accepted=len(accepted),
                n_rejected=gamma - len(accepted),
                rollback_to_position=len(accepted),
            )

    # All draft tokens accepted - take one more free token from target's (gamma+1)-th position
    bonus_token = torch.multinomial(target_probs[gamma], num_samples=1).item()
    return SpeculationResult(
        accepted_tokens=accepted,
        correction_token=bonus_token,
        n_accepted=gamma,
        n_rejected=0,
        rollback_to_position=gamma,
    )