Build a Speculative Decoding Inference Engine: A Hands-On CV Project

mkdir speculative-decoding && cd speculative-decoding
python -m venv venv && source venv/bin/activate
pip install torch transformers numpy tiktoken accelerate

# draft_model.py
import numpy as np
from collections import defaultdict, Counter
from typing import List, Tuple

class NgramDraftModel:
    """Kneser-Ney smoothed trigram language model for speculative drafting."""

    def __init__(self, tokenized_corpus: List[List[int]], vocab_size: int,
                 k: int = 5, discount: float = 0.75):
        self.vocab_size = vocab_size
        self.k = k                    # number of draft tokens to generate
        self.discount = discount

        # Count n-grams
        self.trigram_counts = Counter()
        self.bigram_counts = Counter()
        self.unigram_counts = Counter()
        self.bigram_continuation = defaultdict(Counter)  # for Kneser-Ney

        for sentence in tokenized_corpus:
            padded = [0] * 2 + sentence + [1]  # 0=BOS, 1=EOS
            for i in range(len(padded) - 2):
                w1, w2, w3 = padded[i], padded[i+1], padded[i+2]
                self.trigram_counts[(w1, w2, w3)] += 1
                self.bigram_counts[(w1, w2)] += 1
                self.unigram_counts[w2] += 1
                self.bigram_continuation[(w1, w2)][w3] += 1

        # Precompute log-probabilities (with Laplace smoothing fallback)
        self._log_probs = {}
        self._build_log_probs()

    def _build_log_probs(self):
        total_unigrams = sum(self.unigram_counts.values())
        num_bigrams = len(self.bigram_counts)

        for (w1, w2, w3), count in self.trigram_counts.items():
            bigram_count = self.bigram_counts[(w1, w2)]
            continuation_count = len(self.bigram_continuation[(w1, w2)])

            # Kneser-Ney: (max(c(w1,w2,w3) - D, 0) / c(w1,w2)) * 
            #              (c(w2,w3) / num_unique_successors_of_w2)
            # Plus continuation penalty for unseen trigrams
            numerator = max(count - self.discount, 0)
            denom = bigram_count
            lambda_w2 = self.discount * num_bigrams / bigram_count if bigram_count > 0 else 0

            if denom > 0 and numerator > 0:
                self._log_probs[(w1, w2, w3)] = np.log(numerator / denom)

        # Fallback: for unseen trigrams, use bigram continuation probability
        for (w1, w2), conts in self.bigram_continuation.items():
            for w3 in conts:
                if (w1, w2, w3) not in self._log_probs:
                    self._log_probs[(w1, w2, w3)] = np.log(
                        conts[w3] / sum(conts.values()) + 1e-10
                    )

    def draft_tokens(self, context: List[int], num_tokens: int = None) -> List[int]:
        """Generate `k` candidate tokens autoregressively from context."""
        if num_tokens is None:
            num_tokens = self.k

        draft = []
        current_context = context[-2:]  # last two tokens for trigram

        for _ in range(num_tokens):
            candidates = []
            scores = []

            for token_id in range(self.vocab_size):
                key = (current_context[0], current_context[1], token_id)
                if key in self._log_probs:
                    candidates.append(token_id)
                    scores.append(self._log_probs[key])
                else:
                    # Fallback to unigram
                    score = np.log((self.unigram_counts.get(token_id, 0) + 1) /
                                   (sum(self.unigram_counts.values()) + self.vocab_size))
                    candidates.append(token_id)
                    scores.append(score)

            scores = np.array(scores)
            probs = np.exp(scores - scores.max())  # numerical stability
            probs /= probs.sum()

            next_token = np.random.choice(candidates, p=probs)
            draft.append(next_token)
            current_context = [current_context[1], next_token]

        return draft

# verify_model.py
import torch
import torch.nn.functional as F
from transformers import GPT2LMHeadModel, GPT2Tokenizer
from typing import List

