phase3_profiling.py

"""
Phase 3 — Granular profiling of speculative decoding overhead.

Measures time breakdown per speculative step:
  - Draft forward passes (gamma tokens)
  - Target forward pass (1 pass for gamma+1 positions)
  - KV cache synchronization overhead
  - Token comparison overhead

Also measures:
  - Alpha vs position (does acceptance rate drop later in sequence?)
  - Alpha vs temperature (does sampling help?)
  - Overhead factor vs gamma (confirms Phase 2 model)
"""

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
from transformers import AutoModelForCausalLM, AutoTokenizer

warnings.filterwarnings("ignore")
np.random.seed(42)

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"

DIVIDER = "=" * 72

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

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


# ─────────────────────────────────────────────────────────────────────────────
# Load models
# ─────────────────────────────────────────────────────────────────────────────

def load_models():
    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)

    target = AutoModelForCausalLM.from_pretrained(
        TARGET_MODEL, torch_dtype=DTYPE, trust_remote_code=True,
    ).eval().to(DEVICE)

    return tokenizer, draft, target


def sync():
    if DEVICE == "cuda":
        torch.cuda.synchronize()


def timed(fn):
    sync()
    t0 = time.perf_counter()
    result = fn()
    sync()
    return result, (time.perf_counter() - t0) * 1000.0


# ─────────────────────────────────────────────────────────────────────────────
# Section 1 — Per-phase time breakdown
# ─────────────────────────────────────────────────────────────────────────────

def section1_time_breakdown(tokenizer, draft, target):
    section("1. TIME BREAKDOWN PER SPECULATIVE STEP")

    prompt = (
        "def binary_search(arr, target):\n"
        "    left, right = 0, len(arr) - 1\n"
        "    while left <= right:"
    )
    input_ids = tokenizer(
        prompt, return_tensors="pt", add_special_tokens=True
    ).input_ids.to(DEVICE)

    N_STEPS  = 20
    N_WARMUP = 5

    rows = []

    for gamma in [1, 2, 4, 8]:
        print(f"\n  gamma={gamma}")

        # Prime caches
        with torch.inference_mode():
            d_out = draft( input_ids=input_ids, use_cache=True, return_dict=True)
            t_out = target(input_ids=input_ids, use_cache=True, return_dict=True)
            draft_pkv  = d_out.past_key_values
            target_pkv = t_out.past_key_values

        cur_token = input_ids[:, -1:]

        # Warmup
        for _ in range(N_WARMUP):
            with torch.inference_mode():
                d_out = draft(input_ids=cur_token, past_key_values=draft_pkv,
                              use_cache=True, return_dict=True)

        times_draft  = []
        times_target = []
        times_sync   = []
        n_accepted_list = []

        for step in range(N_STEPS):
            # ── Draft phase ───────────────────────────────────────────────
            draft_tokens = []
            cur = cur_token

            _, t_draft = timed(lambda: None)  # reset
            sync()
            t0 = time.perf_counter()

            with torch.inference_mode():
                cur_dpkv = draft_pkv
                for _ in range(gamma):
                    d_out    = draft(input_ids=cur, past_key_values=cur_dpkv,
                                     use_cache=True, return_dict=True)
                    cur_dpkv = d_out.past_key_values
                    tok      = d_out.logits[:, -1, :].argmax(dim=-1, keepdim=True)
                    draft_tokens.append(tok)
                    cur = tok

            sync()
            t_draft = (time.perf_counter() - t0) * 1000.0
            draft_seq = torch.cat(draft_tokens, dim=1)

            # ── Target phase ──────────────────────────────────────────────
            target_input = torch.cat([cur_token, draft_seq], dim=1)

            sync()
            t0 = time.perf_counter()

            with torch.inference_mode():
                t_out = target(input_ids=target_input,
                               past_key_values=target_pkv,
                               use_cache=True, return_dict=True)

            sync()
            t_target = (time.perf_counter() - t0) * 1000.0
            target_logits = t_out.logits

            # ── Acceptance ────────────────────────────────────────────────
            accepted = 0
            new_tokens = []
            for i in range(gamma):
                target_tok = target_logits[:, i, :].argmax(dim=-1, keepdim=True)
                if target_tok.item() == draft_tokens[i].item():
                    accepted += 1
                    new_tokens.append(draft_tokens[i])
                else:
                    new_tokens.append(target_tok)
                    break
            else:
                bonus = target_logits[:, gamma, :].argmax(dim=-1, keepdim=True)
                new_tokens.append(bonus)

            new_seq   = torch.cat(new_tokens, dim=1)
            cur_token = new_seq[:, -1:]

            # ── KV sync phase ─────────────────────────────────────────────
            sync()
            t0 = time.perf_counter()

            with torch.inference_mode():
                d_sync = draft(input_ids=new_seq,
                               past_key_values=cur_dpkv,
                               use_cache=True, return_dict=True)
                draft_pkv  = d_sync.past_key_values
                target_pkv = t_out.past_key_values

            sync()
            t_sync = (time.perf_counter() - t0) * 1000.0

            times_draft.append(t_draft)
            times_target.append(t_target)
            times_sync.append(t_sync)
            n_accepted_list.append(accepted)

        mean_draft  = np.median(times_draft)
        mean_target = np.median(times_target)
        mean_sync   = np.median(times_sync)
        mean_acc    = np.mean(n_accepted_list) / gamma
        total       = mean_draft + mean_target + mean_sync

        pct_draft  = mean_draft  / total * 100
        pct_target = mean_target / total * 100
        pct_sync   = mean_sync   / total * 100

        print(f"    Draft  ({gamma} fwd): {mean_draft:7.2f}ms  ({pct_draft:.1f}%)")
        print(f"    Target (1 fwd):   {mean_target:7.2f}ms  ({pct_target:.1f}%)")
        print(f"    KV sync:          {mean_sync:7.2f}ms  ({pct_sync:.1f}%)")
        print(f"    Total:            {total:7.2f}ms")
        print(f"    Alpha:            {mean_acc:.3f}")

        rows.append({
            "gamma":          gamma,
            "draft_ms":       mean_draft,
            "target_ms":      mean_target,
            "sync_ms":        mean_sync,
            "total_ms":       total,
            "pct_draft":      pct_draft,
            "pct_target":     pct_target,
            "pct_sync":       pct_sync,
            "alpha":          mean_acc,
        })

    df = pd.DataFrame(rows)
    df.to_csv(os.path.join(RESULTS_DIR, "time_breakdown.csv"), index=False)
    return df


