speculative_decoding.py

"""
Speculative Decoding with Real Models — Qwen2-0.5B (draft) + Qwen2-1.5B (target).

Measures:
1. Token acceptance rate (alpha) per position
2. Real wall-clock speedup vs autoregressive baseline
3. Analytical speedup model validation
4. Regime map: when speculative decoding helps vs hurts
"""

import gc
import json
import os
import time
import warnings
from datetime import datetime, timezone

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import torch
import torch.nn.functional as F
from transformers import AutoModelForCausalLM, AutoTokenizer

warnings.filterwarnings("ignore")

DEVICE      = "cuda" if torch.cuda.is_available() else "cpu"
DTYPE       = torch.float16
RESULTS_DIR = "results"
PLOTS_DIR   = "plots"
os.makedirs(RESULTS_DIR, exist_ok=True)
os.makedirs(PLOTS_DIR,   exist_ok=True)

DRAFT_MODEL  = "Qwen/Qwen2-0.5B"
TARGET_MODEL = "Qwen/Qwen2-1.5B"

# Speculative decoding parameters
GAMMA_VALUES   = [1, 2, 4, 8]   # draft tokens per step
MAX_NEW_TOKENS = 64
N_REPEATS      = 3               # runs per config (take median)
TEMPERATURE    = 1.0             # 1.0 = greedy-equivalent for acceptance

PROMPTS = {
    "code_python": (
        "def binary_search(arr, target):\n"
        "    left, right = 0, len(arr) - 1\n"
        "    while left <= right:"
    ),
    "code_cpp": (
        "#include <vector>\n"
        "#include <algorithm>\n"
        "int findMax(std::vector<int>& nums) {"
    ),
    "technical_ai": (
        "Transformer attention mechanisms compute queries, keys, and values "
        "through learned linear projections. The scaled dot-product attention "
        "formula is defined as"
    ),
    "technical_systems": (
        "In distributed LLM inference, tensor parallelism shards weight matrices "
        "across multiple GPUs. Each device computes a partial result and "
        "synchronizes via AllReduce. The communication overhead is"
    ),
    "creative_story": (
        "Once upon a time, in a land far away, there lived a wise old wizard "
        "who had spent decades studying the ancient arts of magic. One morning,"
    ),
    "factual_science": (
        "The speed of light in vacuum is approximately 299,792,458 meters per "
        "second. This fundamental constant, denoted as c, plays a central role "
        "in Einstein's theory of special relativity because"
    ),
    "conversational": (
        "Hello! I was wondering if you could help me understand how neural "
        "networks learn. I've heard about backpropagation but I'm not sure"
    ),
    "mathematical": (
        "The Fibonacci sequence is defined recursively as F(n) = F(n-1) + F(n-2) "
        "with base cases F(0) = 0 and F(1) = 1. The closed-form solution, "
        "known as Binet's formula, is"
    ),
}

plt.rcParams.update({
    "figure.dpi":        150,
    "axes.spines.top":   False,
    "axes.spines.right": False,
    "axes.grid":         True,
    "grid.alpha":        0.3,
    "font.size":         10,
})

DIVIDER = "=" * 70

def section(title):
    print(f"\n{DIVIDER}\n  {title}\n{DIVIDER}")


# ─────────────────────────────────────────────────────────────────────────────
# Model loading
# ─────────────────────────────────────────────────────────────────────────────

def load_models():
    print(f"  Loading draft:  {DRAFT_MODEL}")
    tokenizer = AutoTokenizer.from_pretrained(DRAFT_MODEL, trust_remote_code=True)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    draft = AutoModelForCausalLM.from_pretrained(
        DRAFT_MODEL, torch_dtype=DTYPE, trust_remote_code=True,
    ).eval().to(DEVICE)

    print(f"  Loading target: {TARGET_MODEL}")
    target = AutoModelForCausalLM.from_pretrained(
        TARGET_MODEL, torch_dtype=DTYPE, trust_remote_code=True,
    ).eval().to(DEVICE)

    draft_params  = sum(p.numel() for p in draft.parameters())  / 1e6
    target_params = sum(p.numel() for p in target.parameters()) / 1e6
    print(f"  Draft:  {draft_params:.0f}M params")
    print(f"  Target: {target_params:.0f}M params")
    print(f"  Ratio:  1:{target_params/draft_params:.1f}")

    return tokenizer, draft, target