class VerificationTransformer:
    """Verification model that scores draft tokens in a single joint forward pass."""

    def __init__(self, model_name: str = "gpt2"):
        self.tokenizer = GPT2Tokenizer.from_pretrained(model_name)
        self.tokenizer.pad_token = self.tokenizer.eos_token
        self.model = GPT2LMHeadModel.from_pretrained(model_name)
        self.model.eval()
        if torch.cuda.is_available():
            self.model.to("cuda")

    def verify(self, prompt_ids: List[int], draft_ids: List[int]) -> List[float]:
        """
        Verify draft tokens against the model.
        Returns a list of acceptance probabilities for each draft position.

        The key insight: we run ONE forward pass over
        [prompt_ids + draft_ids] and get logits for every position.
        """
        input_ids = torch.tensor([prompt_ids + draft_ids],
                                 device=self.model.device)

        with torch.no_grad():
            outputs = self.model(input_ids)
            logits = outputs.logits  # shape: (1, seq_len, vocab_size)

        # Extract logits for each draft position
        # Position i in draft corresponds to position len(prompt_ids) + i
        draft_logits = logits[0, len(prompt_ids):, :]
        draft_input_ids = input_ids[0, len(prompt_ids):]

        acceptance_probs = []
        cumulative_draft = draft_ids.copy()

        for i in range(len(draft_ids)):
            # Position-i logits: probability of draft[i] given
            # prompt + draft[0..i-1]
            pos_logits = draft_logits[i]
            pos_ids = draft_input_ids[i]

            # Compute P_verify(draft[i] | context) using softmax
            probs = F.softmax(pos_logits, dim=-1)
            p_verify = probs[pos_ids].item()

            # Compute P_draft(draft[i] | context) from the draft model
            # This is passed in from outside; we compute the acceptance ratio here
            # For now, return the verification probability
            acceptance_probs.append(p_verify)

        return acceptance_probs

    def generate_single(self, input_ids: List[int], num_tokens: int = 1) -> List[int]:
        """Fallback: generate tokens autoregressively using the verification model."""
        ids = torch.tensor([input_ids], device=self.model.device)
        with torch.no_grad():
            output = self.model.generate(ids, max_new_tokens=num_tokens,
                                         do_sample=False)
        new_tokens = output[0][len(input_ids):].tolist()
        return new_tokens

# speculative_engine.py
import torch
import torch.nn.functional as F
import numpy as np
from typing import List, Tuple

class SpeculativeDecodingEngine:
    """
    Orchestrates speculative decoding between a draft model and a verification model.
    
    Algorithm (Bow-style, from Edward et al., 2022):
    1. Draft k tokens using the n-gram model.
    2. Run verification transformer on [prompt + draft] in ONE forward pass.
    3. For each position i, accept draft[i] with probability:
       min(1, P_verify(draft[i]) / P_draft(draft[i]))
    4. If rejected at position i, use verification model to generate
       a single replacement token and stop this round.
    5. Repeat from step 1.
    """

    def __init__(self, draft_model, verify_model, draft_k: int = 5):
        self.draft_model = draft_model
        self.verify_model = verify_model
        self.k = draft_k

    def _compute_acceptance_prob(self, p_verify: float, p_draft: float) -> float:
        """Acceptance probability per the Bow algorithm."""
        if p_draft == 0:
            return 0.0
        return min(1.0, p_verify / p_draft)

    def generate(self, prompt_text: str, max_new_tokens: int = 50) -> str:
        """Generate text using speculative decoding."""
        # Tokenize prompt
        prompt_ids = self.verify_model.tokenizer.encode(prompt_text)
        generated = prompt_ids.copy()

        tokens_generated = 0
        speculation_rounds = 0
        accepted_total = 0
        draft_total = 0

        while tokens_generated < max_new_tokens:
            # --- Step 1: Draft ---
            draft_tokens = self.draft_model.draft_tokens(
                generated, num_tokens=min(self.k, max_new_tokens - tokens_generated)
            )
            draft_total += len(draft_tokens)

            # --- Step 2: Verify (single joint forward pass) ---
            acceptance_probs = self.verify_model.verify(generated, draft_tokens)

            # --- Step 3: Acceptance Sampling ---
            accepted_in_round = 0
            for i, draft_token in enumerate(draft_tokens):
                # Get draft probability from n-gram model
                context = generated[-2:] if len(generated) >= 2 else [0, 0]
                p_draft = self.draft_model._log_probs.get(
                    (context[0], context[1], draft_token), 1e-10
                )
                p_draft = np.exp(p_draft)

                p_verify = acceptance_probs[i]
                p_accept = self._compute_acceptance_prob(p_verify, p_draft)

                if np.random.random() < p_accept:
                    generated.append(draft_token)
                    accepted_in_round += 1
                    accepted_total += 1
                else:
                    # --- Rejection: generate single replacement ---
                    replacement = self.verify_model.generate_single(
                        generated[:i+1], num_tokens=1
                    )
                    generated.extend(replacement)
                    accepted_total += 1  # still accepted one token
                    break  # stop this round, restart drafting

            tokens_generated += accepted_in_round if accepted_in_round > 0 else 1
            speculation_rounds += 1

            if tokens_generated >= max_new_tokens:
                break

        # Decode output
        output_text = self.verify_model.tokenizer.decode(
            generated[len(prompt_ids):], skip_special_tokens=True
        )

        # Print stats
        speedup = draft_total / max(accepted_total, 1)
        print(f"\n--- Speculative Decoding Stats ---")
        print(f"Tokens generated: {tokens_generated}")
        print(f"Speculation rounds: {speculation_rounds}")
        print(f"Draft tokens proposed: {draft_total}")
        print(f"Tokens accepted: {accepted_total}")
        print(f"Approximate speedup factor: {speedup:.2f}x")

        return output_text

