paper(gqa): rename long-ctx artifacts to gqa_long_ctx_6cases_* and add 4 decode 6-case charts

Filename cleanup so every long-ctx GQA artifact has a consistent
"gqa_long_ctx_6cases_*" prefix (or "gqa_decode_long_ctx_6cases_*"
for decode-only charts). Old "4cases" / mixed names retired.

Renames (long_ctx + figures, content unchanged):
  gqa_hbm_budget.png                  -> gqa_long_ctx_6cases_hbm_budget.png
  gqa_4cases_summary.png              -> gqa_long_ctx_6cases_summary.png
  gqa_4cases_memory_comm_analytical   -> gqa_long_ctx_6cases_memory_comm_analytical.png
  gqa_4cases_memory_comm_paired       -> gqa_long_ctx_6cases_memory_comm_paired.png
  gqa_kv_sharding_6cases_diagram      -> gqa_long_ctx_6cases_kv_sharding_diagram.png
  gqa_kv_sharding_6cases_table        -> gqa_long_ctx_6cases_kv_sharding_table.png
  gqa_3cases_measured_comm.json       -> gqa_long_ctx_6cases_measured_comm.json
  gqa_decode_long_ctx_4cases_*.png    -> gqa_decode_long_ctx_6cases_*.png
                                         (figures dir; long_ctx never had old)

New 4-chart 6-case set in long_ctx output dir (regenerated by
paper_plot_gqa_decode_long_ctx_4cases.py, which now reads all 6
sweep_decode.json panels — Cases 1-6 with the same colour scheme
used elsewhere: red = overflow per-PE HBM, grey = neutral, blue
= Pareto-best ★):

  gqa_decode_long_ctx_6cases_latency.png
  gqa_decode_long_ctx_6cases_memory.png
  gqa_decode_long_ctx_6cases_parallelism.png
  gqa_decode_long_ctx_6cases_traffic.png

Generator scripts updated to write the new filenames + handle the
two new d_head-TP variants (Cases 4, 5) in their per-PE memory and
active-PE-count helpers. Figure widths bumped 10 -> 12 in to fit 6
multi-line case labels.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

scripts/paper/measure_gqa_decode_placement_comm.py

@@ -349,7 +349,7 @@ def main() -> int:
         Path(__file__).resolve().parents[2]
         / "src" / "kernbench" / "benches"
         / "1H_milestone_output" / "gqa" / "long_ctx"
-        / "gqa_3cases_measured_comm.json"
+        / "gqa_long_ctx_6cases_measured_comm.json"
     )
     out_path.parent.mkdir(parents=True, exist_ok=True)
     out_path.write_text(json.dumps(out, indent=2))

scripts/paper/paper_plot_gqa_4cases_summary.py

@@ -241,7 +241,7 @@ _OUT_DIR = (
     / "src" / "kernbench" / "benches"
     / "1H_milestone_output" / "gqa" / "long_ctx"
 )
-_MEASURED_JSON = _OUT_DIR / "gqa_3cases_measured_comm.json"
+_MEASURED_JSON = _OUT_DIR / "gqa_long_ctx_6cases_measured_comm.json"
 
 
 def _load_measured() -> dict[int, float] | None:
@@ -537,7 +537,7 @@ def main() -> Path:
     fig_b, ax_b = plt.subplots(figsize=(4.0, 6.0))
     _plot_budget(ax_b)
     fig_b.tight_layout()
-    out_b = _OUT_DIR / "gqa_hbm_budget.png"
+    out_b = _OUT_DIR / "gqa_long_ctx_6cases_hbm_budget.png"
     fig_b.savefig(out_b, dpi=150)
     plt.close(fig_b)
     print(f"wrote {out_b}")
@@ -552,7 +552,7 @@ def main() -> Path:
     _plot_memory(ax_m)
     _plot_comm(ax_c)
     fig.tight_layout()
-    out = _OUT_DIR / "gqa_4cases_summary.png"
+    out = _OUT_DIR / "gqa_long_ctx_6cases_summary.png"
     fig.savefig(out, dpi=150)
     plt.close(fig)
     print(f"wrote {out}")
@@ -565,7 +565,7 @@ def main() -> Path:
     _plot_memory(ax_m2)
     _plot_comm(ax_c2, mode="analytical")
     fig2.tight_layout()
-    out2 = _OUT_DIR / "gqa_4cases_memory_comm_analytical.png"
+    out2 = _OUT_DIR / "gqa_long_ctx_6cases_memory_comm_analytical.png"
     fig2.savefig(out2, dpi=150)
     plt.close(fig2)
     print(f"wrote {out2}")
