diff --git a/docs/report/1H-codesign-paper/figures/gqa_long_ctx_6cases_kv_sharding_table.png b/docs/report/1H-codesign-paper/figures/gqa_long_ctx_6cases_kv_sharding_table.png index 4c11a41..b9fa974 100644 Binary files a/docs/report/1H-codesign-paper/figures/gqa_long_ctx_6cases_kv_sharding_table.png and b/docs/report/1H-codesign-paper/figures/gqa_long_ctx_6cases_kv_sharding_table.png differ diff --git a/scripts/paper/paper_plot_gqa_kv_sharding_diagram.py b/scripts/paper/paper_plot_gqa_kv_sharding_diagram.py index a00d470..76eabdf 100644 --- a/scripts/paper/paper_plot_gqa_kv_sharding_diagram.py +++ b/scripts/paper/paper_plot_gqa_kv_sharding_diagram.py @@ -163,40 +163,40 @@ def _draw_panel(ax, cfg): def _make_table_png() -> Path: """Slide-14 companion table: per-PE memory + comm for all 6 cases.""" headers = ["Case", "Sharding", "KV / PE", "Fit", - "Comm/tok\n(analytical)", "Comm/tok\n(measured)", "Notes"] + "Comm/tok\n(analytical)", "Notes"] rows = [ ("Case 1", "Cube-Repl · PE-repl", "40 GB", "✗", - "1.2 MB", "1.25 MB", - "no sharding — full KV on every PE"), + "1.2 MB", + "no sharding —\nfull KV on every PE"), ("Case 2", "Cube-SP · PE-repl", "5 GB", "✗", - "3.8 MB", "1.27 MB", - "cube-axis sharded only"), + "3.8 MB", + "cube-axis\nsharded only"), ("Case 3", "Cube-Repl · PE-SP", "5 GB", "✗", - "3.8 MB", "1.39 MB", - "PE-axis sharded only"), + "3.8 MB", + "PE-axis\nsharded only"), ("Case 4", "Cube-SP · PE-TP-d_head", "640 MB", "✓", - "166 MB", "162 MB", - "d_head split intra-cube — partial-score AR ∝ S_kv"), + "166 MB", + "d_head split intra-cube\npartial-score AR ∝ S_kv"), ("Case 5", "Cube-TP-d_head · PE-SP", "640 MB", "✓", - "166 MB", "162 MB", - "d_head split inter-cube — partial-score AR on UCIe"), + "166 MB", + "d_head split inter-cube\npartial-score AR on UCIe"), ("Case 6 ★", "Cube-SP · PE-SP", "640 MB", "✓", - "6.2 MB", "1.41 MB", - "S_kv split on both axes — (m,ℓ,O) AR only, S_kv-indep."), + "6.2 MB", + "S_kv split both axes\n(m,ℓ,O) AR only"), ] accents = [_ACC["red"], _ACC["orange"], _ACC["orange"], _ACC["blue"], _ACC["blue"], _ACC["green"]] - fig, ax = plt.subplots(figsize=(20.0, 4.6)) + fig, ax = plt.subplots(figsize=(15.0, 5.0)) ax.set_axis_off() cell_data = [headers] + [list(r) for r in rows] tbl = ax.table(cellText=cell_data, - colWidths=[0.07, 0.18, 0.08, 0.05, 0.12, 0.12, 0.38], + colWidths=[0.07, 0.20, 0.09, 0.05, 0.14, 0.28], cellLoc="center", loc="center") tbl.auto_set_font_size(False) - tbl.set_fontsize(11) - tbl.scale(1.0, 2.2) + tbl.set_fontsize(10.5) + tbl.scale(1.0, 2.4) n_cols = len(headers) n_rows = len(rows) + 1 # +1 header @@ -237,9 +237,10 @@ def _make_table_png() -> Path: tbl[(0, n_cols - 1)].get_text().set_ha("left") fig.suptitle( - "GQA decode KV-sharding — per-PE memory & communication " - "(LLaMA 70B GQA single KV-head group · S_kv = 1 M, FP16, 80 layers)", - fontsize=12, y=0.97, + "GQA decode KV-sharding — per-PE memory & communication\n" + "(LLaMA 70B GQA single KV-head group · S_kv = 1 M, FP16, " + "80 layers)", + fontsize=11.5, y=0.94, ) out = _OUT_DIR / "gqa_long_ctx_6cases_kv_sharding_table.png" fig.savefig(out, dpi=150, bbox_inches="tight") diff --git a/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/gqa_long_ctx_6cases_kv_sharding_table.png b/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/gqa_long_ctx_6cases_kv_sharding_table.png index 4c11a41..b9fa974 100644 Binary files a/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/gqa_long_ctx_6cases_kv_sharding_table.png and b/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/gqa_long_ctx_6cases_kv_sharding_table.png differ