# ─────────────────────────────────────────────────────────────────────────────
# Core: autoregressive baseline
# ─────────────────────────────────────────────────────────────────────────────

def autoregressive_generate(model, input_ids, max_new_tokens):
    """Standard autoregressive generation. Returns tokens and wall time."""
    generated = input_ids.clone()

    with torch.inference_mode():
        # Warmup KV cache on prompt
        out = model(input_ids=generated, use_cache=True, return_dict=True)
        pkv = out.past_key_values
        cur = out.logits[:, -1, :].argmax(dim=-1, keepdim=True)
        generated = torch.cat([generated, cur], dim=1)

        torch.cuda.synchronize()
        t0 = time.perf_counter()

        for _ in range(max_new_tokens - 1):
            out = model(input_ids=cur, past_key_values=pkv,
                        use_cache=True, return_dict=True)
            pkv = out.past_key_values
            cur = out.logits[:, -1, :].argmax(dim=-1, keepdim=True)
            generated = torch.cat([generated, cur], dim=1)

        torch.cuda.synchronize()
        elapsed = time.perf_counter() - t0

    return generated, elapsed


# ─────────────────────────────────────────────────────────────────────────────
# Core: speculative decoding
# ─────────────────────────────────────────────────────────────────────────────

def speculative_generate(draft_model, target_model, input_ids,
                          max_new_tokens, gamma):
    """
    Speculative decoding with greedy acceptance.

    Algorithm:
    1. Draft model generates gamma tokens autoregressively
    2. Target model scores all gamma+1 positions in one forward pass
    3. Accept draft tokens where target agrees (greedy: argmax matches)
    4. Always accept at least one token (the target's correction)
    5. Repeat until max_new_tokens reached

    Returns:
        generated tokens, elapsed time, acceptance rates per step
    """
    generated       = input_ids.clone()
    prompt_len      = input_ids.shape[1]
    tokens_generated = 0
    acceptance_rates = []

    # Prime target KV cache on prompt
    with torch.inference_mode():
        target_out = target_model(
            input_ids=generated, use_cache=True, return_dict=True
        )
        target_pkv = target_out.past_key_values

        # Prime draft KV cache on prompt
        draft_out = draft_model(
            input_ids=generated, use_cache=True, return_dict=True
        )
        draft_pkv = draft_out.past_key_values

    torch.cuda.synchronize()
    t0 = time.perf_counter()

    with torch.inference_mode():
        while tokens_generated < max_new_tokens:
            actual_gamma = min(gamma, max_new_tokens - tokens_generated)

            # ── Step 1: Draft generates gamma tokens ──────────────────────
            draft_tokens   = []
            draft_logits_list = []
            cur_draft_pkv  = draft_pkv
            cur_token      = generated[:, -1:]

            for _ in range(actual_gamma):
                d_out = draft_model(
                    input_ids=cur_token,
                    past_key_values=cur_draft_pkv,
                    use_cache=True,
                    return_dict=True,
                )
                cur_draft_pkv = d_out.past_key_values
                d_logits      = d_out.logits[:, -1, :]
                d_token       = d_logits.argmax(dim=-1, keepdim=True)
                draft_tokens.append(d_token)
                draft_logits_list.append(d_logits)
                cur_token = d_token

            draft_seq = torch.cat(draft_tokens, dim=1)  # [1, gamma]

            # ── Step 2: Target scores draft tokens in ONE forward pass ────
            target_input = torch.cat([generated[:, -1:], draft_seq], dim=1)
            t_out = target_model(
                input_ids=target_input,
                past_key_values=target_pkv,
                use_cache=True,
                return_dict=True,
            )
            # target_logits: [1, gamma+1, vocab]
            target_logits = t_out.logits

            # ── Step 3: Greedy acceptance ─────────────────────────────────
            accepted    = 0
            new_tokens  = []

            for i in range(actual_gamma):
                target_token = target_logits[:, i, :].argmax(dim=-1, keepdim=True)
                if target_token.item() == draft_tokens[i].item():
                    new_tokens.append(draft_tokens[i])
                    accepted += 1
                else:
                    # Reject: use target's correction and stop
                    new_tokens.append(target_token)
                    break
            else:
                # All accepted: append target's next token too
                bonus = target_logits[:, actual_gamma, :].argmax(dim=-1, keepdim=True)
                new_tokens.append(bonus)

            acceptance_rates.append(accepted / actual_gamma)

            # ── Step 4: Update state ──────────────────────────────────────
            new_seq    = torch.cat(new_tokens, dim=1)
            generated  = torch.cat([generated, new_seq], dim=1)
            n_accepted = new_seq.shape[1]
            tokens_generated += n_accepted

            # Re-sync KV caches to the accepted prefix
            # Draft: recompute from the accepted prefix
            accepted_input = new_seq
            d_out = draft_model(
                input_ids=accepted_input,
                past_key_values=cur_draft_pkv,
                use_cache=True,
                return_dict=True,
            )
            draft_pkv = d_out.past_key_values

            # Target: its KV cache already contains up to bonus token
            # We need to rewind if not all draft tokens were accepted
            if accepted < actual_gamma:
                # Recompute target KV on the accepted prefix
                target_pkv = t_out.past_key_values
                # Fast-forward target to accepted position
                if n_accepted > 1:
                    t_out2 = target_model(
                        input_ids=new_seq[:, 1:],
                        past_key_values=target_pkv,
                        use_cache=True,
                        return_dict=True,
                    )
                    target_pkv = t_out2.past_key_values
                else:
                    target_pkv = t_out.past_key_values
            else:
                target_pkv = t_out.past_key_values

    torch.cuda.synchronize()
    elapsed = time.perf_counter() - t0

    return generated, elapsed, acceptance_rates


