diff --git a/docs/sweeps/plot_scripts/plot_short_context.py b/docs/sweeps/plot_scripts/plot_short_context.py new file mode 100644 index 0000000..6ae61bb --- /dev/null +++ b/docs/sweeps/plot_scripts/plot_short_context.py @@ -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.") diff --git a/docs/sweeps/plot_scripts/plot_short_context_extra.py b/docs/sweeps/plot_scripts/plot_short_context_extra.py deleted file mode 100644 index 0c0b997..0000000 --- a/docs/sweeps/plot_scripts/plot_short_context_extra.py +++ /dev/null @@ -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.") diff --git a/docs/sweeps/plot_scripts/plot_short_context_wall.py b/docs/sweeps/plot_scripts/plot_short_context_wall.py deleted file mode 100644 index ef6f9c8..0000000 --- a/docs/sweeps/plot_scripts/plot_short_context_wall.py +++ /dev/null @@ -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.")