sweep: 6 short-context CSVs (3 variants x prefill+decode) + 6 wall/util figures

docs/sweeps/plot_scripts/plot_short_context_extra.py

@@ -0,0 +1,165 @@
+"""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.")

docs/sweeps/plot_scripts/plot_short_context_wall.py

@@ -0,0 +1,156 @@
+"""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.")

docs/sweeps/short_context_decode_composite_fused_sweep.csv

@@ -0,0 +1,17 @@
+variant,mode,kv_per_cube,C,S_kv,wall_us,n_pe,gemm_util,math_count,math_pipeline_us,hbm_bw_util,hbm_read_mb,hbm_write_kb,ipcq_kb,kv_cache_per_cube_mb
+composite_fused,A1,1,8,8192,146.95153450000072,64,0.0017838806576053131,656,0.0,0.013997812319543966,16.0625,8.0,57.75,2.0
+composite_fused,A1,1,8,16384,279.867342500003,64,0.007173570099630665,1296,13.903615999992937,0.0146676634841736,32.0625,8.0,57.75,4.0
+composite_fused,A1,1,8,32768,544.5763825000042,64,0.010097125355742134,2576,41.71084799996763,0.015059411798858053,64.0625,8.0,57.75,8.0
+composite_fused,A1,1,8,65536,1074.101862500392,64,0.011619635377527878,5136,97.32531200143696,0.015262053416273654,128.0625,8.0,57.75,16.0
+composite_fused,A2,2,4,8192,140.35490875000028,32,0.01430408111743636,624,6.951807999995537,0.029254409671653128,16.03125,8.0,24.75,4.0
+composite_fused,A2,2,4,16384,273.9919487500021,32,0.020068677291975947,1264,20.855423999989405,0.029935186188568382,32.03125,8.0,24.75,8.0
+composite_fused,A2,2,4,32768,541.3734287504049,32,0.0230537210304644,2544,48.66265599996224,0.030282239817052973,64.03125,8.0,24.75,16.0
+composite_fused,A2,2,4,65536,1076.851308751879,32,0.024557433125924955,5104,104.27712000153959,0.030438742780552714,128.03125,8.0,24.75,32.0
+composite_fused,A4,4,2,8192,140.2291675000001,16,0.03921192786066565,608,10.427711999993306,0.058504233792873325,16.015625,8.0,8.25,8.0
+composite_fused,A4,4,2,16384,276.4849475001752,16,0.04514051167338893,1248,24.33132799998764,0.05930160085835996,32.015625,8.0,8.25,16.0
+composite_fused,A4,4,2,32768,549.3539675010354,16,0.04813782290440163,2528,52.13855999995954,0.059670088757369746,64.015625,8.0,8.25,32.0
+composite_fused,A4,4,2,65536,1095.4730075019588,16,0.04963405545386165,5088,107.7530240015909,0.05983534012350623,128.015625,8.0,8.25,64.0
+composite_fused,B,8,1,8192,144.10828500005923,8,0.08660620726926088,600,12.16566399999219,0.11380331116974475,16.0078125,8.0,0.0,16.0
+composite_fused,B,8,1,16384,285.6728750004253,8,0.09256988084807763,1240,26.069279999986758,0.11476063311909189,32.0078125,8.0,0.0,32.0
+composite_fused,B,8,1,32768,568.9925550010347,8,0.09555971782205692,2520,53.8765119999582,0.1152071313127828,64.0078125,8.0,0.0,64.0
+composite_fused,B,8,1,65536,1135.6319150019874,8,0.09706392938508733,5080,109.49097600161657,0.11543176822374755,128.0078125,8.0,0.0,128.0

docs/sweeps/short_context_decode_composite_sweep.csv

