plot: merge plot scripts into single plot_short_context.py
This commit is contained in:
eusonice
@@ -0,0 +1,231 @@
|
|||||||
|
"""Generate 6 comparison figures for the short-context GQA attention sweep.
|
||||||
|
|
||||||
|
Sweep matrix: 3 variants × 2 phases × 4 modes × 4 contexts (8K/16K/32K/64K).
|
||||||
|
|
||||||
|
Figures
|
||||||
|
Wall-clock latency:
|
||||||
|
1. wall_variant1_mode_compare.png — (1) without composite
|
||||||
|
2. wall_variant2_mode_compare.png — (2) with composite (GEMM-only)
|
||||||
|
3. wall_variant3_mode_compare.png — (3) with composite + softmax_merge
|
||||||
|
4. wall_a1_variant_compare.png — 1-kv-per-cube: variant comparison
|
||||||
|
|
||||||
|
Mode trade-off + composite ablation:
|
||||||
|
5. per_cube_tradeoff_mode_compare.png — HBM BW + IPCQ (variant-invariant)
|
||||||
|
6. gemm_util_a1_variant_ablation.png — 1-kv-per-cube GEMM util
|
||||||
|
|
||||||
|
Output: docs/report/1H-codesign-paper/figures/gqa_short_context/
|
||||||
|
"""
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import csv
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
import matplotlib.pyplot as plt
|
||||||
|
import numpy as np
|
||||||
|
|
||||||
|
ROOT = Path("/Users/lalicorne/Desktop/ahbm/kernbench2")
|
||||||
|
SWEEPS = ROOT / "docs" / "sweeps"
|
||||||
|
OUT = (ROOT / "docs" / "report" / "1H-codesign-paper"
|
||||||
|
/ "figures" / "gqa_short_context")
|
||||||
|
OUT.mkdir(parents=True, exist_ok=True)
|
||||||
|
|
||||||
|
CONTEXTS = [8192, 16384, 32768, 65536]
|
||||||
|
CONTEXT_LABELS = ["8K", "16K", "32K", "64K"]
|
||||||
|
MODES = ["A1", "A2", "A4", "B"]
|
||||||
|
MODE_LABEL = {
|
||||||
|
"A1": "1-kv-per-cube",
|
||||||
|
"A2": "2-kv-per-cube",
|
||||||
|
"A4": "4-kv-per-cube",
|
||||||
|
"B": "8-kv-per-cube",
|
||||||
|
}
|
||||||
|
MODE_COLOR = {
|
||||||
|
MODE_LABEL["A1"]: "#1f77b4",
|
||||||
|
MODE_LABEL["A2"]: "#2ca02c",
|
||||||
|
MODE_LABEL["A4"]: "#ffd000",
|
||||||
|
MODE_LABEL["B"]: "#d62728",
|
||||||
|
}
|
||||||
|
VARIANTS = [
|
||||||
|
("baseline", "(1) without composite"),
|
||||||
|
("composite", "(2) with composite (GEMM-only)"),
|
||||||
|
("composite_fused", "(3) with composite + softmax_merge"),
|
||||||
|
]
|
||||||
|
VARIANT_COLOR = {
|
||||||
|
"baseline": "#1f77b4",
|
||||||
|
"composite": "#2ca02c",
|
||||||
|
"composite_fused": "#d62728",
|
||||||
|
}
|
||||||
|
|
||||||
|
CSV_MAP = {
|
||||||
|
("decode", "baseline"): SWEEPS / "short_context_decode_sweep.csv",
|
||||||
|
("decode", "composite"): SWEEPS / "short_context_decode_composite_sweep.csv",
|
||||||
|
("decode", "composite_fused"): SWEEPS / "short_context_decode_composite_fused_sweep.csv",
|
||||||
|
("prefill", "baseline"): SWEEPS / "short_context_prefill_sweep.csv",
|
||||||
|
("prefill", "composite"): SWEEPS / "short_context_prefill_composite_sweep.csv",
|
||||||
|
("prefill", "composite_fused"): SWEEPS / "short_context_prefill_composite_fused_sweep.csv",
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def load(phase, variant):
|
||||||
|
"""Return {(mode, S_kv): row dict}."""
|
||||||
|
rows = {}
|
||||||
|
with CSV_MAP[(phase, variant)].open() as f:
|
||||||
|
for r in csv.DictReader(f):
|
||||||
|
rows[(r["mode"], int(r["S_kv"]))] = r
|
||||||
|
return rows
|
||||||
|
|
||||||
|
|
||||||
|
def grouped_bars(ax, *, x_labels, groups, group_color, ylabel, title,
|
||||||
|
y_log=False, value_fmt="{:.3g}"):
|
||||||
|
"""Grouped bars: x = x_labels, groups = list of (label, values)."""
|
||||||
|
n = len(groups)
|
||||||
|
nx = len(x_labels)
|
||||||
|
width = 0.8 / n
|
||||||
|
x = np.arange(nx)
|
||||||
|
for i, (label, values) in enumerate(groups):
|
||||||
|
offset = (i - (n - 1) / 2) * width
|
||||||
|
bars = ax.bar(x + offset, values, width, label=label,
|
||||||
|
color=group_color[label], edgecolor="black",
|
||||||
|
linewidth=0.4)
|
||||||
|
for bar, v in zip(bars, values):
|
||||||
|
if v <= 0 and y_log:
|
||||||
|
continue
|
||||||
|
ax.text(bar.get_x() + bar.get_width() / 2, bar.get_height(),
|
||||||
|
value_fmt.format(v), ha="center", va="bottom",
|
||||||
|
fontsize=7)
|
||||||
|
ax.set_xticks(x)
|
||||||
|
ax.set_xticklabels(x_labels)
|
||||||
|
ax.set_xlabel("S_kv")
|
||||||
|
ax.set_ylabel(ylabel)
|
||||||
|
ax.set_title(title)
|
||||||
|
if y_log:
|
||||||
|
ax.set_yscale("log")
|
||||||
|
ax.grid(True, axis="y", alpha=0.3, which="both")
|
||||||
|
ax.legend(fontsize=8, loc="upper left")
|
||||||
|
|
||||||
|
|
||||||
|
# ── 1-3. Wall mode comparison per variant ──────────────────────────
|
||||||
|
|
||||||
|
def fig_wall_mode_compare(variant_key, suptitle, fname):
|
||||||
|
fig, axes = plt.subplots(1, 2, figsize=(13, 4.8))
|
||||||
|
for ax, phase in zip(axes, ("prefill", "decode")):
|
||||||
|
rows = load(phase, variant_key)
|
||||||
|
groups = [
|
||||||
|
(MODE_LABEL[m], [float(rows[(m, s)]["wall_us"]) for s in CONTEXTS])
|
||||||
|
for m in MODES
|
||||||
|
]
|
||||||
|
grouped_bars(ax, x_labels=CONTEXT_LABELS, groups=groups,
|
||||||
|
group_color=MODE_COLOR,
|
||||||
|
ylabel="Wall-clock latency (μs)",
|
||||||
|
title=phase.capitalize(),
|
||||||
|
y_log=True, value_fmt="{:.0f}")
|
||||||
|
fig.suptitle(suptitle, fontsize=12, fontweight="bold")
|
||||||
|
fig.tight_layout()
|
||||||
|
out = OUT / fname
|
||||||
|
fig.savefig(out, dpi=140, bbox_inches="tight")
|
||||||
|
plt.close(fig)
|
||||||
|
print(f" ✓ {out}")
|
||||||
|
|
||||||
|
|
||||||
|
# ── 4. Wall 1-kv-per-cube variant comparison ───────────────────────
|
||||||
|
|
||||||
|
def fig_wall_a1_variant_compare(fname):
|
||||||
|
fig, axes = plt.subplots(1, 2, figsize=(13, 4.8))
|
||||||
|
for ax, phase in zip(axes, ("prefill", "decode")):
|
||||||
|
groups = []
|
||||||
|
for variant_key, _ in VARIANTS:
|
||||||
|
rows = load(phase, variant_key)
|
||||||
|
values = [float(rows[("A1", s)]["wall_us"]) for s in CONTEXTS]
|
||||||
|
groups.append((variant_key, values))
|
||||||
|
grouped_bars(ax, x_labels=CONTEXT_LABELS, groups=groups,
|
||||||
|
group_color=VARIANT_COLOR,
|
||||||
|
ylabel="Wall-clock latency (μs)",
|
||||||
|
title=phase.capitalize(),
|
||||||
|
y_log=True, value_fmt="{:.0f}")
|
||||||
|
ax.legend([v[1] for v in VARIANTS], fontsize=8, loc="upper left")
|
||||||
|
fig.suptitle("1-kv-per-cube: variant comparison",
|
||||||
|
fontsize=12, fontweight="bold")
|
||||||
|
fig.tight_layout()
|
||||||
|
out = OUT / fname
|
||||||
|
fig.savefig(out, dpi=140, bbox_inches="tight")
|
||||||
|
plt.close(fig)
|
||||||
|
print(f" ✓ {out}")
|
||||||
|
|
||||||
|
|
||||||
|
# ── 5. Per-cube trade-off (HBM BW + IPCQ, mode compare) ────────────
|
||||||
|
|
||||||
|
def fig_per_cube_tradeoff(fname):
|
||||||
|
fig, axes = plt.subplots(2, 2, figsize=(13, 9))
|
||||||
|
metrics = [
|
||||||
|
("hbm_bw_util", "HBM BW utilization (per PE)", False, "{:.3f}"),
|
||||||
|
("ipcq_kb", "IPCQ traffic (KB)", True, "{:.1f}"),
|
||||||
|
]
|
||||||
|
for row, (key, ylabel, log, fmt) in enumerate(metrics):
|
||||||
|
for col, phase in enumerate(("prefill", "decode")):
|
||||||
|
ax = axes[row][col]
|
||||||
|
rows = load(phase, "baseline")
|
||||||
|
groups = [
|
||||||
|
(MODE_LABEL[m], [float(rows[(m, s)][key]) for s in CONTEXTS])
|
||||||
|
for m in MODES
|
||||||
|
]
|
||||||
|
grouped_bars(ax, x_labels=CONTEXT_LABELS, groups=groups,
|
||||||
|
group_color=MODE_COLOR, ylabel=ylabel,
|
||||||
|
title=phase.capitalize(),
|
||||||
|
y_log=log, value_fmt=fmt)
|
||||||
|
fig.suptitle("Per-cube trade-off: HBM BW + IPCQ mode comparison",
|
||||||
|
fontsize=12, fontweight="bold")
|
||||||
|
fig.tight_layout()
|
||||||
|
out = OUT / fname
|
||||||
|
fig.savefig(out, dpi=140, bbox_inches="tight")
|
||||||
|
plt.close(fig)
|
||||||
|
print(f" ✓ {out}")
|
||||||
|
|
||||||
|
|
||||||
|
# ── 6. GEMM util ablation (1-kv-per-cube, variant compare) ─────────
|
||||||
|
|
||||||
|
def fig_gemm_util_a1_ablation(fname):
|
||||||
|
fig, axes = plt.subplots(1, 2, figsize=(13, 5.4))
|
||||||
|
for ax, phase in zip(axes, ("prefill", "decode")):
|
||||||
|
groups = []
|
||||||
|
for variant_key, _ in VARIANTS:
|
||||||
|
rows = load(phase, variant_key)
|
||||||
|
values = [float(rows[("A1", s)]["gemm_util"]) for s in CONTEXTS]
|
||||||
|
groups.append((variant_key, values))
|
||||||
|
grouped_bars(ax, x_labels=CONTEXT_LABELS, groups=groups,
|
||||||
|
group_color=VARIANT_COLOR,
|
||||||
|
ylabel="GEMM engine utilization (per PE)",
|
||||||
|
title=phase.capitalize(),
|
||||||
|
y_log=False, value_fmt="{:.4f}")
|
||||||
|
ax.legend([v[1] for v in VARIANTS], fontsize=8, loc="upper left")
|
||||||
|
fig.suptitle("1-kv-per-cube: GEMM utilization variant comparison",
|
||||||
|
fontsize=12, fontweight="bold")
|
||||||
|
fig.text(
|
||||||
|
0.5, -0.04,
|
||||||
|
"⚠ Higher gemm_util on composite/fused is largely supertile "
|
||||||
|
"padding overhead (M=G=8 padded to TILE_M=32), not pure fusion "
|
||||||
|
"gain. See ADR-0070 limitation #4.",
|
||||||
|
ha="center", fontsize=9, style="italic", color="#555")
|
||||||
|
fig.tight_layout()
|
||||||
|
out = OUT / fname
|
||||||
|
fig.savefig(out, dpi=140, bbox_inches="tight")
|
||||||
|
plt.close(fig)
|
||||||
|
print(f" ✓ {out}")
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
print(f"Output: {OUT}")
|
||||||
|
fig_wall_mode_compare(
|
||||||
|
"baseline",
|
||||||
|
"(1) Without composite: latency mode comparison",
|
||||||
|
"wall_variant1_mode_compare.png")
|
||||||
|
fig_wall_mode_compare(
|
||||||
|
"composite",
|
||||||
|
"(2) With composite (GEMM-only): latency mode comparison",
|
||||||
|
"wall_variant2_mode_compare.png")
|
||||||
|
fig_wall_mode_compare(
|
||||||
|
"composite_fused",
|
||||||
|
"(3) With composite + softmax_merge: latency mode comparison",
|
||||||
|
"wall_variant3_mode_compare.png")
|
||||||
|
fig_wall_a1_variant_compare("wall_a1_variant_compare.png")
|
||||||
|
fig_per_cube_tradeoff("per_cube_tradeoff_mode_compare.png")
|
||||||
|
fig_gemm_util_a1_ablation("gemm_util_a1_variant_ablation.png")
|
||||||
|
print("\nDone.")
|
||||||
@@ -1,165 +0,0 @@
|
|||||||
"""B: per-cube trade-off (mode comparison, variant-invariant)
|
|
||||||
C: GEMM util ablation (A1 only, variant comparison)
|
|
||||||
|
|
||||||
Output: docs/report/1H-codesign-paper/figures/gqa_short_context/
|
|
||||||
"""
|
|
||||||
from __future__ import annotations
|
|
||||||
|
|
||||||
import csv
|
|
||||||
from pathlib import Path
|
|
||||||
|
|
||||||
import matplotlib.pyplot as plt
|
|
||||||
import numpy as np
|
|
||||||
|
|
||||||
ROOT = Path("/Users/lalicorne/Desktop/ahbm/kernbench2")
|
|
||||||
SWEEPS = ROOT / "docs" / "sweeps"
|
|
||||||
OUT = ROOT / "docs" / "report" / "1H-codesign-paper" / "figures" / "gqa_short_context"
|
|
||||||
OUT.mkdir(parents=True, exist_ok=True)
|
|
||||||
|
|
||||||
CONTEXTS = [8192, 16384, 32768, 65536]
|
|
||||||
CONTEXT_LABELS = ["8K", "16K", "32K", "64K"]
|
|
||||||
MODES = ["A1", "A2", "A4", "B"]
|
|
||||||
MODE_LABEL = {"A1": "1-kv-per-cube", "A2": "2-kv-per-cube",
|
|
||||||
"A4": "4-kv-per-cube", "B": "8-kv-per-cube"}
|
|
||||||
MODE_COLOR = {MODE_LABEL["A1"]: "#1f77b4",
|
|
||||||
MODE_LABEL["A2"]: "#2ca02c",
|
|
||||||
MODE_LABEL["A4"]: "#ffd000",
|
|
||||||
MODE_LABEL["B"]: "#d62728"}
|
|
||||||
VARIANTS = [
|
|
||||||
("baseline", "(1) without composite"),
|
|
||||||
("composite", "(2) with composite (GEMM-only)"),
|
|
||||||
("composite_fused", "(3) with composite + softmax_merge"),
|
|
||||||
]
|
|
||||||
VARIANT_COLOR = {
|
|
||||||
"baseline": "#1f77b4",
|
|
||||||
"composite": "#2ca02c",
|
|
||||||
"composite_fused": "#d62728",
|
|
||||||
}
|
|
||||||
|
|
||||||
CSV_MAP = {
|
|
||||||
("decode", "baseline"): SWEEPS / "short_context_decode_sweep.csv",
|
|
||||||
("decode", "composite"): SWEEPS / "short_context_decode_composite_sweep.csv",
|
|
||||||
("decode", "composite_fused"): SWEEPS / "short_context_decode_composite_fused_sweep.csv",
|
|
||||||
("prefill", "baseline"): SWEEPS / "short_context_prefill_sweep.csv",
|
|
||||||
("prefill", "composite"): SWEEPS / "short_context_prefill_composite_sweep.csv",
|
|
||||||
("prefill", "composite_fused"): SWEEPS / "short_context_prefill_composite_fused_sweep.csv",
|
|
||||||
}
|
|
||||||
|
|
||||||
|
|
||||||
def load(phase, variant):
|
|
||||||
"""{(mode, S_kv): row dict}."""
|
|
||||||
rows = {}
|
|
||||||
with CSV_MAP[(phase, variant)].open() as f:
|
|
||||||
for r in csv.DictReader(f):
|
|
||||||
rows[(r["mode"], int(r["S_kv"]))] = r
|
|
||||||
return rows
|
|
||||||
|
|
||||||
|
|
||||||
def _grouped_bars(ax, x_labels, groups, group_color, ylabel, title,
|
|
||||||
log=False, value_fmt="{:.3g}"):
|
|
||||||
n = len(groups)
|
|
||||||
nx = len(x_labels)
|
|
||||||
width = 0.8 / n
|
|
||||||
x = np.arange(nx)
|
|
||||||
for i, (label, values) in enumerate(groups):
|
|
||||||
offset = (i - (n - 1) / 2) * width
|
|
||||||
bars = ax.bar(x + offset, values, width, label=label,
|
|
||||||
color=group_color[label], edgecolor="black",
|
|
||||||
linewidth=0.4)
|
|
||||||
for bar, v in zip(bars, values):
|
|
||||||
if v <= 0 and log:
|
|
||||||
continue
|
|
||||||
ax.text(bar.get_x() + bar.get_width() / 2,
|
|
||||||
bar.get_height(),
|
|
||||||
value_fmt.format(v), ha="center", va="bottom",
|
|
||||||
fontsize=7)
|
|
||||||
ax.set_xticks(x)
|
|
||||||
ax.set_xticklabels(x_labels)
|
|
||||||
ax.set_xlabel("S_kv")
|
|
||||||
ax.set_ylabel(ylabel)
|
|
||||||
ax.set_title(title)
|
|
||||||
if log:
|
|
||||||
ax.set_yscale("log")
|
|
||||||
ax.grid(True, axis="y", alpha=0.3, which="both")
|
|
||||||
ax.legend(fontsize=8, loc="upper left")
|
|
||||||
|
|
||||||
|
|
||||||
# ── B: per-cube trade-off (mode comparison, variant 1 baseline) ──
|
|
||||||
|
|
||||||
def fig_b_per_cube_tradeoff():
|
|
||||||
"""3 metrics × 2 phases — mode comparison.
|
|
||||||
|
|
||||||
Compute path doesn't affect HBM/IPCQ/KV cost (variant-invariant);
|
|
||||||
use variant 1 baseline rows.
|
|
||||||
"""
|
|
||||||
fig, axes = plt.subplots(2, 2, figsize=(13, 9))
|
|
||||||
|
|
||||||
METRICS = [
|
|
||||||
("hbm_bw_util", "HBM BW utilization (per PE)", False, "{:.3f}"),
|
|
||||||
("ipcq_kb", "IPCQ traffic (KB)", True, "{:.1f}"),
|
|
||||||
]
|
|
||||||
for row, (key, ylabel, log, fmt) in enumerate(METRICS):
|
|
||||||
for col, phase in enumerate(("prefill", "decode")):
|
|
||||||
ax = axes[row][col]
|
|
||||||
rows = load(phase, "baseline")
|
|
||||||
groups = []
|
|
||||||
for mode in MODES:
|
|
||||||
values = [float(rows[(mode, s)][key]) for s in CONTEXTS]
|
|
||||||
groups.append((MODE_LABEL[mode], values))
|
|
||||||
_grouped_bars(ax, CONTEXT_LABELS, groups, MODE_COLOR,
|
|
||||||
ylabel=ylabel,
|
|
||||||
title=f"{phase.capitalize()}",
|
|
||||||
log=log, value_fmt=fmt)
|
|
||||||
|
|
||||||
fig.suptitle(
|
|
||||||
"Per-cube trade-off: HBM BW + IPCQ mode comparison",
|
|
||||||
fontsize=12, fontweight="bold")
|
|
||||||
fig.tight_layout()
|
|
||||||
out = OUT / "per_cube_tradeoff_mode_compare.png"
|
|
||||||
fig.savefig(out, dpi=140, bbox_inches="tight")
|
|
||||||
plt.close(fig)
|
|
||||||
print(f" ✓ {out}")
|
|
||||||
|
|
||||||
|
|
||||||
# ── C: GEMM util ablation (A1, variant compare) ──
|
|
||||||
|
|
||||||
def fig_c_gemm_util_ablation():
|
|
||||||
"""A1 only; x=S_kv, grouped bars per variant; 2 panels (prefill/decode).
|
|
||||||
|
|
||||||
NOTE annotation: gemm_util gap is largely supertile padding overhead
|
|
||||||
(M=8 < TILE_M=32), not pure fusion gain. See ADR-0070 limitation #4.
|
|
||||||
"""
|
|
||||||
fig, axes = plt.subplots(1, 2, figsize=(13, 5.4))
|
|
||||||
for ax, phase in zip(axes, ("prefill", "decode")):
|
|
||||||
groups = []
|
|
||||||
for variant_key, _ in VARIANTS:
|
|
||||||
rows = load(phase, variant_key)
|
|
||||||
values = [float(rows[("A1", s)]["gemm_util"]) for s in CONTEXTS]
|
|
||||||
groups.append((variant_key, values))
|
|
||||||
_grouped_bars(ax, CONTEXT_LABELS, groups, VARIANT_COLOR,
|
|
||||||
ylabel="GEMM engine utilization (per PE)",
|
|
||||||
title=f"{phase.capitalize()}",
|
|
||||||
log=False, value_fmt="{:.4f}")
|
|
||||||
# Replace legend with readable labels
|
|
||||||
ax.legend([v[1] for v in VARIANTS], fontsize=8, loc="upper left")
|
|
||||||
|
|
||||||
fig.suptitle("1-kv-per-cube: GEMM utilization variant comparison",
|
|
||||||
fontsize=12, fontweight="bold")
|
|
||||||
fig.text(
|
|
||||||
0.5, -0.04,
|
|
||||||
"⚠ Higher gemm_util on composite/fused is largely supertile "
|
|
||||||
"padding overhead (M=G=8 padded to TILE_M=32), not pure fusion "
|
|
||||||
"gain. See ADR-0070 limitation #4.",
|
|
||||||
ha="center", fontsize=9, style="italic", color="#555")
|
|
||||||
fig.tight_layout()
|
|
||||||
out = OUT / "gemm_util_a1_variant_ablation.png"
|
|
||||||
fig.savefig(out, dpi=140, bbox_inches="tight")
|
|
||||||
plt.close(fig)
|
|
||||||
print(f" ✓ {out}")
|
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
|
||||||
print(f"Output: {OUT}")
|
|
||||||
fig_b_per_cube_tradeoff()
|
|
||||||
fig_c_gemm_util_ablation()
|
|
||||||
print("\nDone.")
|
|
||||||
@@ -1,156 +0,0 @@
|
|||||||
"""Generate 4 wall-clock comparison figures for the short-context GQA
|
|
||||||
attention sweep (prefill + decode × A1/A2/A4/B × 3 variants × 4 contexts).
|
|
||||||
|
|
||||||
Figures 1-3: mode comparison within one variant
|
|
||||||
• 2 panels (prefill, decode), grouped bars per S_kv, hue=mode (A1/A2/A4/B)
|
|
||||||
Figure 4: A1 variant comparison (best mode)
|
|
||||||
• 2 panels, grouped bars per S_kv, hue=variant
|
|
||||||
|
|
||||||
Output: docs/report/1H-codesign-paper/figures/gqa_short_context/
|
|
||||||
"""
|
|
||||||
from __future__ import annotations
|
|
||||||
|
|
||||||
import csv
|
|
||||||
from pathlib import Path
|
|
||||||
|
|
||||||
import matplotlib.pyplot as plt
|
|
||||||
import numpy as np
|
|
||||||
|
|
||||||
ROOT = Path("/Users/lalicorne/Desktop/ahbm/kernbench2")
|
|
||||||
SWEEPS = ROOT / "docs" / "sweeps"
|
|
||||||
OUT = ROOT / "docs" / "report" / "1H-codesign-paper" / "figures" / "gqa_short_context"
|
|
||||||
OUT.mkdir(parents=True, exist_ok=True)
|
|
||||||
|
|
||||||
CONTEXTS = [8192, 16384, 32768, 65536]
|
|
||||||
CONTEXT_LABELS = ["8K", "16K", "32K", "64K"]
|
|
||||||
MODES = ["A1", "A2", "A4", "B"]
|
|
||||||
MODE_LABEL = {"A1": "1-kv-per-cube", "A2": "2-kv-per-cube",
|
|
||||||
"A4": "4-kv-per-cube", "B": "8-kv-per-cube"}
|
|
||||||
MODE_COLOR = {MODE_LABEL["A1"]: "#1f77b4",
|
|
||||||
MODE_LABEL["A2"]: "#2ca02c",
|
|
||||||
MODE_LABEL["A4"]: "#ffd000",
|
|
||||||
MODE_LABEL["B"]: "#d62728"}
|
|
||||||
VARIANTS = [
|
|
||||||
("baseline", "(1) without composite"),
|
|
||||||
("composite", "(2) with composite (GEMM-only)"),
|
|
||||||
("composite_fused", "(3) with composite + softmax_merge"),
|
|
||||||
]
|
|
||||||
VARIANT_COLOR = {
|
|
||||||
"baseline": "#1f77b4",
|
|
||||||
"composite": "#2ca02c",
|
|
||||||
"composite_fused": "#d62728",
|
|
||||||
}
|
|
||||||
|
|
||||||
CSV_MAP = {
|
|
||||||
("decode", "baseline"): SWEEPS / "short_context_decode_sweep.csv",
|
|
||||||
("decode", "composite"): SWEEPS / "short_context_decode_composite_sweep.csv",
|
|
||||||
("decode", "composite_fused"): SWEEPS / "short_context_decode_composite_fused_sweep.csv",
|
|
||||||
("prefill", "baseline"): SWEEPS / "short_context_prefill_sweep.csv",
|
|
||||||
("prefill", "composite"): SWEEPS / "short_context_prefill_composite_sweep.csv",
|
|
||||||
("prefill", "composite_fused"): SWEEPS / "short_context_prefill_composite_fused_sweep.csv",
|
|
||||||
}
|
|
||||||
|
|
||||||
|
|
||||||
def load(phase, variant):
|
|
||||||
"""Return {(mode, S_kv): wall_us}."""
|
|
||||||
rows = {}
|
|
||||||
with CSV_MAP[(phase, variant)].open() as f:
|
|
||||||
for r in csv.DictReader(f):
|
|
||||||
rows[(r["mode"], int(r["S_kv"]))] = float(r["wall_us"])
|
|
||||||
return rows
|
|
||||||
|
|
||||||
|
|
||||||
def _bar_chart(ax, *, x_labels, groups, group_color, ylabel, title,
|
|
||||||
y_log=True):
|
|
||||||
"""Grouped bars: x = x_labels, groups = list of (label, values).
|
|
||||||
|
|
||||||
values is a list aligned with x_labels.
|
|
||||||
"""
|
|
||||||
n_groups = len(groups)
|
|
||||||
n_x = len(x_labels)
|
|
||||||
width = 0.8 / n_groups
|
|
||||||
x = np.arange(n_x)
|
|
||||||
for i, (label, values) in enumerate(groups):
|
|
||||||
offset = (i - (n_groups - 1) / 2) * width
|
|
||||||
bars = ax.bar(x + offset, values, width, label=label,
|
|
||||||
color=group_color[label], edgecolor="black",
|
|
||||||
linewidth=0.4)
|
|
||||||
for bar, v in zip(bars, values):
|
|
||||||
ax.text(bar.get_x() + bar.get_width() / 2,
|
|
||||||
bar.get_height(),
|
|
||||||
f"{v:.0f}", ha="center", va="bottom", fontsize=7,
|
|
||||||
rotation=0)
|
|
||||||
ax.set_xticks(x)
|
|
||||||
ax.set_xticklabels(x_labels)
|
|
||||||
ax.set_xlabel("S_kv")
|
|
||||||
ax.set_ylabel(ylabel)
|
|
||||||
ax.set_title(title)
|
|
||||||
if y_log:
|
|
||||||
ax.set_yscale("log")
|
|
||||||
ax.grid(True, axis="y", alpha=0.3, which="both")
|
|
||||||
ax.legend(fontsize=8, loc="upper left")
|
|
||||||
|
|
||||||
|
|
||||||
def _fig_mode_compare(variant_key, variant_label, fname):
|
|
||||||
"""Figure: 2 panels (prefill/decode), x=S_kv, grouped bars per mode."""
|
|
||||||
fig, axes = plt.subplots(1, 2, figsize=(13, 4.8))
|
|
||||||
for ax, phase in zip(axes, ("prefill", "decode")):
|
|
||||||
rows = load(phase, variant_key)
|
|
||||||
groups = []
|
|
||||||
for mode in MODES:
|
|
||||||
values = [rows[(mode, s)] for s in CONTEXTS]
|
|
||||||
groups.append((MODE_LABEL[mode], values))
|
|
||||||
_bar_chart(ax, x_labels=CONTEXT_LABELS, groups=groups,
|
|
||||||
group_color=MODE_COLOR,
|
|
||||||
ylabel="Wall-clock latency (μs)",
|
|
||||||
title=f"{phase.capitalize()}")
|
|
||||||
fig.suptitle(variant_label,
|
|
||||||
fontsize=12, fontweight="bold")
|
|
||||||
fig.tight_layout()
|
|
||||||
out = OUT / fname
|
|
||||||
fig.savefig(out, dpi=140, bbox_inches="tight")
|
|
||||||
plt.close(fig)
|
|
||||||
print(f" ✓ {out}")
|
|
||||||
|
|
||||||
|
|
||||||
def _fig_a1_variant_compare(fname):
|
|
||||||
"""Figure: A1 (best mode), 2 panels, x=S_kv, grouped bars per variant."""
|
|
||||||
fig, axes = plt.subplots(1, 2, figsize=(13, 4.8))
|
|
||||||
for ax, phase in zip(axes, ("prefill", "decode")):
|
|
||||||
groups = []
|
|
||||||
for variant_key, variant_label in VARIANTS:
|
|
||||||
rows = load(phase, variant_key)
|
|
||||||
values = [rows[("A1", s)] for s in CONTEXTS]
|
|
||||||
groups.append((variant_key, values))
|
|
||||||
_bar_chart(
|
|
||||||
ax, x_labels=CONTEXT_LABELS, groups=groups,
|
|
||||||
group_color=VARIANT_COLOR,
|
|
||||||
ylabel="Wall-clock latency (μs)",
|
|
||||||
title=f"{phase.capitalize()}",
|
|
||||||
)
|
|
||||||
# Custom legend labels
|
|
||||||
handles, _ = ax.get_legend_handles_labels()
|
|
||||||
readable = [v[1] for v in VARIANTS]
|
|
||||||
ax.legend(handles, readable, fontsize=8, loc="upper left")
|
|
||||||
fig.suptitle("1-kv-per-cube: variant comparison",
|
|
||||||
fontsize=12, fontweight="bold")
|
|
||||||
fig.tight_layout()
|
|
||||||
out = OUT / fname
|
|
||||||
fig.savefig(out, dpi=140, bbox_inches="tight")
|
|
||||||
plt.close(fig)
|
|
||||||
print(f" ✓ {out}")
|
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
|
||||||
print(f"Output dir: {OUT}")
|
|
||||||
_fig_mode_compare("baseline",
|
|
||||||
"(1) Without composite: latency mode comparison",
|
|
||||||
"wall_variant1_mode_compare.png")
|
|
||||||
_fig_mode_compare("composite",
|
|
||||||
"(2) With composite (GEMM-only): latency mode comparison",
|
|
||||||
"wall_variant2_mode_compare.png")
|
|
||||||
_fig_mode_compare("composite_fused",
|
|
||||||
"(3) With composite + softmax_merge: latency mode comparison",
|
|
||||||
"wall_variant3_mode_compare.png")
|
|
||||||
_fig_a1_variant_compare("wall_a1_variant_compare.png")
|
|
||||||
print("\nDone.")
|
|
||||||
Reference in New Issue
Block a user