ywkang/kernbench2
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

paper(gqa): table cleanup — drop measured-comm column, narrower fig, wrapped title

Browse Source 
Removes the (measured) comm column, widens the Notes column with
shorter wrapped text, scales the table figure down, and wraps the
heading to two lines so it sits flush with the table.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
Mukesh Garg
2026-06-18 14:02:42 -07:00
parent 1552028c25
commit 73e0b315fe
3 changed files with 21 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scripts/paper/paper_plot_gqa_kv_sharding_diagram.py
+21 -20
View File
@@ -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")