@@ -578,7 +578,7 @@ def main() -> Path:
     _plot_memory(ax_m3)
     _plot_comm(ax_c3, mode="paired")
     fig3.tight_layout()
-    out3 = _OUT_DIR / "gqa_4cases_memory_comm_paired.png"
+    out3 = _OUT_DIR / "gqa_long_ctx_6cases_memory_comm_paired.png"
     fig3.savefig(out3, dpi=150)
     plt.close(fig3)
     print(f"wrote {out3}")

scripts/paper/paper_plot_gqa_decode_long_ctx_4cases.py

@@ -1,13 +1,25 @@
 """Comparative figures for milestone-gqa-decode-long-ctx-4cases.
 
-Reads sweep.json (emitted by ``kernbench run --bench
+Reads sweep_decode.json (emitted by the milestone-1h-gqa bench) and
-milestone-gqa-decode-long-ctx-4cases``) and writes four PNGs into
+writes four PNGs into the same bench-output dir
-``docs/report/1H-codesign-paper/figures/``:
+(src/kernbench/benches/1H_milestone_output/gqa/long_ctx/):
 
-  gqa_decode_long_ctx_4cases_latency.png       end-to-end latency per case
+  gqa_decode_long_ctx_6cases_latency.png       end-to-end latency per case
-  gqa_decode_long_ctx_4cases_traffic.png       ipcq/dma op-count breakdown
+  gqa_decode_long_ctx_6cases_traffic.png       ipcq/dma op-count breakdown
-  gqa_decode_long_ctx_4cases_memory.png        per-PE KV bytes per case
+  gqa_decode_long_ctx_6cases_memory.png        per-PE KV bytes per case
-  gqa_decode_long_ctx_4cases_parallelism.png   per-PE S_local (compute work)
+  gqa_decode_long_ctx_6cases_parallelism.png   per-PE S_local (compute work)
+
+Filename still says "4cases" for backwards compat, but the script now
+covers all SIX kv-sharding strategies from the analytical chart
+(`gqa_4cases_summary.png`) — the original 4 plus the two new
+d_head-TP variants:
+
+  Case 1   Cube-Repl × PE-repl      (PE-TP doesn't shard KV)
+  Case 2   Cube-SP   × PE-repl
+  Case 3   Cube-Repl × PE-SP
+  Case 4   Cube-SP   × PE-TP-d_head      ← NEW
+  Case 5   Cube-TP-d_head × PE-SP        ← NEW
+  Case 6 ★ Cube-SP   × PE-SP             (Pareto-best)
 
 Run (after the bench):
     GQA_DECODE_LONG_CTX_4CASES_RUN=1 python -m kernbench.cli.main run \\
@@ -32,16 +44,27 @@ _FIG_DIR = (
 )
 _SWEEP_JSON = _FIG_DIR / "sweep_decode.json"
 
-# Panel name → (short label, case ordinal for left-to-right plot order).
+# Panel name → (short label, case ordinal, accent flag) using the
+# analytical chart's memory-descending ordering. PE-TP doesn't shard
+# KV memory, so the cube_repl_pe_tp panel maps to Case 1 (no
+# sharding, KV-wise) and cube_sp_pe_tp panel maps to Case 2.
+_NORMAL, _OVERFLOW, _PARETO = "normal", "overflow", "pareto"
+
 _CASE_INFO = {
-    "single_kv_group_decode_long_ctx_gqa_cube_sp_pe_tp": (
+    # panel name                                                       label                            ord  flag
-        "Case 1\nCube-SP × PE-TP", 1),
+    "single_kv_group_decode_long_ctx_gqa_cube_repl_pe_tp":            ("Case 1\nCube-Repl × PE-repl",   1, _OVERFLOW),
-    "single_kv_group_decode_long_ctx_gqa_cube_repl_pe_tp": (
+    "single_kv_group_decode_long_ctx_gqa_cube_sp_pe_tp":              ("Case 2\nCube-SP × PE-repl",     2, _OVERFLOW),
-        "Case 2\nCube-Repl × PE-TP", 2),
+    "single_kv_group_decode_long_ctx_gqa_cube_repl_pe_sp":            ("Case 3\nCube-Repl × PE-SP",     3, _OVERFLOW),
-    "single_kv_group_decode_long_ctx_gqa_cube_repl_pe_sp": (
+    "single_kv_group_decode_long_ctx_gqa_cube_sp_pe_tp_dhead":        ("Case 4\nCube-SP × PE-TP-d_head", 4, _NORMAL),
-        "Case 3\nCube-Repl × PE-SP", 3),
+    "single_kv_group_decode_long_ctx_gqa_cube_tp_dhead_pe_sp":        ("Case 5\nCube-TP-d_head × PE-SP", 5, _NORMAL),
-    "single_kv_group_decode_long_ctx_gqa_cube_sp_pe_sp": (
+    "single_kv_group_decode_long_ctx_gqa_cube_sp_pe_sp":              ("Case 6 ★\nCube-SP × PE-SP",     6, _PARETO),
-        "Case 4 ★\nCube-SP × PE-SP", 4),
+}
+
+# Bar fill colour per flag (used by every panel).
+_FLAG_COLOR = {
+    _NORMAL:   "#888888",     # neutral grey
+    _OVERFLOW: "#c0504d",     # red — fails the per-PE HBM budget
+    _PARETO:   "#3b6ea5",     # blue — Pareto-best
 }
 
 
@@ -53,23 +76,26 @@ def _sorted_by_case(rows: list[dict]) -> list[dict]:
     return sorted(rows, key=lambda r: _CASE_INFO[r["panel"]][1])
 
 
+def _bar_colors(rows: list[dict]) -> list[str]:
+    return [_FLAG_COLOR[_CASE_INFO[r["panel"]][2]] for r in rows]
+
+
 def _plot_latency(rows: list[dict]) -> Path:
     rows = _sorted_by_case(rows)
     labels = [_CASE_INFO[r["panel"]][0] for r in rows]
     lat_us = [r["latency_ns"] / 1e3 for r in rows]
-    colors = ["#888", "#888", "#888", "#3b6ea5"]  # Case 4 highlighted
+    fig, ax = plt.subplots(figsize=(12.0, 4.8))
-    fig, ax = plt.subplots(figsize=(8.0, 4.5))
+    bars = ax.bar(labels, lat_us, color=_bar_colors(rows), width=0.6)
-    bars = ax.bar(labels, lat_us, color=colors, width=0.6)
     ax.set_ylabel("end-to-end latency (µs)")
     ax.set_title(
-        "Long-context decode 4-cases — end-to-end latency per case\n"
+        "Long-context decode 6-cases — end-to-end latency per case\n"
         "LLaMA-3.1-70B single-KV-head group (8 cubes × 8 PEs)"
     )
     ax.bar_label(bars, fmt="%.1f", padding=3, fontsize=9)
     ax.grid(axis="y", ls=":", alpha=0.5)
     ax.set_ylim(0, max(lat_us) * 1.15)
     fig.tight_layout()
-    out = _FIG_DIR / "gqa_decode_long_ctx_4cases_latency.png"
+    out = _FIG_DIR / "gqa_decode_long_ctx_6cases_latency.png"
     fig.savefig(out, dpi=150)
     plt.close(fig)
     return out
@@ -83,18 +109,18 @@ def _plot_traffic(rows: list[dict]) -> Path:
     disp = ["IPCQ copy", "DMA read", "DMA write"]
     colors = ["#c0504d", "#9bbb59", "#8064a2"]
     w = 0.25
-    fig, ax = plt.subplots(figsize=(9.0, 4.5))
+    fig, ax = plt.subplots(figsize=(11.0, 4.5))
     for i, (k, d, c) in enumerate(zip(keys, disp, colors)):
         vals = [r["op_log_summary"][k] for r in rows]
         ax.bar([xi + (i - 1) * w for xi in x], vals, width=w, label=d, color=c)
     ax.set_xticks(list(x))
     ax.set_xticklabels(labels, fontsize=9)
     ax.set_ylabel("op count")
-    ax.set_title("Long-context decode 4-cases — op-count breakdown per case")
+    ax.set_title("Long-context decode 6-cases — op-count breakdown per case")
     ax.legend(fontsize=9)
     ax.grid(axis="y", ls=":", alpha=0.5)
     fig.tight_layout()
-    out = _FIG_DIR / "gqa_decode_long_ctx_4cases_traffic.png"
+    out = _FIG_DIR / "gqa_decode_long_ctx_6cases_traffic.png"
     fig.savefig(out, dpi=150)
     plt.close(fig)
     return out
@@ -103,27 +129,41 @@ def _plot_traffic(rows: list[dict]) -> Path:
 def _s_local_per_pe(panel: str, *, S_kv: int, C: int, P: int) -> int:
     """S_local (token count) each PE attends over locally.
 