# ─────────────────────────────────────────────────────────────────────────────
# Section 2 — Alpha vs token position
# ─────────────────────────────────────────────────────────────────────────────

def section2_alpha_vs_position(tokenizer, draft, target):
    section("2. ALPHA vs TOKEN POSITION")

    prompts = {
        "code_python": (
            "def binary_search(arr, target):\n"
            "    left, right = 0, len(arr) - 1\n"
            "    while left <= right:"
        ),
        "creative": (
            "Once upon a time in a land far away there lived a brave knight"
        ),
    }

    gamma    = 4
    n_tokens = 64

    all_rows = []

    for prompt_name, prompt_text in prompts.items():
        input_ids = tokenizer(
            prompt_text, return_tensors="pt", add_special_tokens=True
        ).input_ids.to(DEVICE)

        with torch.inference_mode():
            d_out = draft( input_ids=input_ids, use_cache=True, return_dict=True)
            t_out = target(input_ids=input_ids, use_cache=True, return_dict=True)
            draft_pkv  = d_out.past_key_values
            target_pkv = t_out.past_key_values

        cur_token      = input_ids[:, -1:]
        tokens_so_far  = 0
        step           = 0

        while tokens_so_far < n_tokens:
            with torch.inference_mode():
                draft_tokens = []
                cur = cur_token
                cur_dpkv = draft_pkv

                for _ in range(gamma):
                    d_out    = draft(input_ids=cur, past_key_values=cur_dpkv,
                                     use_cache=True, return_dict=True)
                    cur_dpkv = d_out.past_key_values
                    tok      = d_out.logits[:, -1, :].argmax(dim=-1, keepdim=True)
                    draft_tokens.append(tok)
                    cur = tok

                draft_seq     = torch.cat(draft_tokens, dim=1)
                target_input  = torch.cat([cur_token, draft_seq], dim=1)
                t_out         = target(input_ids=target_input,
                                       past_key_values=target_pkv,
                                       use_cache=True, return_dict=True)
                target_logits = t_out.logits

                accepted   = 0
                new_tokens = []
                for i in range(gamma):
                    target_tok = target_logits[:, i, :].argmax(dim=-1, keepdim=True)
                    if target_tok.item() == draft_tokens[i].item():
                        accepted += 1
                        new_tokens.append(draft_tokens[i])
                    else:
                        new_tokens.append(target_tok)
                        break
                else:
                    bonus = target_logits[:, gamma, :].argmax(dim=-1, keepdim=True)
                    new_tokens.append(bonus)

                new_seq   = torch.cat(new_tokens, dim=1)
                cur_token = new_seq[:, -1:]

                d_sync     = draft(input_ids=new_seq, past_key_values=cur_dpkv,
                                   use_cache=True, return_dict=True)
                draft_pkv  = d_sync.past_key_values
                target_pkv = t_out.past_key_values

            alpha_step     = accepted / gamma
            tokens_so_far += new_seq.shape[1]

            all_rows.append({
                "prompt":      prompt_name,
                "step":        step,
                "token_pos":   tokens_so_far,
                "alpha":       alpha_step,
                "n_accepted":  accepted,
            })
            step += 1

        print(f"\n  {prompt_name}:")
        prompt_rows = [r for r in all_rows if r["prompt"] == prompt_name]
        alphas      = [r["alpha"] for r in prompt_rows]
        # Rolling mean
        window = 3
        rolling = [np.mean(alphas[max(0,i-window):i+1])
                   for i in range(len(alphas))]
        print(f"    Steps: {step}  Mean alpha: {np.mean(alphas):.3f}")
        print(f"    First 5 steps alpha: {[f'{a:.2f}' for a in alphas[:5]]}")
        print(f"    Last  5 steps alpha: {[f'{a:.2f}' for a in alphas[-5:]]}")

    df = pd.DataFrame(all_rows)
    df.to_csv(os.path.join(RESULTS_DIR, "alpha_vs_position.csv"), index=False)
    return df