# ─────────────────────────────────────────────────────────────────────────────
# Analytical speedup model
# ─────────────────────────────────────────────────────────────────────────────

def analytical_speedup(alpha, gamma, c):
    """
    Expected speedup for speculative decoding.

    Formula from Leviathan et al. (2023):
        E[speedup] = (1 - alpha^(gamma+1)) / ((1 - alpha) * (1 + gamma/c))

    Args:
        alpha: per-token acceptance rate (0 to 1)
        gamma: draft tokens per speculation step
        c:     cost ratio = time(target) / time(draft)
               i.e. how many draft steps fit in one target step
    """
    numerator   = 1 - alpha ** (gamma + 1)
    denominator = (1 - alpha) * (1 + gamma / c)
    return numerator / denominator if denominator > 0 else 1.0


# ─────────────────────────────────────────────────────────────────────────────
# Benchmark runner
# ─────────────────────────────────────────────────────────────────────────────

def run_benchmark(tokenizer, draft, target):
    section("RUNNING BENCHMARK")

    rows = []

    for prompt_name, prompt_text in PROMPTS.items():
        print(f"\n  Prompt: {prompt_name}")

        input_ids = tokenizer(
            prompt_text, return_tensors="pt", add_special_tokens=True
        ).input_ids.to(DEVICE)

        prompt_len = input_ids.shape[1]

        # ── Baseline: target autoregressive ──────────────────────────────
        ar_times = []
        for _ in range(N_REPEATS):
            gc.collect(); torch.cuda.empty_cache()
            _, t = autoregressive_generate(target, input_ids, MAX_NEW_TOKENS)
            ar_times.append(t)
        ar_time = float(np.median(ar_times))
        ar_tok_s = MAX_NEW_TOKENS / ar_time

        print(f"    Baseline AR: {ar_time*1000:.1f}ms  "
              f"({ar_tok_s:.1f} tok/s)")

        # ── Draft only: measure draft speed ──────────────────────────────
        draft_times = []
        for _ in range(N_REPEATS):
            gc.collect(); torch.cuda.empty_cache()
            _, t = autoregressive_generate(draft, input_ids, MAX_NEW_TOKENS)
            draft_times.append(t)
        draft_time = float(np.median(draft_times))
        draft_tok_s = MAX_NEW_TOKENS / draft_time

        # Cost ratio: target / draft per token
        cost_ratio = ar_time / draft_time

        print(f"    Draft AR:    {draft_time*1000:.1f}ms  "
              f"({draft_tok_s:.1f} tok/s)  "
              f"cost_ratio={cost_ratio:.2f}")

        # ── Speculative decoding for each gamma ───────────────────────────
        for gamma in GAMMA_VALUES:
            spec_times = []
            all_acceptance = []

            for _ in range(N_REPEATS):
                gc.collect(); torch.cuda.empty_cache()
                _, t, acc = speculative_generate(
                    draft, target, input_ids, MAX_NEW_TOKENS, gamma
                )
                spec_times.append(t)
                all_acceptance.extend(acc)

            spec_time = float(np.median(spec_times))
            alpha     = float(np.mean(all_acceptance))
            speedup   = ar_time / spec_time
            spec_tok_s = MAX_NEW_TOKENS / spec_time

            # Analytical prediction
            predicted_speedup = analytical_speedup(alpha, gamma, cost_ratio)

            print(f"    gamma={gamma}: "
                  f"spec={spec_time*1000:.1f}ms  "
                  f"speedup={speedup:.3f}x  "
                  f"alpha={alpha:.3f}  "
                  f"predicted={predicted_speedup:.3f}x")

            rows.append({
                "prompt":            prompt_name,
                "prompt_len":        prompt_len,
                "gamma":             gamma,
                "ar_time_s":         ar_time,
                "draft_time_s":      draft_time,
                "spec_time_s":       spec_time,
                "ar_tok_s":          ar_tok_s,
                "draft_tok_s":       draft_tok_s,
                "spec_tok_s":        spec_tok_s,
                "speedup":           speedup,
                "alpha":             alpha,
                "cost_ratio":        cost_ratio,
                "predicted_speedup": predicted_speedup,
                "model_error":       abs(speedup - predicted_speedup),
            })

    return pd.DataFrame(rows)