-    Encodes the cube/pe sharding axes from the panel name:
+      cube_repl_pe_tp  (Case 1): S_kv        (no sharding, KV-wise)
-      cube_sp_pe_tp (Case 1): S_kv / C       (pe=replicate within cube)
+      cube_sp_pe_tp    (Case 2): S_kv / C    (cube splits S_kv, PEs replicate)
-      cube_repl_pe_tp (Case 2): S_kv         (pe=replicate; only 1 PE works)
+      cube_repl_pe_sp  (Case 3): S_kv / P
-      cube_repl_pe_sp (Case 3): S_kv / P     (pe=row_wise within cube)
+      cube_sp_pe_tp_dhead (Case 4): S_kv / C (cube splits S_kv, PE splits d_head)
-      cube_sp_pe_sp (Case 4): S_kv / (C·P)   (★ 64-way split)
+      cube_tp_dhead_pe_sp (Case 5): S_kv / P (cube splits d_head, PE splits S_kv)
+      cube_sp_pe_sp    (Case 6 ★): S_kv / (C·P)
     """
-    S_per_cube = S_kv if "cube_repl" in panel else S_kv // C
+    cube_splits_s = "cube_sp" in panel
-    return S_per_cube // P if "pe_sp" in panel else S_per_cube
+    pe_splits_s   = "pe_sp" in panel
+    S_per_cube = S_kv // C if cube_splits_s else S_kv
+    return S_per_cube // P if pe_splits_s else S_per_cube
+
+
+def _d_head_per_pe(panel: str, *, d_head: int, C: int, P: int) -> int:
+    """d_head dims each PE owns (Cases 4 and 5 shard d_head)."""
+    if "cube_tp_dhead" in panel:    # Case 5: cube shards d_head
+        return d_head // C
+    if "pe_tp_dhead" in panel:      # Case 4: PE shards d_head
+        return d_head // P
+    return d_head                   # Cases 1, 2, 3, 6: full d_head per PE
 
 
 def _active_pe_count(panel: str, *, C: int, P: int) -> int:
     """Number of PEs doing non-idle attention work.
 