# main.py
from draft_model import NgramDraftModel
from verify_model import VerificationTransformer
from speculative_engine import SpeculativeDecodingEngine

# --- Load a small training corpus for the n-gram model ---
# In practice, use a large text corpus (e.g., WikiText-2 or a subset of
# your target domain). Here we use a minimal example.

training_corpus = [
    [1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
    [1, 2, 3, 10, 9, 8, 7, 6, 5, 4],
    [1, 5, 9, 2, 6, 10, 3, 7, 8, 4],
    # ... in practice, tokenize a real corpus with tiktoken or sentencepiece
]

# For a real project, tokenize properly:
# import tiktoken
# enc = tiktoken.get_encoding("p50k_base")
# tokenized_corpus = [enc.encode(text) for text in open("corpus.txt")]

VOCAB_SIZE = 50257  # GPT-2 vocabulary size

# Initialize components
draft_model = NgramDraftModel(
    tokenized_corpus=training_corpus,
    vocab_size=VOCAB_SIZE,
    k=5
)

verify_model = VerificationTransformer(model_name="gpt2-medium")

# Build and run the engine
engine = SpeculativeDecodingEngine(
    draft_model=draft_model,
    verify_model=verify_model,
    draft_k=5
)

prompt = "The future of artificial intelligence depends on"
output = engine.generate(prompt, max_new_tokens=30)
print(f"\nPrompt: {prompt}")
print(f"Output: {output}")

# Activate your environment
source venv/bin/activate

# Run the full pipeline
python main.py

# test_engine.py
import pytest
import torch
from speculative_engine import SpeculativeDecodingEngine

def test_acceptance_probability_bounds():
    """Acceptance probability must be in [0, 1]."""
    engine = SpeculativeDecodingEngine.__new__(SpeculativeDecodingEngine)
    # p_verify > p_draft → accept with prob 1
    assert engine._compute_acceptance_prob(0.9, 0.3) == 1.0
    # p_verify < p_draft → accept with prob ratio
    assert engine._compute_acceptance_prob(0.3, 0.9) == pytest.approx(0.333, rel=0.01)
    # p_draft = 0 → never accept
    assert engine._compute_acceptance_prob(0.5, 0.0) == 0.0

def test_draft_model_returns_correct_length():
    """Draft model should always return exactly k tokens."""
    from draft_model import NgramDraftModel
    model = NgramDraftModel(
        tokenized_corpus=[[1,2,3,4,5]] * 100,
        vocab_size=100,
        k=5
    )
    draft = model.draft_tokens([1, 2], num_tokens=5)
    assert len(draft) == 5

def test_distribution_preservation():
    """Verify that speculative decoding approximately preserves
    the target distribution. Run 1000 samples and check
    that token frequencies are within a tolerance."""
    # This is a statistical test — run with sufficient samples
    # and compare against baseline GPT-2 generation
    pass  # Implement as part of your benchmarking suite

if __name__ == "__main__":
    pytest.main([__file__, "-v"])

python -m pytest test_engine.py -v

# benchmark.py
import time
from verify_model import VerificationTransformer

def benchmark_baseline(prompt: str, num_tokens: int = 50):
    model = VerificationTransformer("gpt2-medium")
    start = time.time()
    output = model.generate_single(
        model.tokenizer.encode(prompt), num_tokens=num_tokens
    )
    elapsed = time.time() - start
    return elapsed, len(output)

def benchmark_speculative(prompt: str, num_tokens: int = 50):
    # ... instantiate engine and run
    # compare wall-clock time and tokens/second
    pass

# Compare tokens/second
baseline_time, baseline_tokens = benchmark_baseline(
    "The key to scalable inference is", 50
)
print(f"Baseline: {baseline_tokens / baseline_time:.1f} tokens/sec")
# Speculative should show 1.2x–3x speedup depending on draft quality