# ─────────────────────────────────────────────────────────────────────────────
# Analysis
# ─────────────────────────────────────────────────────────────────────────────

def analyze(df):
    section("RESULTS ANALYSIS")

    # ── 1. Speedup summary ────────────────────────────────────────────────
    print(f"\n  {'Prompt':24} {'gamma':>6} {'alpha':>7} "
          f"{'Speedup':>9} {'Predicted':>10} {'Error':>7}")
    print(f"  {'-'*24} {'-'*6} {'-'*7} {'-'*9} {'-'*10} {'-'*7}")

    for _, r in df.sort_values(["prompt", "gamma"]).iterrows():
        marker = "►" if r["speedup"] > 1.0 else " "
        print(
            f"  {marker}{r['prompt']:23} {int(r['gamma']):>6} "
            f"{r['alpha']:>7.3f} {r['speedup']:>9.3f}x "
            f"{r['predicted_speedup']:>10.3f}x "
            f"{r['model_error']:>7.3f}"
        )

    # ── 2. Best gamma per prompt ──────────────────────────────────────────
    section("BEST GAMMA PER PROMPT")
    best = df.loc[df.groupby("prompt")["speedup"].idxmax()]
    print(f"\n  {'Prompt':24} {'Best gamma':>10} {'Speedup':>9} "
          f"{'alpha':>7} {'Helpful?':>9}")
    print(f"  {'-'*24} {'-'*10} {'-'*9} {'-'*7} {'-'*9}")

    for _, r in best.sort_values("speedup", ascending=False).iterrows():
        helpful = "YES ✓" if r["speedup"] > 1.0 else "NO  ✗"
        print(
            f"  {r['prompt']:24} {int(r['gamma']):>10} "
            f"{r['speedup']:>9.3f}x {r['alpha']:>7.3f} {helpful:>9}"
        )

    # ── 3. Regime analysis ────────────────────────────────────────────────
    section("REGIME ANALYSIS")

    helpful   = df[df["speedup"] > 1.0]
    harmful   = df[df["speedup"] <= 1.0]
    n_total   = len(df)
    n_helpful = len(helpful)

    print(f"""
  Total configurations:  {n_total}
  Helpful (speedup > 1): {n_helpful} ({100*n_helpful/n_total:.1f}%)
  Harmful (speedup <= 1): {len(harmful)} ({100*len(harmful)/n_total:.1f}%)

  Mean speedup (all):     {df['speedup'].mean():.3f}x
  Mean speedup (helpful): {helpful['speedup'].mean():.3f}x  (when it works)
  Max speedup:            {df['speedup'].max():.3f}x  ({df.loc[df['speedup'].idxmax(), 'prompt']})
  Min speedup:            {df['speedup'].min():.3f}x  ({df.loc[df['speedup'].idxmin(), 'prompt']})

  Mean alpha (helpful):   {helpful['alpha'].mean():.3f}
  Mean alpha (harmful):   {harmful['alpha'].mean():.3f}

  Analytical model:
    Mean absolute error:  {df['model_error'].mean():.4f}
    Max absolute error:   {df['model_error'].max():.4f}
    Correlation (real vs predicted): {df['speedup'].corr(df['predicted_speedup']):.4f}
    """)

    # ── 4. Alpha threshold ────────────────────────────────────────────────
    print("  Alpha threshold analysis:")
    for threshold in [0.5, 0.6, 0.7, 0.8, 0.9]:
        above = df[df["alpha"] >= threshold]
        if len(above) > 0:
            pct_helpful = (above["speedup"] > 1.0).mean() * 100
            mean_sp     = above["speedup"].mean()
            print(f"    alpha >= {threshold:.1f}: "
                  f"{len(above):>3} configs, "
                  f"{pct_helpful:.0f}% helpful, "
                  f"mean speedup={mean_sp:.3f}x")

    return helpful, harmful


