Files
kernbench2/docs/sweeps/plot_scripts/plot_batch_scaling.py
T
ywkang baf22f01ab plot: extend batch-scaling curves with saturated-throughput ceiling
Beyond a mapping's SIP capacity (16/C) the throughput plateaus: extra
users run in waves at the saturated rate, they are not un-runnable. Draw
a dashed horizontal extension at the ceiling so 1-kv-per-cube (cap 2)
reads as saturating, not stopping, at B=2. Caption updated.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-22 17:39:13 -07:00

70 lines
2.8 KiB
Python

"""Batch-scaling figure for concurrent decode users on one SIP.
Reads ``docs/sweeps/decode_batch_sweep.csv`` (from
``tests/attention/test_gqa_decode_batch_sweep.py``) and plots aggregate
throughput (requests / us) versus batch size B for the four KV mappings,
one panel per context length. Each mapping's curve runs up to its SIP
capacity 16/C (A1=2 ... B=16 users).
"""
from pathlib import Path
import csv
import matplotlib.pyplot as plt
ROOT = Path(__file__).resolve().parents[3]
CSV = ROOT / "docs" / "sweeps" / "decode_batch_sweep.csv"
OUT = (ROOT / "docs" / "report" / "1H-codesign-paper"
/ "figures" / "gqa_short_context" / "batch_scaling.png")
MODE_LABEL = {"A1": "1-kv-per-cube", "A2": "2-kv-per-cube",
"A4": "4-kv-per-cube", "B": "8-kv-per-cube"}
MODE_COLOR = {"A1": "#1f77b4", "A2": "#2ca02c",
"A4": "#ffd000", "B": "#d62728"}
MODES = ["A1", "A2", "A4", "B"]
def main():
rows = list(csv.DictReader(CSV.open()))
contexts = sorted({int(r["S_kv"]) for r in rows})
fig, axes = plt.subplots(1, len(contexts), figsize=(6.5 * len(contexts), 4.5),
squeeze=False)
xmax = max(int(r["B"]) for r in rows) # largest SIP capacity (= 16)
for col, S in enumerate(contexts):
ax = axes[0][col]
for m in MODES:
pts = sorted(
((int(r["B"]), float(r["throughput_users_per_us"]))
for r in rows if r["mode"] == m and int(r["S_kv"]) == S),
key=lambda t: t[0],
)
if not pts:
continue
xs, ys = zip(*pts)
# Measured: concurrent users up to the SIP capacity 16/C.
ax.plot(xs, ys, marker="o", color=MODE_COLOR[m],
label=MODE_LABEL[m])
# Beyond capacity the SIP is full: extra users run in waves at
# the saturated rate, so throughput plateaus. Draw that ceiling
# as a dashed extension (projected, not concurrently measured).
if xs[-1] < xmax:
ax.plot([xs[-1], xmax], [ys[-1], ys[-1]],
linestyle="--", color=MODE_COLOR[m], alpha=0.55)
ax.annotate(f"{ys[-1]:.2f}", (xs[-1], ys[-1]),
textcoords="offset points", xytext=(4, 4), fontsize=8)
ax.set_title(f"S_kv = {S // 1024}K")
ax.set_xlabel("batch size B (concurrent users)")
ax.set_ylabel("throughput (requests / µs)")
ax.grid(True, alpha=0.3)
ax.legend(fontsize=9)
fig.suptitle("Batch scaling: aggregate throughput vs concurrent users "
"(one SIP)", fontsize=12, fontweight="bold")
fig.tight_layout()
OUT.parent.mkdir(parents=True, exist_ok=True)
fig.savefig(OUT, dpi=140, bbox_inches="tight")
print(f"{OUT}")
if __name__ == "__main__":
main()