pythia-12b vs Greedy on A100

# Speculative Decoding vs. Greedy Decoding — pythia-12b + pythia-410m on one A100 **Paper:** Leviathan, Kalman & Matias, *Fast Inference from Transformers via Speculative Decoding* (ICML 2023) — [arXiv:2211.17192](https://arxiv.org/abs/2211.17192) · [PMLR v202](https://proceedings.mlr.press/v202/leviathan23a.html) ## The idea Standard autoregressive decoding is **serial**: K tokens cost K forward passes of the big *target* model, and at batch size 1 each pass is memory-bandwidth-bound — the model's weights are re-read from HBM just to emit a single token. Speculative decoding breaks the serial bottleneck (Algorithm 1 in the paper): 1. **Draft** — a small, fast approximation model (here `pythia-410m`) autoregressively proposes γ candidate tokens. 2. **Verify** — the target model (`pythia-12b`) scores the whole candidate block in **one parallel forward pass**. 3. **Accept** — take the longest prefix that matches the target model's own greedy choice, plus one "bonus" token from the target, then repeat. Because accepted tokens are exactly what the target would have picked, **greedy speculative decoding returns exactly the same token IDs as plain greedy decoding** — no retraining, no architecture change, no approximation. The paper reports 2–3× wall-clock speedups on T5-XXL with identical outputs. **Expected speedup** (Theorem 3.8): `(1 − α^(γ+1)) / ((1 − α)(γc + 1))`, where `α` = draft-token acceptance rate and `c` = draft/target per-step cost ratio. ## This notebook - *…

import gc
import time

import torch
import pandas as pd
import matplotlib.pyplot as plt
from transformers import AutoModelForCausalLM, AutoTokenizer

# Exactness requirement: the verify pass (1,gamma+1 tokens) and greedy decode (1,1)
# use different matmul shapes, so low-precision rounding flips near-tie argmaxes
# (measured bf16 logit noise ~1.75 with top-2 gaps as small as 0.125 -> IDs diverge).
# fp32 with TF32 OFF keeps the two paths argmax-identical.
torch.backends.cuda.matmul.allow_tf32 = False
torch.backends.cudnn.allow_tf32 = False