# ─────────────────────────────────────────────────────────────────────────────
# Plots
# ─────────────────────────────────────────────────────────────────────────────

def generate_plots(df):
    section("GENERATING PLOTS")

    PROMPT_COLORS = {
        "code_python":       "#E74C3C",
        "code_cpp":          "#C0392B",
        "technical_ai":      "#3498DB",
        "technical_systems": "#2980B9",
        "creative_story":    "#2ECC71",
        "factual_science":   "#27AE60",
        "conversational":    "#9B59B6",
        "mathematical":      "#8E44AD",
    }

    # ── Plot 1: Real vs predicted speedup ─────────────────────────────────
    fig, ax = plt.subplots(figsize=(9, 7))

    for prompt in df["prompt"].unique():
        sub = df[df["prompt"] == prompt]
        ax.scatter(
            sub["predicted_speedup"], sub["speedup"],
            s=80, color=PROMPT_COLORS.get(prompt, "#888"),
            label=prompt, zorder=5, alpha=0.85,
        )

    mn = min(df["predicted_speedup"].min(), df["speedup"].min()) * 0.95
    mx = max(df["predicted_speedup"].max(), df["speedup"].max()) * 1.05
    ax.plot([mn, mx], [mn, mx], "--", color="gray",
            linewidth=1.5, label="Perfect prediction", alpha=0.6)
    ax.axhline(1.0, color="red", linewidth=1, linestyle=":", alpha=0.5,
               label="Speedup = 1 (break-even)")
    ax.axvline(1.0, color="red", linewidth=1, linestyle=":", alpha=0.5)

    corr = df["speedup"].corr(df["predicted_speedup"])
    mae  = df["model_error"].mean()
    ax.text(0.05, 0.92,
            f"Correlation: {corr:.4f}\nMAE: {mae:.4f}",
            transform=ax.transAxes, fontsize=10,
            bbox=dict(boxstyle="round", facecolor="wheat", alpha=0.6))

    ax.set_xlabel("Predicted speedup (analytical model)")
    ax.set_ylabel("Real speedup (measured)")
    ax.set_title("Speculative Decoding: Real vs Predicted Speedup\n"
                 f"Qwen2-0.5B (draft) → Qwen2-1.5B (target)")
    ax.legend(fontsize=7, ncol=2, loc="lower right")
    plt.tight_layout()
    plt.savefig(os.path.join(PLOTS_DIR, "real_vs_predicted.png"),
                dpi=180, bbox_inches="tight")
    plt.close()
    print("  plot: plots/real_vs_predicted.png")

    # ── Plot 2: Speedup vs alpha, colored by gamma ────────────────────────
    fig, ax = plt.subplots(figsize=(10, 6))

    GAMMA_COLORS = {1: "#E74C3C", 2: "#E67E22", 4: "#3498DB", 8: "#2ECC71"}

    for gamma in GAMMA_VALUES:
        sub = df[df["gamma"] == gamma]
        ax.scatter(sub["alpha"], sub["speedup"],
                   s=90, color=GAMMA_COLORS[gamma],
                   label=f"gamma={gamma}", zorder=5, alpha=0.85)

    ax.axhline(1.0, color="gray", linestyle="--",
               linewidth=1.5, alpha=0.6, label="Break-even")

    # Overlay analytical curves
    alpha_range = np.linspace(0.3, 0.99, 200)
    cost_ratio  = df["cost_ratio"].mean()
    for gamma, col in GAMMA_COLORS.items():
        curve = [analytical_speedup(a, gamma, cost_ratio) for a in alpha_range]
        ax.plot(alpha_range, curve, "-", color=col,
                linewidth=1.5, alpha=0.4)

    ax.set_xlabel("Token acceptance rate (alpha)")
    ax.set_ylabel("Speedup over autoregressive baseline")
    ax.set_title("Speedup vs Acceptance Rate\n"
                 f"Dots=measured, Lines=analytical model "
                 f"(cost_ratio={cost_ratio:.2f})")
    ax.legend(fontsize=9)
    plt.tight_layout()
    plt.savefig(os.path.join(PLOTS_DIR, "speedup_vs_alpha.png"),
                dpi=180, bbox_inches="tight")
    plt.close()
    print("  plot: plots/speedup_vs_alpha.png")

    # ── Plot 3: Speedup heatmap (prompt × gamma) ──────────────────────────
    pivot = df.pivot_table(
        index="prompt", columns="gamma",
        values="speedup", aggfunc="mean"
    )

    fig, ax = plt.subplots(figsize=(10, 6))
    im = ax.imshow(pivot.values, aspect="auto",
                   cmap="RdYlGn", vmin=0.5, vmax=2.0)
    ax.set_xticks(range(len(pivot.columns)))
    ax.set_xticklabels([f"gamma={g}" for g in pivot.columns])
    ax.set_yticks(range(len(pivot.index)))
    ax.set_yticklabels(pivot.index, fontsize=9)
    plt.colorbar(im, ax=ax, label="Speedup (>1 = helpful)")

    for i in range(len(pivot.index)):
        for j in range(len(pivot.columns)):
            v   = pivot.values[i, j]
            col = "white" if v < 0.9 or v > 1.6 else "black"
            ax.text(j, i, f"{v:.2f}x",
                    ha="center", va="center", fontsize=9, color=col)

    ax.set_title("Speedup Heatmap — Prompt × Gamma\n"
                 "Green > 1.0 (helpful) · Red < 1.0 (harmful)")
    plt.tight_layout()
    plt.savefig(os.path.join(PLOTS_DIR, "speedup_heatmap.png"),
                dpi=180, bbox_inches="tight")
    plt.close()
    print("  plot: plots/speedup_heatmap.png")

    # ── Plot 4: Alpha heatmap (prompt × gamma) ────────────────────────────
    pivot_alpha = df.pivot_table(
        index="prompt", columns="gamma",
        values="alpha", aggfunc="mean"
    )

    fig, ax = plt.subplots(figsize=(10, 6))
    im = ax.imshow(pivot_alpha.values, aspect="auto",
                   cmap="Blues", vmin=0, vmax=1)
    ax.set_xticks(range(len(pivot_alpha.columns)))
    ax.set_xticklabels([f"gamma={g}" for g in pivot_alpha.columns])
    ax.set_yticks(range(len(pivot_alpha.index)))
    ax.set_yticklabels(pivot_alpha.index, fontsize=9)
    plt.colorbar(im, ax=ax, label="Token acceptance rate (alpha)")

    for i in range(len(pivot_alpha.index)):
        for j in range(len(pivot_alpha.columns)):
            v   = pivot_alpha.values[i, j]
            col = "white" if v > 0.7 else "black"
            ax.text(j, i, f"{v:.3f}",
                    ha="center", va="center", fontsize=9, color=col)

    ax.set_title("Token Acceptance Rate (alpha) — Prompt × Gamma\n"
                 "Higher = draft model agrees more with target")
    plt.tight_layout()
    plt.savefig(os.path.join(PLOTS_DIR, "alpha_heatmap.png"),
                dpi=180, bbox_inches="tight")
    plt.close()
    print("  plot: plots/alpha_heatmap.png")

    # ── Plot 5: Analytical model curves ───────────────────────────────────
    fig, axes = plt.subplots(1, 2, figsize=(16, 6))

    cost_ratio = df["cost_ratio"].mean()
    alpha_rng  = np.linspace(0.01, 0.999, 300)

    # Left: speedup vs alpha for different gamma
    ax = axes[0]
    for gamma, col in GAMMA_COLORS.items():
        curve = [analytical_speedup(a, gamma, cost_ratio) for a in alpha_rng]
        ax.plot(alpha_rng, curve, linewidth=2, color=col,
                label=f"gamma={gamma}")

    ax.axhline(1.0, color="gray", linestyle="--", alpha=0.5)
    ax.set_xlabel("Token acceptance rate (alpha)")
    ax.set_ylabel("Expected speedup")
    ax.set_title(f"Analytical Speedup Model\n"
                 f"cost_ratio={cost_ratio:.2f} "
                 f"(Qwen2-0.5B → Qwen2-1.5B)")
    ax.legend()
    ax.set_ylim(0, 4)

    # Right: speedup vs gamma for different alpha values
    ax = axes[1]
    gamma_rng = np.linspace(1, 16, 100)
    for alpha_val, col in [(0.5, "#E74C3C"), (0.7, "#E67E22"),
                            (0.85, "#3498DB"), (0.95, "#2ECC71")]:
        curve = [analytical_speedup(alpha_val, g, cost_ratio)
                 for g in gamma_rng]
        ax.plot(gamma_rng, curve, linewidth=2, color=col,
                label=f"alpha={alpha_val}")

    ax.axhline(1.0, color="gray", linestyle="--", alpha=0.5)
    ax.set_xlabel("Draft tokens per step (gamma)")
    ax.set_ylabel("Expected speedup")
    ax.set_title("Analytical Speedup vs Gamma\n"
                 "Diminishing returns at high gamma")
    ax.legend()
    ax.set_ylim(0, 4)

    plt.suptitle("Leviathan et al. (2023) Analytical Speedup Model",
                 fontsize=12, fontweight="bold")
    plt.tight_layout()
    plt.savefig(os.path.join(PLOTS_DIR, "analytical_model.png"),
                dpi=180, bbox_inches="tight")
    plt.close()
    print("  plot: plots/analytical_model.png")

    # ── Plot 6: Regime map ────────────────────────────────────────────────
    fig, ax = plt.subplots(figsize=(10, 6))

    helpful = df[df["speedup"] > 1.0]
    harmful = df[df["speedup"] <= 1.0]

    ax.scatter(harmful["alpha"], harmful["cost_ratio"],
               s=100, color="#E74C3C", marker="x",
               linewidths=2, label="Harmful (speedup ≤ 1.0)", zorder=5)
    ax.scatter(helpful["alpha"], helpful["cost_ratio"],
               s=100, color="#2ECC71", marker="o",
               label="Helpful (speedup > 1.0)", zorder=5, alpha=0.8)

    ax.set_xlabel("Token acceptance rate (alpha)")
    ax.set_ylabel("Cost ratio (target / draft time)")
    ax.set_title("Regime Map: When Does Speculative Decoding Help?\n"
                 "Green=helpful · Red=harmful")
    ax.legend()
    plt.tight_layout()
    plt.savefig(os.path.join(PLOTS_DIR, "regime_map.png"),
                dpi=180, bbox_inches="tight")
    plt.close()
    print("  plot: plots/regime_map.png")


# ─────────────────────────────────────────────────────────────────────────────
# Main
# ─────────────────────────────────────────────────────────────────────────────

def main():
    metadata = {
        "timestamp":    datetime.now(timezone.utc).isoformat(),
        "device":       DEVICE,
        "gpu":          torch.cuda.get_device_name(0) if torch.cuda.is_available() else None,
        "draft_model":  DRAFT_MODEL,
        "target_model": TARGET_MODEL,
        "gamma_values": GAMMA_VALUES,
        "max_new_tokens": MAX_NEW_TOKENS,
        "n_repeats":    N_REPEATS,
        "n_prompts":    len(PROMPTS),
    }
    print(json.dumps(metadata, indent=2))

    section("LOADING MODELS")
    tokenizer, draft, target = load_models()

    df = run_benchmark(tokenizer, draft, target)
    df.to_csv(os.path.join(RESULTS_DIR, "speculative_results.csv"), index=False)

    helpful, harmful = analyze(df)
    generate_plots(df)

    print(f"\n{DIVIDER}")
    print("  speculative_decoding.py complete.")
    print(f"  Saved: results/speculative_results.csv")
    print(DIVIDER)


if __name__ == "__main__":
    main()