@@ -0,0 +1,17 @@
+variant,mode,kv_per_cube,C,S_kv,wall_us,n_pe,gemm_util,math_count,math_pipeline_us,hbm_bw_util,hbm_read_mb,hbm_write_kb,ipcq_kb,kv_cache_per_cube_mb
+composite,A1,1,8,8192,147.23658450000076,64,0.0057446592018384314,656,0.0,0.013970712557516501,16.0625,8.0,57.75,2.0
+composite,A1,1,8,16384,279.9533665000026,64,0.006042606385378475,1360,0.0,0.014663156408229735,32.0625,8.0,57.75,4.0
+composite,A1,1,8,32768,545.3869305000062,64,0.006203478321011016,2768,0.0,0.015037030668265908,64.0625,8.0,57.75,8.0
+composite,A1,1,8,65536,1076.3614585003581,64,0.006286542449928883,5584,0.0,0.015230013923798023,128.0625,8.0,57.75,16.0
+composite,A2,2,4,8192,140.44093275000077,32,0.012045263206857853,656,0.0,0.02923649052736712,16.03125,8.0,24.75,4.0
+composite,A2,2,4,16384,274.80249675000204,32,0.012311736756599198,1360,0.0,0.029846890392199246,32.03125,8.0,24.75,8.0
+composite,A2,2,4,32768,543.6330247504077,32,0.012446984807344498,2768,0.0,0.030156372504276757,64.03125,8.0,24.75,16.0
+composite,A2,2,4,65536,1082.0090007518106,32,0.012507459726542472,5584,0.0,0.030293648183356066,128.03125,8.0,24.75,32.0
+composite,A4,4,2,8192,141.0397155000009,16,0.023988250316491475,656,0.0,0.05816801296653173,16.015625,8.0,8.25,8.0
+composite,A4,4,2,16384,278.7445435001758,16,0.02427524469220757,1360,0.0,0.05882088235384476,32.015625,8.0,8.25,16.0
+composite,A4,4,2,32768,554.51165950104,16,0.024405589617097562,2768,0.0,0.059115077994024615,64.015625,8.0,8.25,32.0
+composite,A4,4,2,65536,1106.4268915018213,16,0.024462861675913012,5584,0.0,0.05924295631591859,128.015625,8.0,8.25,64.0
+composite,B,8,1,8192,146.3478710000607,8,0.0462363541999099,656,0.0,0.11206176002378058,16.0078125,8.0,0.0,16.0
+composite,B,8,1,16384,290.8105570004271,8,0.046536082251474285,1360,0.0,0.1127331838917796,32.0078125,8.0,0.0,32.0
+composite,B,8,1,32768,579.9264290010428,8,0.046672071914328234,2768,0.0,0.11303502775846438,64.0078125,8.0,0.0,64.0
+composite,B,8,1,65536,1158.158173001712,8,0.04674036523209764,5584,0.0,0.11318661220534877,128.0078125,8.0,0.0,128.0

docs/sweeps/short_context_decode_sweep.4ctx.csv

@@ -0,0 +1,17 @@
+mode,kv_per_cube,C,S_kv,wall_us,n_pe,gemm_util,math_count,hbm_bw_util,hbm_read_mb,hbm_write_kb,ipcq_kb,kv_cache_per_cube_mb
+A1,1,8,8192,149.2685345000007,64,0.001756190618994587,656,0.013780533230866485,16.0625,8.0,57.75,2.0
+A1,1,8,16384,282.2252665000025,64,0.0018576933472394743,1360,0.014545118695104373,32.0625,8.0,57.75,4.0
+A1,1,8,32768,548.138730500006,64,0.0019129755692392889,2768,0.014961540835691614,64.0625,8.0,57.75,8.0
+A1,1,8,65536,1080.073058500367,64,0.0019416760593127277,5584,0.01517767698303756,128.0625,8.0,57.75,16.0
+A2,2,4,8192,142.16883275000066,32,0.0036877843748065785,656,0.028881154333033524,16.03125,8.0,24.75,4.0
+A2,2,4,16384,277.0102967500018,32,0.0037853322143696724,1360,0.029609007665885423,32.03125,8.0,24.75,8.0
+A2,2,4,32768,546.8006247504073,32,0.0038353138330044275,2768,0.02998167752182655,64.03125,8.0,24.75,16.0
+A2,2,4,65536,1087.0962007518285,32,0.003858263874988177,5584,0.03015188534127058,128.03125,8.0,24.75,32.0
+A4,4,2,8192,142.21417550000066,16,0.007373217165681769,656,0.05768763888097753,16.015625,8.0,8.25,8.0
+A4,4,2,16384,280.8788035001758,16,0.007466394665122732,1360,0.058373931374247456,32.015625,8.0,8.25,16.0
+A4,4,2,32768,558.5655195010398,16,0.007509063580845671,2768,0.058686042828569145,64.015625,8.0,8.25,32.0
+A4,4,2,65536,1114.319951501857,16,0.007528006645388841,5584,0.05882332081702009,128.015625,8.0,8.25,64.0
+B,8,1,8192,147.93485600006025,8,0.014176185766516818,656,0.11085960701508589,16.0078125,8.0,0.0,16.0
+B,8,1,16384,294.3171420004266,8,0.014250967413897551,1360,0.1113900460475132,32.0078125,8.0,0.0,32.0
+B,8,1,32768,587.2722140010417,8,0.01428401991446495,2768,0.11162115018757507,64.0078125,8.0,0.0,64.0
+B,8,1,65536,1173.1823580017838,8,0.014300603725891686,5584,0.11173710472707321,128.0078125,8.0,0.0,128.0