assert torch.cuda.is_available(), "This notebook needs a CUDA GPU"
device = torch.device("cuda")
print(f"GPU: {torch.cuda.get_device_name(0)} "
…

GPU: NVIDIA A100 80GB PCIe (79 GiB)
models loaded in 7.4s
target: 11.85B params (fp32) | draft: 0.41B params (bf16) | VRAM used: 68.5 GiB

from transformers import DynamicCache

PROMPTS = [
    "The key difference between supervised and unsupervised learning is",
    "In a surprising discovery announced on Tuesday, scientists reported that",
    "The recipe for a perfect chocolate chip cookie starts with",
    "def binary_search(arr, target):",
    "The economic history of the late Roman Empire shows that",
]
MAX_NEW_TOKENS = 256
GAMMA = 5  # draft proposes 5 tokens per iteration (same as HF's default num_assistant_tokens)
EOS_ID = tokenizer.eos_token_id


@torch.inference_mode()
def timed_generate(prompt, max_new=MAX_NEW_TOKENS):
    """Baseline: plain greedy decoding with the 12B target (HF generate).
    Returns (new_token_ids, seconds)."""
…

warm-up done

# --- Baseline: plain greedy decoding with the 12B target model only ---
baseline = []
for p in PROMPTS:
    ids, dt = timed_generate(p)
    baseline.append(dict(prompt=p, ids=ids, time_s=dt, tok_s=len(ids) / dt))
    print(f"{dt:7.2f}s  {len(ids) / dt:6.2f} tok/s  | {p[:70]}")

8.98s   28.50 tok/s  | The key difference between supervised and unsupervised learning is
   8.99s   28.49 tok/s  | In a surprising discovery announced on Tuesday, scientists reported th
   8.99s   28.47 tok/s  | The recipe for a perfect chocolate chip cookie starts with
   8.98s   28.49 tok/s  | def binary_search(arr, target):
   8.98s   28.50 tok/s  | The economic history of the late Roman Empire shows that

# --- Speculative decoding: Algorithm 1 from the paper (12B target + 410M draft) ---
speculative = []
for p in PROMPTS:
    ids, dt, st = speculative_generate(p)
    speculative.append(dict(prompt=p, ids=ids, time_s=dt, tok_s=len(ids) / dt, **st))
    print(f"{dt:7.2f}s  {len(ids) / dt:6.2f} tok/s  accept={st['accept_rate']:.2f}  "
          f"({st['tok_per_round']:.1f} tok/round)  | {p[:52]}")

# Cost ratio c from the paper: draft step time / target step time
_din = tokenizer(PROMPTS[0], return_tensors="pt").to(device)
torch.cuda.synchronize(); _t0 = time.perf_counter()
with torch.inference_mode():
    draft.generate(**_din, do_sample=False, max_new_tokens=MAX_NEW_TOKENS,
                   min_new_tokens=MAX_NEW_TOKENS, pad_token_id=EOS_ID)
torch.cuda.synchronize()
draft_step_s = (time.perf_counter() - _t0) / MAX_NEW_TOKENS
target_step_s = sum(r["time_s"] for r in baseline) / (len(baseline) * MAX_NEW_TOKENS)
C_RATIO = draft_step_s / target_step_s
…

9.07s   28.23 tok/s  accept=0.47  (3.3 tok/round)  | The key difference between supervised and unsupervis
   7.52s   34.05 tok/s  accept=0.61  (4.0 tok/round)  | In a surprising discovery announced on Tuesday, scie
   6.08s   42.14 tok/s  accept=0.80  (4.9 tok/round)  | The recipe for a perfect chocolate chip cookie start
   6.57s   38.95 tok/s  accept=0.73  (4.6 tok/round)  | def binary_search(arr, target):
   6.46s   39.65 tok/s  accept=0.74  (4.7 tok/round)  | The economic history of the late…

# --- Verification + results table ---
rows = []
for i, (b, s) in enumerate(zip(baseline, speculative), 1):
    identical = b["ids"] == s["ids"]
    div = next((j for j, (x, y) in enumerate(zip(b["ids"], s["ids"])) if x != y), None)
    rows.append({
        "#": i,
        "prompt": b["prompt"][:44] + "...",
        "new tokens": len(b["ids"]),
        "greedy s": round(b["time_s"], 2),
        "spec s": round(s["time_s"], 2),
        "greedy tok/s": round(b["tok_s"], 1),
        "spec tok/s": round(s["tok_s"], 1),
        "speedup": f"{b['time_s'] / s['time_s']:.2f}x",
        "accept rate": round(s["accept_rate"], 2),
        "IDs identical": "YES" if identical else f"NO (diverge@{div})",
    })
df = pd.DataFrame(rows)
…

#                                          prompt  new tokens  greedy s  spec s  greedy tok/s  spec tok/s speedup  accept rate IDs identical
 1 The key difference between supervised and un...         256      8.98    9.07          28.5        28.2   0.99x         0.47           YES
 2 In a surprising discovery announced on Tuesd...         256      8.99    7.52          28.5        34.0   1.20x         0.61           YES
 3 The recipe for a perfect chocolate chip cook...         256      8.99…

# --- Visualization: throughput, acceptance dynamics, and a traced generation ---
import numpy as np

labels = [f"P{i}" for i in range(1, len(PROMPTS) + 1)]
b_tok = np.array([r["tok_s"] for r in baseline])
s_tok = np.array([r["tok_s"] for r in speculative])
acc = np.array([r["accept_rate"] for r in speculative])
tpr = np.array([r["tok_per_round"] for r in speculative])
spd = s_tok / b_tok

# one traced run (per-round timings) on the median-speedup prompt
i_med = int(np.argsort(spd)[len(PROMPTS) // 2])
ids_tr, dt_tr, st_tr = speculative_generate(PROMPTS[i_med], trace=True)
assert ids_tr == baseline[i_med]["ids"], "traced run must also match greedy exactly"
det = [d for d in st_tr["detail"] if "t_cum" in d]

C_G, C_S, C_A = "#7f9cbf", "#2ca25f", "#e08214"
fig, axes = plt.subplots(1, 3, figsize=(16.5, 4.8))
…