-      cube_sp_pe_tp (Case 1): C       (PE 0 of each cube; 7 PEs idle per cube)
+      cube_repl_pe_tp  (Case 1): 1     (PE-TP idle for B=1; only one PE works)
-      cube_repl_pe_tp (Case 2): 1     (only PE 0 of CUBE 0)
+      cube_sp_pe_tp    (Case 2): C     (PE 0 of each cube; 7 PEs idle per cube)
-      cube_repl_pe_sp (Case 3): C·P   (all PEs busy, but cubes are redundant)
+      cube_repl_pe_sp  (Case 3): C·P   (all PEs busy, cube-side redundant)
-      cube_sp_pe_sp (Case 4): C·P     (all 64 PEs doing unique work)
+      cube_sp_pe_tp_dhead (Case 4): C·P (PE shards d_head — all 64 active)
+      cube_tp_dhead_pe_sp (Case 5): C·P (PE shards S_kv — all active)
+      cube_sp_pe_sp    (Case 6 ★): C·P (all 64 PEs doing unique work)
     """
-    if "cube_repl" in panel and "pe_tp" in panel:
+    if "cube_repl" in panel and "pe_tp" in panel and "dhead" not in panel:
         return 1
-    if "cube_sp" in panel and "pe_tp" in panel:
+    if "cube_sp" in panel and "pe_tp" in panel and "dhead" not in panel:
         return C
     return C * P
 
@@ -132,11 +172,12 @@ def _kv_bytes_per_pe(panel: str, *, S_kv: int, h_kv: int,
                      d_head: int, C: int, P: int) -> int:
     """KV bytes a single PE references (K + V, f16, 2 B/elem)."""
     s_local = _s_local_per_pe(panel, S_kv=S_kv, C=C, P=P)
-    return 2 * s_local * h_kv * d_head * 2
+    d_local = _d_head_per_pe(panel, d_head=d_head, C=C, P=P)
+    return 2 * s_local * h_kv * d_local * 2
 
 
 def _plot_memory(rows: list[dict]) -> Path:
-    """Per-PE KV bytes — Case 4 wins (64-way split)."""
+    """Per-PE KV bytes — Case 6 ★ wins (64-way split)."""
     rows = _sorted_by_case(rows)
     labels = [_CASE_INFO[r["panel"]][0] for r in rows]
     mib_per_pe = [
@@ -146,26 +187,25 @@ def _plot_memory(rows: list[dict]) -> Path:
         ) / (1024 * 1024)
         for r in rows
     ]
-    colors = ["#888", "#c0504d", "#888", "#3b6ea5"]  # 4 highlighted, 2 marked red
+    fig, ax = plt.subplots(figsize=(12.0, 4.8))
-    fig, ax = plt.subplots(figsize=(8.0, 4.5))
+    bars = ax.bar(labels, mib_per_pe, color=_bar_colors(rows), width=0.6)
-    bars = ax.bar(labels, mib_per_pe, color=colors, width=0.6)
     ax.set_ylabel("KV bytes per PE (MiB, K + V, f16)")
     ax.set_title(
-        "Long-context decode 4-cases — KV memory per PE\n"
+        "Long-context decode 6-cases — KV memory per PE\n"
         "(one KV-head group; per-layer, per-token state)"
     )
     ax.bar_label(bars, fmt="%.3f", padding=3, fontsize=9)
     ax.grid(axis="y", ls=":", alpha=0.5)
     ax.set_ylim(0, max(mib_per_pe) * 1.15)
     fig.tight_layout()
-    out = _FIG_DIR / "gqa_decode_long_ctx_4cases_memory.png"
+    out = _FIG_DIR / "gqa_decode_long_ctx_6cases_memory.png"
     fig.savefig(out, dpi=150)
     plt.close(fig)
     return out
 
 
 def _plot_parallelism(rows: list[dict]) -> Path:
-    """Total active PE-token compute load — exposes Case 3's redundancy."""
+    """Total active PE-token compute load — exposes redundant-work cases."""
     rows = _sorted_by_case(rows)
     labels = [_CASE_INFO[r["panel"]][0] for r in rows]
     total_work = [
@@ -173,19 +213,19 @@ def _plot_parallelism(rows: list[dict]) -> Path:
         * _s_local_per_pe(r["panel"], S_kv=r["S_kv"], C=r["C"], P=r["P"])
         for r in rows
     ]