docs/sweeps/short_context_decode_sweep.csv

@@ -1,17 +1,17 @@
 mode,kv_per_cube,C,S_kv,wall_us,n_pe,gemm_util,math_count,hbm_bw_util,hbm_read_mb,hbm_write_kb,ipcq_kb,kv_cache_per_cube_mb
-A1,1,8,8192,149.2685345000007,64,0.001756190618994587,656,0.013780533230866485,16.0625,8.0,57.75,2.0
-A1,1,8,16384,282.2252665000025,64,0.0018576933472394743,1360,0.014545118695104373,32.0625,8.0,57.75,4.0
-A1,1,8,32768,548.138730500006,64,0.0019129755692392889,2768,0.014961540835691614,64.0625,8.0,57.75,8.0
-A1,1,8,65536,1080.073058500367,64,0.0019416760593127277,5584,0.01517767698303756,128.0625,8.0,57.75,16.0
-A2,2,4,8192,142.16883275000066,32,0.0036877843748065785,656,0.028881154333033524,16.03125,8.0,24.75,4.0
-A2,2,4,16384,277.0102967500018,32,0.0037853322143696724,1360,0.029609007665885423,32.03125,8.0,24.75,8.0
-A2,2,4,32768,546.8006247504073,32,0.0038353138330044275,2768,0.02998167752182655,64.03125,8.0,24.75,16.0
-A2,2,4,65536,1087.0962007518285,32,0.003858263874988177,5584,0.03015188534127058,128.03125,8.0,24.75,32.0
-A4,4,2,8192,142.21417550000066,16,0.007373217165681769,656,0.05768763888097753,16.015625,8.0,8.25,8.0
-A4,4,2,16384,280.8788035001758,16,0.007466394665122732,1360,0.058373931374247456,32.015625,8.0,8.25,16.0
-A4,4,2,32768,558.5655195010398,16,0.007509063580845671,2768,0.058686042828569145,64.015625,8.0,8.25,32.0
-A4,4,2,65536,1114.319951501857,16,0.007528006645388841,5584,0.05882332081702009,128.015625,8.0,8.25,64.0
-B,8,1,8192,147.93485600006025,8,0.014176185766516818,656,0.11085960701508589,16.0078125,8.0,0.0,16.0
-B,8,1,16384,294.3171420004266,8,0.014250967413897551,1360,0.1113900460475132,32.0078125,8.0,0.0,32.0
-B,8,1,32768,587.2722140010417,8,0.01428401991446495,2768,0.11162115018757507,64.0078125,8.0,0.0,64.0
-B,8,1,65536,1173.1823580017838,8,0.014300603725891686,5584,0.11173710472707321,128.0078125,8.0,0.0,128.0
+A1,1,8,8192,146.95153450000072,64,0.0017838806576053131,656,0.013997812319543966,16.0625,8.0,57.75,2.0
+A1,1,8,16384,279.3832665000025,64,0.0018765905580817512,1360,0.014693077546933053,32.0625,8.0,57.75,4.0
+A1,1,8,32768,544.246730500006,64,0.0019266555796063503,2768,0.015068533333155061,64.0625,8.0,57.75,8.0
+A1,1,8,65536,1074.081058500367,64,0.0019525081309290412,5584,0.01526234902874828,128.0625,8.0,57.75,16.0
+A2,2,4,8192,139.87083275000066,32,0.003748372621310069,656,0.029355655637933423,16.03125,8.0,24.75,4.0
+A2,2,4,16384,273.6622967500018,32,0.0038316421825465713,1360,0.029971245938539923,32.03125,8.0,24.75,8.0
+A2,2,4,32768,541.3526247504072,32,0.0038739112070761647,2768,0.030283403553383558,64.03125,8.0,24.75,16.0
+A2,2,4,65536,1077.4482007518286,32,0.003892812663356761,5584,0.030421880121130614,128.03125,8.0,24.75,32.0
+A4,4,2,8192,139.89951550000075,16,0.007495208230366452,656,0.05864209015076937,16.015625,8.0,8.25,8.0
+A4,4,2,16384,276.4641435001754,16,0.007585620230706115,1360,0.05930606331952627,32.015625,8.0,8.25,16.0
+A4,4,2,32768,549.9508595010393,16,0.007626688689616412,2768,0.059605325518975856,64.015625,8.0,8.25,32.0
+A4,4,2,65536,1097.305291501857,16,0.0076447348472311214,5584,0.059735426874945555,128.015625,8.0,8.25,64.0
+B,8,1,8192,144.06747100006032,8,0.014556735017573057,656,0.11383555139933797,16.0078125,8.0,0.0,16.0
+B,8,1,16384,286.24975700042637,8,0.014652602831705778,1360,0.1145293548666704,32.0078125,8.0,0.0,32.0
+B,8,1,32768,570.8048290010414,8,0.014696105522939606,2768,0.11484135499470417,64.0078125,8.0,0.0,64.0
+B,8,1,65536,1139.9149730017834,8,0.014717953880200899,5584,0.11499805082373886,128.0078125,8.0,0.0,128.0