# ─────────────────────────────────────────────────────────────────────────────
# Section 3 — Alpha vs temperature
# ─────────────────────────────────────────────────────────────────────────────

def section3_alpha_vs_temperature(tokenizer, draft, target):
    section("3. ALPHA vs TEMPERATURE — Does Sampling Help?")

    prompt = (
        "def binary_search(arr, target):\n"
        "    left, right = 0, len(arr) - 1\n"
        "    while left <= right:"
    )
    input_ids = tokenizer(
        prompt, return_tensors="pt", add_special_tokens=True
    ).input_ids.to(DEVICE)

    gamma       = 4
    n_steps     = 30
    temperatures = [0.0, 0.3, 0.5, 0.7, 1.0, 1.5]

    rows = []

    for temp in temperatures:
        gc.collect()
        torch.cuda.empty_cache()

        with torch.inference_mode():
            d_out = draft( input_ids=input_ids, use_cache=True, return_dict=True)
            t_out = target(input_ids=input_ids, use_cache=True, return_dict=True)
            draft_pkv  = d_out.past_key_values
            target_pkv = t_out.past_key_values

        cur_token   = input_ids[:, -1:]
        step_alphas = []

        for _ in range(n_steps):
            with torch.inference_mode():
                draft_tokens = []
                cur = cur_token
                cur_dpkv = draft_pkv

                for _ in range(gamma):
                    d_out = draft(input_ids=cur, past_key_values=cur_dpkv,
                                  use_cache=True, return_dict=True)
                    cur_dpkv = d_out.past_key_values
                    logits   = d_out.logits[:, -1, :]

                    if temp == 0.0:
                        tok = logits.argmax(dim=-1, keepdim=True)
                    else:
                        probs = torch.softmax(logits / temp, dim=-1)
                        tok   = torch.multinomial(probs, 1)

                    draft_tokens.append(tok)
                    cur = tok

                draft_seq    = torch.cat(draft_tokens, dim=1)
                target_input = torch.cat([cur_token, draft_seq], dim=1)
                t_out        = target(input_ids=target_input,
                                      past_key_values=target_pkv,
                                      use_cache=True, return_dict=True)
                target_logits = t_out.logits

                # Greedy acceptance regardless of sampling
                accepted   = 0
                new_tokens = []
                for i in range(gamma):
                    target_tok = target_logits[:, i, :].argmax(dim=-1, keepdim=True)
                    if target_tok.item() == draft_tokens[i].item():
                        accepted += 1
                        new_tokens.append(draft_tokens[i])
                    else:
                        new_tokens.append(target_tok)
                        break
                else:
                    bonus = target_logits[:, gamma, :].argmax(dim=-1, keepdim=True)
                    new_tokens.append(bonus)

                new_seq   = torch.cat(new_tokens, dim=1)
                cur_token = new_seq[:, -1:]
                d_sync    = draft(input_ids=new_seq, past_key_values=cur_dpkv,
                                  use_cache=True, return_dict=True)
                draft_pkv  = d_sync.past_key_values
                target_pkv = t_out.past_key_values

            step_alphas.append(accepted / gamma)

        mean_alpha = np.mean(step_alphas)
        rows.append({"temperature": temp, "alpha": mean_alpha,
                     "n_steps": n_steps})
        temp_str = "greedy" if temp == 0.0 else f"T={temp}"
        print(f"  {temp_str:>10}: alpha={mean_alpha:.3f}")

    df = pd.DataFrame(rows)
    df.to_csv(os.path.join(RESULTS_DIR, "alpha_vs_temperature.csv"), index=False)

    print(f"""
  Key insight:
    Greedy (T=0): alpha={df[df['temperature']==0.0]['alpha'].values[0]:.3f}
    T=1.0:        alpha={df[df['temperature']==1.0]['alpha'].values[0]:.3f}
    Higher temperature lowers alpha because sampling introduces
    randomness that diverges from the target's greedy prediction.
    Greedy acceptance requires EXACT token match — sampling makes
    this harder, not easier.
    """)

    return df


