speculative-decoding-impl/plot_speculative.py at main · JohnScheuer/speculative-decoding-impl · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
import os
import matplotlib.pyplot as plt
import pandas as pd

os.makedirs("plots", exist_ok=True)

df = pd.read_csv("results/speculative_decoding_results.csv")

baseline = df[df["method"] == "baseline"].copy()
spec = df[df["method"] == "speculative"].copy()

# Save best-k summary
best_rows = []
for prompt, sub in spec.groupby("prompt"):
    best = sub.loc[sub["speedup_vs_baseline"].idxmax()]
    base = baseline[baseline["prompt"] == prompt].iloc[0]
    best_rows.append({
        "prompt": prompt,
        "baseline_tput": base["throughput_tok_s"],
        "best_k": int(best["k"]),
        "best_tput": best["throughput_tok_s"],
        "best_speedup": best["speedup_vs_baseline"],
        "best_accept_rate": best["mean_accept_rate"],
    })

best_df = pd.DataFrame(best_rows)
best_df.to_csv("results/best_k_by_prompt.csv", index=False)
print("saved: results/best_k_by_prompt.csv")
print(best_df.to_string(index=False))

# ── Plot 1: speedup vs K ──────────────────────────────────────────────────
plt.figure(figsize=(10, 6))
for prompt, sub in spec.groupby("prompt"):
    sub = sub.sort_values("k")
    plt.plot(sub["k"], sub["speedup_vs_baseline"], marker="o", linewidth=2, label=prompt)
plt.axhline(1.0, color="gray", linestyle="--", alpha=0.5, label="baseline")
plt.xlabel("Draft tokens K")
plt.ylabel("Speedup vs baseline (×)")
plt.title("Speculative decoding speedup vs K")
plt.legend()
plt.grid(True, alpha=0.3)
plt.tight_layout()
plt.savefig("plots/speedup_vs_k.png", dpi=180, bbox_inches="tight")
plt.close()
print("plot: plots/speedup_vs_k.png")

# ── Plot 2: acceptance rate vs K ──────────────────────────────────────────
plt.figure(figsize=(10, 6))
for prompt, sub in spec.groupby("prompt"):
    sub = sub.sort_values("k")
    plt.plot(sub["k"], sub["mean_accept_rate"], marker="o", linewidth=2, label=prompt)
plt.xlabel("Draft tokens K")
plt.ylabel("Mean acceptance rate")
plt.title("Acceptance rate vs K")
plt.legend()
plt.grid(True, alpha=0.3)
plt.tight_layout()
plt.savefig("plots/acceptance_vs_k.png", dpi=180, bbox_inches="tight")
plt.close()
print("plot: plots/acceptance_vs_k.png")

# ── Plot 3: throughput vs K with baseline lines ───────────────────────────
fig, axes = plt.subplots(2, 2, figsize=(14, 10))
axes = axes.flatten()

for ax, prompt in zip(axes, sorted(spec["prompt"].unique())):
    sub = spec[spec["prompt"] == prompt].sort_values("k")
    base = baseline[baseline["prompt"] == prompt]["throughput_tok_s"].iloc[0]

    ax.plot(sub["k"], sub["throughput_tok_s"], marker="o", linewidth=2, color="#4C9BE8")
    ax.axhline(base, color="#E74C3C", linestyle="--", alpha=0.7, label="baseline")
    ax.set_title(prompt)
    ax.set_xlabel("Draft tokens K")
    ax.set_ylabel("Throughput (tok/s)")
    ax.grid(True, alpha=0.3)
    ax.legend()

plt.suptitle("Throughput vs K by prompt", y=1.02)
plt.tight_layout()
plt.savefig("plots/throughput_vs_k_by_prompt.png", dpi=180, bbox_inches="tight")
plt.close()
print("plot: plots/throughput_vs_k_by_prompt.png")

# ── Plot 4: acceptance vs speedup scatter ─────────────────────────────────
plt.figure(figsize=(10, 6))
for prompt, sub in spec.groupby("prompt"):
    plt.scatter(
        sub["mean_accept_rate"],
        sub["speedup_vs_baseline"],
        s=90,
        label=prompt,
    )
    for _, r in sub.iterrows():
        plt.annotate(f"K={int(r['k'])}", (r["mean_accept_rate"], r["speedup_vs_baseline"]),
                     textcoords="offset points", xytext=(5, 5), fontsize=8)

plt.axhline(1.0, color="gray", linestyle="--", alpha=0.5)
plt.xlabel("Acceptance rate")
plt.ylabel("Speedup vs baseline (×)")
plt.title("Acceptance rate vs speedup")
plt.legend()
plt.grid(True, alpha=0.3)
plt.tight_layout()
plt.savefig("plots/acceptance_vs_speedup.png", dpi=180, bbox_inches="tight")
plt.close()
print("plot: plots/acceptance_vs_speedup.png")