docs/sweeps/short_context_prefill_composite_fused_sweep.csv

@@ -0,0 +1,17 @@
+variant,mode,kv_per_cube,C,S_kv,wall_us,n_pe,gemm_util,math_count,math_pipeline_us,hbm_bw_util,hbm_read_mb,hbm_write_kb,ipcq_kb,kv_cache_per_cube_mb
+composite_fused,A1,1,8,8192,1101.4436329998505,64,0.011331194467560951,4800,97.32531200143696,0.0018739043362378583,16.0625,64.0,114688.0,2.0
+composite_fused,A1,1,8,16384,2178.44091796638,64,0.012139279878164511,9920,208.55424000307917,0.0018875884886695197,32.0625,64.0,229376.0,4.0
+composite_fused,A1,1,8,32768,4335.060767847123,64,0.012542561892233649,20160,431.01209600543973,0.0018933990639481287,64.0625,64.0,458752.0,8.0
+composite_fused,A1,1,8,65536,8649.804197613514,64,0.012743513431742019,40640,875.9278079059123,0.0018959966752224023,128.0625,64.0,917504.0,16.0
+composite_fused,A2,2,4,8192,1093.9720835019473,32,0.01853321150177446,2400,96.40870400077105,0.0037734052470386023,16.0625,64.0,49152.0,4.0
+composite_fused,A2,2,4,16384,2175.4760434340433,32,0.019695389489835093,4960,206.59008000165224,0.00378032202414797,32.0625,64.0,98304.0,8.0
+composite_fused,A2,2,4,32768,4338.482933349601,32,0.020281493171481216,10080,426.9528320015669,0.003783811127574436,64.0625,64.0,196608.0,16.0
+composite_fused,A2,2,4,65536,8664.499802991075,32,0.020575802877008553,20320,867.678335950613,0.0037855618611332994,128.0625,64.0,393216.0,32.0
+composite_fused,A4,4,2,8192,1102.9720360028073,16,0.03256574312686088,1200,92.28038400042057,0.007485230568419403,16.0625,64.0,16384.0,8.0
+composite_fused,A4,4,2,16384,2196.357595426575,16,0.03449854256839416,2480,197.74368000090124,0.007488762319145705,32.0625,64.0,32768.0,16.0
+composite_fused,A4,4,2,32768,4383.130075368471,16,0.03547127219906801,5040,408.6702720000148,0.00749053745507198,64.0625,64.0,65536.0,32.0
+composite_fused,A4,4,2,65536,8756.675034884938,16,0.03595923415220853,10160,830.5234559737444,0.007491427937962982,128.0625,64.0,131072.0,64.0
+composite_fused,B,8,1,8192,1125.1019960018205,8,0.12366810519999533,712,121.6138560003899,0.014676002761240575,16.0625,64.0,0.0,16.0
+composite_fused,B,8,1,16384,2241.8354359431773,8,0.13192705729555512,1480,260.6011200008355,0.0146736907948643,32.0625,64.0,0.0,32.0
+composite_fused,B,8,1,32768,4475.302315825972,8,0.1360796989728368,3016,538.575647996068,0.014672528326811126,64.0625,64.0,0.0,64.0
+composite_fused,B,8,1,65536,8942.236075603707,8,0.13816184916758104,6088,1094.5247039788662,0.014671945460927953,128.0625,64.0,0.0,128.0