# ─────────────────────────────────────────────────────────────────────────────
# Section 4 — Draft-only timing isolation
# ─────────────────────────────────────────────────────────────────────────────

def section4_timing_isolation(tokenizer, draft, target):
    section("4. TIMING ISOLATION — Draft vs Target per Token")

    prompt = (
        "def binary_search(arr, target):\n"
        "    left, right = 0, len(arr) - 1\n"
        "    while left <= right:"
    )
    input_ids = tokenizer(
        prompt, return_tensors="pt", add_special_tokens=True
    ).input_ids.to(DEVICE)

    N = 50
    WARMUP = 10

    # Warmup
    with torch.inference_mode():
        d_out = draft( input_ids=input_ids, use_cache=True, return_dict=True)
        t_out = target(input_ids=input_ids, use_cache=True, return_dict=True)
        dpkv  = d_out.past_key_values
        tpkv  = t_out.past_key_values
        cur   = input_ids[:, -1:]

    for _ in range(WARMUP):
        with torch.inference_mode():
            o = draft(input_ids=cur, past_key_values=dpkv,
                      use_cache=True, return_dict=True)

    # Measure draft single token
    draft_times = []
    with torch.inference_mode():
        dpkv2 = dpkv
        for _ in range(N):
            sync()
            t0 = time.perf_counter()
            o  = draft(input_ids=cur, past_key_values=dpkv2,
                       use_cache=True, return_dict=True)
            sync()
            draft_times.append((time.perf_counter() - t0) * 1000.0)
            dpkv2 = o.past_key_values
            cur   = o.logits[:, -1, :].argmax(dim=-1, keepdim=True)

    # Measure target single token
    cur = input_ids[:, -1:]
    target_times = []
    with torch.inference_mode():
        tpkv2 = tpkv
        for _ in range(N):
            sync()
            t0 = time.perf_counter()
            o  = target(input_ids=cur, past_key_values=tpkv2,
                        use_cache=True, return_dict=True)
            sync()
            target_times.append((time.perf_counter() - t0) * 1000.0)
            tpkv2 = o.past_key_values
            cur   = o.logits[:, -1, :].argmax(dim=-1, keepdim=True)

    d_med = np.median(draft_times)
    t_med = np.median(target_times)
    d_p99 = np.percentile(draft_times, 99)
    t_p99 = np.percentile(target_times, 99)
    real_ratio = t_med / d_med

    print(f"""
  Single-token decode times (N={N} repetitions):

                   Draft (0.5B)    Target (1.5B)    Ratio
  Median:          {d_med:>10.3f}ms    {t_med:>10.3f}ms    {real_ratio:.3f}x
  P99:             {d_p99:>10.3f}ms    {t_p99:>10.3f}ms
  Std:             {np.std(draft_times):>10.3f}ms    {np.std(target_times):>10.3f}ms

  cost_ratio (for analytical model): {real_ratio:.4f}x

  Required for speedup > 1 at alpha=0.7, gamma=4:  2.26x
  Measured cost_ratio:                              {real_ratio:.2f}x
  Gap:                                              {2.26/real_ratio:.2f}x short
    """)

    rows = [{"model": "draft",  "times": draft_times,
             "median": d_med,   "p99": d_p99},
            {"model": "target", "times": target_times,
             "median": t_med,   "p99": t_p99}]

    df = pd.DataFrame([
        {"model": "draft",  "time_ms": t} for t in draft_times
    ] + [
        {"model": "target", "time_ms": t} for t in target_times
    ])
    df.to_csv(os.path.join(RESULTS_DIR, "timing_isolation.csv"), index=False)

    return d_med, t_med, real_ratio, draft_times, target_times


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

