long-context-benchmark/plot_long_context.py at main · JohnScheuer/long-context-benchmark · 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
107
108
109
110
111
112
113
114
import os
import matplotlib.pyplot as plt
import pandas as pd
import numpy as np

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

df = pd.read_csv("results/long_context_final.csv")
ok = df[df["status"] == "ok"].copy()

fig, ax = plt.subplots(figsize=(11, 6))
for bs in sorted(ok["batch_size"].unique()):
    sub = ok[ok["batch_size"] == bs].sort_values("context_len")
    ax.plot(sub["context_len"], sub["prefill_ms"], marker="o", linewidth=2, label=f"bs={bs}")
ax.set_xlabel("Context length (tokens)")
ax.set_ylabel("Prefill latency (ms)")
ax.set_title("Prefill latency vs context length (Qwen2-0.5B on RTX 2070)")
ax.set_xscale("log", base=2)
ax.set_yscale("log")
ax.grid(True, alpha=0.3, which="both")
ax.legend()
plt.tight_layout()
plt.savefig("plots/prefill_vs_context.png", dpi=180, bbox_inches="tight")
plt.close()
print("plot: plots/prefill_vs_context.png")

fig, ax = plt.subplots(figsize=(11, 6))
for bs in sorted(ok["batch_size"].unique()):
    sub = ok[ok["batch_size"] == bs].sort_values("context_len")
    ax.plot(sub["context_len"], sub["throughput_tok_s"], marker="o", linewidth=2, label=f"bs={bs}")
ax.set_xlabel("Context length (tokens)")
ax.set_ylabel("Decode throughput (tok/s)")
ax.set_title("Decode throughput vs context length")
ax.set_xscale("log", base=2)
ax.grid(True, alpha=0.3, which="both")
ax.legend()
plt.tight_layout()
plt.savefig("plots/throughput_vs_context.png", dpi=180, bbox_inches="tight")
plt.close()
print("plot: plots/throughput_vs_context.png")

fig, ax = plt.subplots(figsize=(11, 6))
for bs in sorted(ok["batch_size"].unique()):
    sub = ok[ok["batch_size"] == bs].sort_values("context_len")
    ax.plot(sub["context_len"], sub["peak_total_mb"], marker="o", linewidth=2, label=f"bs={bs}")
ax.axhline(8600, color="red", linestyle="--", alpha=0.7, label="RTX 2070 VRAM (8.6 GB)")
ax.set_xlabel("Context length (tokens)")
ax.set_ylabel("Peak memory (MB)")
ax.set_title("Peak VRAM vs context length")
ax.set_xscale("log", base=2)
ax.grid(True, alpha=0.3, which="both")
ax.legend()
plt.tight_layout()
plt.savefig("plots/memory_vs_context.png", dpi=180, bbox_inches="tight")
plt.close()
print("plot: plots/memory_vs_context.png")

max_bs = ok.groupby("context_len")["batch_size"].max().reset_index()
fig, ax = plt.subplots(figsize=(10, 6))
ax.bar(max_bs["context_len"].astype(str), max_bs["batch_size"], color="#2ECC71")
ax.set_xlabel("Context length")
ax.set_ylabel("Max batch size on RTX 2070")
ax.set_title("Maximum serving batch at each context length")
ax.grid(True, alpha=0.3, axis="y")
plt.tight_layout()
plt.savefig("plots/max_batch_capacity.png", dpi=180, bbox_inches="tight")
plt.close()
print("plot: plots/max_batch_capacity.png")

all_ctx = sorted(df["context_len"].unique())
all_bs = sorted(df["batch_size"].unique())
matrix = []
for ctx in all_ctx:
    row = []
    for bs in all_bs:
        sub = df[(df["context_len"] == ctx) & (df["batch_size"] == bs)]
        if len(sub) > 0 and sub.iloc[0]["status"] == "ok":
            row.append(1)
        else:
            row.append(0)
    matrix.append(row)

fig, ax = plt.subplots(figsize=(8, 6))
im = ax.imshow(matrix, aspect="auto", cmap="RdYlGn", vmin=0, vmax=1, origin="lower")
ax.set_xticks(range(len(all_bs)))
ax.set_xticklabels(all_bs)
ax.set_yticks(range(len(all_ctx)))
ax.set_yticklabels(all_ctx)
ax.set_xlabel("Batch size")
ax.set_ylabel("Context length")
ax.set_title("RTX 2070 capacity map (green=fits, red=OOM)")
for i in range(len(all_ctx)):
    for j in range(len(all_bs)):
        label = "OK" if matrix[i][j] else "OOM"
        color = "black" if matrix[i][j] else "white"
        ax.text(j, i, label, ha="center", va="center", fontsize=10, color=color, fontweight="bold")
plt.tight_layout()
plt.savefig("plots/capacity_heatmap.png", dpi=180, bbox_inches="tight")
plt.close()
print("plot: plots/capacity_heatmap.png")

print("\n=== Key findings ===")
print("\nMax batch per context:")
for _, r in max_bs.iterrows():
    sub = ok[(ok["context_len"] == r["context_len"]) & (ok["batch_size"] == r["batch_size"])]
    if len(sub) > 0:
        peak = sub.iloc[0]["peak_total_mb"]
        tput = sub.iloc[0]["throughput_tok_s"]
        print(f"  ctx={int(r['context_len']):6d}: bs={int(r['batch_size'])}  peak={peak:.0f}MB  tput={tput:.1f}tok/s")

print("\nPrefill scaling (bs=1):")
bs1 = ok[ok["batch_size"] == 1].sort_values("context_len")
for _, r in bs1.iterrows():
    print(f"  ctx={int(r['context_len']):6d}: {r['prefill_ms']:.1f}ms")