docs/sweeps/short_context_prefill_composite_sweep.csv

@@ -0,0 +1,17 @@
+variant,mode,kv_per_cube,C,S_kv,wall_us,n_pe,gemm_util,math_count,math_pipeline_us,hbm_bw_util,hbm_read_mb,hbm_write_kb,ipcq_kb,kv_cache_per_cube_mb
+composite,A1,1,8,8192,1101.224596999844,64,0.006144606666764116,5248,0.0,0.0018742770599413814,16.0625,64.0,114688.0,2.0
+composite,A1,1,8,16384,2178.2218819663735,64,0.006212950164859459,10880,0.0,0.0018877782993750494,32.0625,64.0,229376.0,4.0
+composite,A1,1,8,32768,4334.841731847132,64,0.006243911466154239,22144,0.0,0.0018934947358510517,64.0625,64.0,458752.0,8.0
+composite,A1,1,8,65536,8649.585161613555,64,0.006258419909799432,44672,0.0,0.0018960446881062475,128.0625,64.0,917504.0,16.0
+composite,A2,2,4,8192,1094.5704665014161,32,0.007139940496760185,2624,0.0,0.0037713423907684606,16.0625,64.0,49152.0,4.0
+composite,A2,2,4,16384,2177.5659939366815,32,0.007177893135813299,5440,0.0,0.0037766938053309506,32.0625,64.0,98304.0,8.0
+composite,A2,2,4,32768,4341.645937350417,32,0.007200189155116074,11072,0.0,0.0037810545209999826,64.0625,64.0,196608.0,16.0
+composite,A2,2,4,65536,8671.42409207618,32,0.007210043393449057,22336,0.0,0.003782539021470782,128.0625,64.0,393216.0,32.0
+composite,A4,4,2,8192,1104.2310135028845,16,0.008976672344028575,1312,0.0,0.007476696360673657,16.0625,64.0,16384.0,8.0
+composite,A4,4,2,16384,2197.6415934265688,16,0.009020870400324496,2720,0.0,0.0074843869215062645,32.0625,64.0,32768.0,16.0
+composite,A4,4,2,32768,4384.414073368468,16,0.00904323344692278,5536,0.0,0.007488343813013936,64.0625,64.0,65536.0,32.0
+composite,A4,4,2,65536,8757.965852885032,16,0.009054449551299243,11168,0.0,0.0074903237922982055,128.0625,64.0,131072.0,64.0
+composite,B,8,1,8192,1125.5021760018114,8,0.020408994750147753,656,0.0,0.014670784608037422,16.0625,64.0,0.0,16.0
+composite,B,8,1,16384,2242.2356159431683,8,0.02048880843562789,1360,0.0,0.014671071927542597,32.0625,64.0,0.0,32.0
+composite,B,8,1,32768,4475.702495825993,8,0.020528949826292484,2768,0.0,0.014671216431663579,64.0625,64.0,0.0,64.0
+composite,B,8,1,65536,8942.63625560376,8,0.020549079566244906,5584,0.0,0.014671288896245289,128.0625,64.0,0.0,128.0