def generate_plots(breakdown_df, position_df, temp_df,
                   d_med, t_med, ratio, draft_times, target_times):
    section("GENERATING PLOTS")

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

    # ── Plot 1: Time breakdown stacked bar ────────────────────────────────
    fig, ax = plt.subplots(figsize=(10, 6))
    x   = np.arange(len(breakdown_df))
    w   = 0.5

    ax.bar(x, breakdown_df["draft_ms"],
           w, label="Draft forward passes", color="#3498DB")
    ax.bar(x, breakdown_df["target_ms"],
           w, bottom=breakdown_df["draft_ms"],
           label="Target forward pass", color="#E74C3C")
    ax.bar(x, breakdown_df["sync_ms"],
           w, bottom=breakdown_df["draft_ms"] + breakdown_df["target_ms"],
           label="KV cache sync", color="#E67E22")

    ax.set_xticks(x)
    ax.set_xticklabels([f"gamma={g}" for g in breakdown_df["gamma"]])
    ax.set_ylabel("Time per speculative step (ms)")
    ax.set_title("Time Breakdown per Speculative Step\n"
                 "code_python prompt — Qwen2-0.5B → Qwen2-1.5B")
    ax.legend()

    for i, row in breakdown_df.iterrows():
        total = row["total_ms"]
        ax.text(i, total + 0.5, f"{total:.1f}ms", ha="center", fontsize=9)

    plt.tight_layout()
    plt.savefig(os.path.join(PLOTS_DIR, "time_breakdown.png"),
                dpi=180, bbox_inches="tight")
    plt.close()
    print("  plot: plots/time_breakdown.png")

    # ── Plot 2: Alpha vs position ─────────────────────────────────────────
    fig, axes = plt.subplots(1, 2, figsize=(14, 5))

    for ax, prompt in zip(axes, position_df["prompt"].unique()):
        sub = position_df[position_df["prompt"] == prompt].sort_values("step")
        ax.scatter(sub["step"], sub["alpha"],
                   s=50, alpha=0.6, color="#3498DB", label="Per step")

        # Rolling mean
        window  = 4
        rolling = sub["alpha"].rolling(window, min_periods=1).mean()
        ax.plot(sub["step"], rolling,
                linewidth=2, color="#E74C3C", label=f"Rolling mean (w={window})")

        ax.axhline(sub["alpha"].mean(), color="gray", linestyle="--",
                   alpha=0.5, label=f"Mean={sub['alpha'].mean():.3f}")
        ax.set_xlabel("Speculative step")
        ax.set_ylabel("Acceptance rate (alpha)")
        ax.set_title(f"Alpha vs Token Position\n{prompt}")
        ax.set_ylim(-0.05, 1.05)
        ax.legend(fontsize=8)

    plt.suptitle("Does Acceptance Rate Change as Sequence Grows?",
                 fontsize=11, fontweight="bold")
    plt.tight_layout()
    plt.savefig(os.path.join(PLOTS_DIR, "alpha_vs_position.png"),
                dpi=180, bbox_inches="tight")
    plt.close()
    print("  plot: plots/alpha_vs_position.png")

    # ── Plot 3: Alpha vs temperature ──────────────────────────────────────
    fig, ax = plt.subplots(figsize=(9, 5))
    ax.plot(temp_df["temperature"], temp_df["alpha"],
            "o-", linewidth=2.5, color="#3498DB", markersize=8)
    ax.axhline(temp_df["alpha"].max(), color="gray",
               linestyle="--", alpha=0.5, label="Max alpha (greedy)")
    ax.set_xlabel("Sampling temperature")
    ax.set_ylabel("Mean token acceptance rate (alpha)")
    ax.set_title("Alpha vs Temperature — Greedy Acceptance\n"
                 "Higher temperature → more randomness → lower alpha")
    ax.annotate("Greedy\n(T=0)", xy=(0, temp_df[temp_df["temperature"]==0]["alpha"].values[0]),
                xytext=(0.2, temp_df[temp_df["temperature"]==0]["alpha"].values[0] - 0.08),
                arrowprops=dict(arrowstyle="->"), fontsize=9)
    plt.tight_layout()
    plt.savefig(os.path.join(PLOTS_DIR, "alpha_vs_temperature.png"),
                dpi=180, bbox_inches="tight")
    plt.close()
    print("  plot: plots/alpha_vs_temperature.png")

    # ── Plot 4: Draft vs target latency distribution ───────────────────────
    fig, axes = plt.subplots(1, 2, figsize=(14, 5))

    ax = axes[0]
    ax.hist(draft_times,  bins=30, alpha=0.7, color="#3498DB",
            label=f"Draft 0.5B (median={d_med:.2f}ms)")
    ax.hist(target_times, bins=30, alpha=0.7, color="#E74C3C",
            label=f"Target 1.5B (median={t_med:.2f}ms)")
    ax.axvline(d_med, color="#3498DB", linewidth=2, linestyle="--")
    ax.axvline(t_med, color="#E74C3C", linewidth=2, linestyle="--")
    ax.set_xlabel("Time per token (ms)")
    ax.set_ylabel("Count")
    ax.set_title("Single-Token Decode Latency Distribution\n"
                 f"cost_ratio = {ratio:.3f}x  (expected 3.1x from params)")
    ax.legend()

    ax2 = axes[1]
    # KV sync overhead as % of total step time
    for _, row in breakdown_df.iterrows():
        g     = row["gamma"]
        total = row["total_ms"]
        sync_ = row["sync_ms"]
        ax2.bar(g, sync_ / total * 100,
                color=GAMMA_COLORS.get(g, "#888"), alpha=0.85,
                label=f"gamma={g}")
        ax2.text(g, sync_ / total * 100 + 0.5,
                 f"{sync_/total*100:.1f}%",
                 ha="center", fontsize=9)

    ax2.set_xlabel("Gamma (draft tokens per step)")
    ax2.set_ylabel("KV sync as % of total step time")
    ax2.set_title("KV Cache Sync Overhead\n"
                  "Grows at higher gamma due to more tokens to re-sync")
    ax2.set_xticks([1, 2, 4, 8])
    ax2.legend()

    plt.suptitle("Single GPU Timing Analysis — RTX 2070",
                 fontsize=11, fontweight="bold")
    plt.tight_layout()
    plt.savefig(os.path.join(PLOTS_DIR, "timing_analysis.png"),
                dpi=180, bbox_inches="tight")
    plt.close()
    print("  plot: plots/timing_analysis.png")

    # ── Plot 5: KV sync % per phase per gamma ────────────────────────────
    fig, ax = plt.subplots(figsize=(10, 5))
    x = np.arange(len(breakdown_df))
    w = 0.25

    ax.bar(x - w, breakdown_df["pct_draft"],
           w, label="Draft", color="#3498DB")
    ax.bar(x,     breakdown_df["pct_target"],
           w, label="Target", color="#E74C3C")
    ax.bar(x + w, breakdown_df["pct_sync"],
           w, label="KV sync", color="#E67E22")

    ax.set_xticks(x)
    ax.set_xticklabels([f"gamma={g}" for g in breakdown_df["gamma"]])
    ax.set_ylabel("% of total step time")
    ax.set_title("Time Distribution by Phase\n"
                 "KV sync grows from 20% to 40% as gamma increases")
    ax.legend()
    plt.tight_layout()
    plt.savefig(os.path.join(PLOTS_DIR, "phase_distribution.png"),
                dpi=180, bbox_inches="tight")
    plt.close()
    print("  plot: plots/phase_distribution.png")


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

def main():
    print(json.dumps({
        "timestamp":    datetime.now(timezone.utc).isoformat(),
        "script":       "phase3_profiling.py",
        "draft_model":  DRAFT_MODEL,
        "target_model": TARGET_MODEL,
    }, indent=2))

    print("\nLoading models...")
    tokenizer, draft, target = load_models()

    breakdown_df              = section1_time_breakdown(tokenizer, draft, target)
    position_df               = section2_alpha_vs_position(tokenizer, draft, target)
    temp_df                   = section3_alpha_vs_temperature(tokenizer, draft, target)
    d_med, t_med, ratio, dt, tt = section4_timing_isolation(tokenizer, draft, target)

    generate_plots(breakdown_df, position_df, temp_df,
                   d_med, t_med, ratio, dt, tt)

    print(f"\n{DIVIDER}")
    print("  phase3_profiling.py complete.")
    print(DIVIDER)


if __name__ == "__main__":
    main()