-    colors = ["#888", "#888", "#c0504d", "#3b6ea5"]  # 4 highlighted, 3 marked red
+    fig, ax = plt.subplots(figsize=(12.0, 4.8))
-    fig, ax = plt.subplots(figsize=(8.0, 4.5))
+    bars = ax.bar(labels, total_work, color=_bar_colors(rows), width=0.6)
-    bars = ax.bar(labels, total_work, color=colors, width=0.6)
     ax.set_ylabel("active-PE × S_local (PE-tokens; lower ⇒ less wasted work)")
     ax.set_title(
-        "Long-context decode 4-cases — total compute load across active PEs\n"
+        "Long-context decode 6-cases — total compute load across active PEs\n"
-        "(Case 3 replicates the full K/V across 8 cubes ⇒ 8× wasted PE-tokens)"
+        "(Case 3 replicates KV across 8 cubes → 8× wasted PE-tokens; "
+        "Case 6 ★ is fully parallel without replication)"
     )
     ax.bar_label(bars, fmt="%d", padding=3, fontsize=9)
     ax.grid(axis="y", ls=":", alpha=0.5)
     ax.set_ylim(0, max(total_work) * 1.15)
     fig.tight_layout()
-    out = _FIG_DIR / "gqa_decode_long_ctx_4cases_parallelism.png"
+    out = _FIG_DIR / "gqa_decode_long_ctx_6cases_parallelism.png"
     fig.savefig(out, dpi=150)
     plt.close(fig)
     return out

scripts/paper/paper_plot_gqa_kv_sharding_diagram.py

@@ -241,7 +241,7 @@ def _make_table_png() -> Path:
         "(LLaMA 70B GQA single KV-head group · S_kv = 1 M, FP16, 80 layers)",
         fontsize=12, y=0.97,
     )
-    out = _OUT_DIR / "gqa_kv_sharding_6cases_table.png"
+    out = _OUT_DIR / "gqa_long_ctx_6cases_kv_sharding_table.png"
     fig.savefig(out, dpi=150, bbox_inches="tight")
     plt.close(fig)
     print(f"wrote {out}")
@@ -309,7 +309,7 @@ def main() -> Path:
         fontsize=12, y=0.99,
     )
 
-    out = _OUT_DIR / "gqa_kv_sharding_6cases_diagram.png"
+    out = _OUT_DIR / "gqa_long_ctx_6cases_kv_sharding_diagram.png"
     fig.savefig(out, dpi=150, bbox_inches="tight")
     plt.close(fig)
     print(f"wrote {out}")