docs/sweeps/short_context_prefill_sweep.4ctx.csv

@@ -0,0 +1,17 @@
+mode,kv_per_cube,C,S_kv,wall_us,n_pe,gemm_util,math_count,hbm_bw_util,hbm_read_mb,hbm_write_kb,ipcq_kb,kv_cache_per_cube_mb
+A1,1,8,8192,1103.359903999566,64,0.0019006962210579498,5248,0.001870649814732448,16.0625,64.0,114688.0,2.0
+A1,1,8,16384,2182.457823968016,64,0.0019218259129421755,10880,0.0018841143021603984,32.0625,64.0,229376.0,4.0
+A1,1,8,32768,4343.2822238495155,64,0.0019313983221991736,22144,0.0018898150239762977,64.0625,64.0,458752.0,8.0
+A1,1,8,65536,8666.440783650422,64,0.001935883071136998,44672,0.0018923570136126979,128.0625,64.0,917504.0,16.0
+A2,2,4,8192,1097.2535885006619,32,0.003822547535007557,2624,0.0037621203004135906,16.0625,64.0,49152.0,4.0
+A2,2,4,16384,2179.78425743804,32,0.003848366172648072,5440,0.003772850442394648,32.0625,64.0,98304.0,8.0
+A2,2,4,32768,4344.834905351148,32,0.003861416225357031,11072,0.003778279349528763,64.0625,64.0,196608.0,16.0
+A2,2,4,65536,8674.946611121853,32,0.003867969856656271,22336,0.003781003096658628,128.0625,64.0,393216.0,32.0
+A4,4,2,8192,1103.675383502932,16,0.007600611670231015,1312,0.007480460399321815,16.0625,64.0,16384.0,8.0
+A4,4,2,16384,2197.0616234262543,16,0.00763620638634706,2720,0.007486362614786297,32.0625,64.0,32768.0,16.0
+A4,4,2,32768,4383.834103368879,16,0.007654129058897454,5536,0.007489334501679554,64.0625,64.0,65536.0,32.0
+A4,4,2,65536,8757.379062883949,16,0.007663121981838932,11168,0.007490825682997995,128.0625,64.0,131072.0,64.0
+B,8,1,8192,1125.9457520018113,8,0.014900554462921552,656,0.01466500492643045,16.0625,64.0,0.0,16.0
+B,8,1,16384,2242.543191943168,8,0.014962669223294344,1360,0.014669059716747552,32.0625,64.0,0.0,32.0
+B,8,1,32768,4475.650071825993,8,0.014994216018466134,2768,0.014671388277951352,64.0625,64.0,0.0,64.0
+B,8,1,65536,8942.79183160376,8,0.015008481750139946,5584,0.014671033662702533,128.0625,64.0,0.0,128.0

docs/sweeps/short_context_prefill_sweep.csv

@@ -1,17 +1,17 @@
 mode,kv_per_cube,C,S_kv,wall_us,n_pe,gemm_util,math_count,hbm_bw_util,hbm_read_mb,hbm_write_kb,ipcq_kb,kv_cache_per_cube_mb
-A1,1,8,8192,1103.359903999566,64,0.0019006962210579498,5248,0.001870649814732448,16.0625,64.0,114688.0,2.0
-A1,1,8,16384,2182.457823968016,64,0.0019218259129421755,10880,0.0018841143021603984,32.0625,64.0,229376.0,4.0
-A1,1,8,32768,4343.2822238495155,64,0.0019313983221991736,22144,0.0018898150239762977,64.0625,64.0,458752.0,8.0
-A1,1,8,65536,8666.440783650422,64,0.001935883071136998,44672,0.0018923570136126979,128.0625,64.0,917504.0,16.0
-A2,2,4,8192,1097.2535885006619,32,0.003822547535007557,2624,0.0037621203004135906,16.0625,64.0,49152.0,4.0
-A2,2,4,16384,2179.78425743804,32,0.003848366172648072,5440,0.003772850442394648,32.0625,64.0,98304.0,8.0
-A2,2,4,32768,4344.834905351148,32,0.003861416225357031,11072,0.003778279349528763,64.0625,64.0,196608.0,16.0
-A2,2,4,65536,8674.946611121853,32,0.003867969856656271,22336,0.003781003096658628,128.0625,64.0,393216.0,32.0
-A4,4,2,8192,1103.675383502932,16,0.007600611670231015,1312,0.007480460399321815,16.0625,64.0,16384.0,8.0
-A4,4,2,16384,2197.0616234262543,16,0.00763620638634706,2720,0.007486362614786297,32.0625,64.0,32768.0,16.0
-A4,4,2,32768,4383.834103368879,16,0.007654129058897454,5536,0.007489334501679554,64.0625,64.0,65536.0,32.0
-A4,4,2,65536,8757.379062883949,16,0.007663121981838932,11168,0.007490825682997995,128.0625,64.0,131072.0,64.0
-B,8,1,8192,1125.9457520018113,8,0.014900554462921552,656,0.01466500492643045,16.0625,64.0,0.0,16.0
-B,8,1,16384,2242.543191943168,8,0.014962669223294344,1360,0.014669059716747552,32.0625,64.0,0.0,32.0
-B,8,1,32768,4475.650071825993,8,0.014994216018466134,2768,0.014671388277951352,64.0625,64.0,0.0,64.0
-B,8,1,65536,8942.79183160376,8,0.015008481750139946,5584,0.014671033662702533,128.0625,64.0,0.0,128.0
+A1,1,8,8192,1100.677128999906,64,0.0019053289513740923,5248,0.001875209310359147,16.0625,64.0,114688.0,2.0
+A1,1,8,16384,2177.674413966317,64,0.0019260473343055045,10880,0.0018882528874050511,32.0625,64.0,229376.0,4.0
+A1,1,8,32768,4334.294263847072,64,0.001935403433484189,22144,0.0018937339046091134,64.0625,64.0,458752.0,8.0
+A1,1,8,65536,8649.037693613021,64,0.0019397783423315907,44672,0.0018961647042087428,128.0625,64.0,917504.0,16.0
+A2,2,4,8192,1093.6804565019731,32,0.003835036070239803,2624,0.003774411415563731,16.0625,64.0,49152.0,4.0
+A2,2,4,16384,2175.2205604341793,32,0.0038564401939640384,5440,0.0037807660287830624,32.0625,64.0,98304.0,8.0
+A2,2,4,32768,4338.30076834928,32,0.003867232102121915,11072,0.003783970009586559,64.0625,64.0,196608.0,16.0
+A2,2,4,65536,8664.461183990821,32,0.003872650738157816,22336,0.0037855787340364574,128.0625,64.0,393216.0,32.0
+A4,4,2,8192,1103.915289502884,16,0.0075989598837576024,1312,0.007478834724463186,16.0625,64.0,16384.0,8.0
+A4,4,2,16384,2197.52387042685,16,0.007634600117789289,2720,0.007484787865719574,32.0625,64.0,32768.0,16.0
+A4,4,2,32768,4384.131109368463,16,0.0076536105246211355,5536,0.007488827131524694,64.0625,64.0,65536.0,32.0
+A4,4,2,65536,8757.668308885037,16,0.007662868886230499,11168,0.007490578277947104,128.0625,64.0,131072.0,64.0
+B,8,1,8192,1125.2657520018113,8,0.014909558893223249,656,0.014673867013748253,16.0625,64.0,0.0,16.0
+B,8,1,16384,2241.999191943168,8,0.01496629977413786,1360,0.014672619026008048,32.0625,64.0,0.0,32.0
+B,8,1,32768,4475.466071825993,8,0.014994832476216706,2768,0.014671991463273241,64.0625,64.0,0.0,64.0
+B,8,1,65536,8942.39983160376,8,0.01500913966355881,5584,0.014671676783710771,128.0625,64.0,0.0,128.0