"""GQA 4-cases (+ 2 d_head-TP variants) — combined memory + comm summary. Single PNG with two panels side-by-side: (left) Per-PE KV memory at S_kv = 1 M tokens (across 80 layers). (right) Per-PE communication per output token at S_kv = 1 M (decode). ================================================================ SYSTEM INPUTS (LLaMA-3.1-70B single-KV-head group) ================================================================ N_layers = 80 (transformer layers per token) h_kv = 1 (per KV group) h_q = 8 (query heads per KV group) d_head = 128 d_model = 8192 (LLaMA-3.1-70B hidden dim — for Wo / FFN AR) bytes = 2 (FP16) C = 8 cubes per KV group P = 8 PEs per cube HBM_per_PE = 6.0 GB Attn weights = 1.76 GB → KV headroom = 4.24 GB / PE B = 1 user S_q = 1 token (decode) ================================================================ (1) PER-PE KV MEMORY (left panel) ================================================================ KV bytes per token across all 80 layers, single KV group: KV/tok = 2 (K+V) · h_kv · d_head · bytes · N_layers = 2 · 1 · 128 · 2 · 80 = 40 KB / token Per-PE share = KV/tok ÷ divisor, where divisor depends on the K/V tensor placement (NOT on the compute-side label): Case 1 Cube-SP × PE-replicate divisor = C = 8 Case 1' Cube-SP × PE-TP (on d_head) divisor = C·P = 64 Case 2 Cube-Repl × PE-replicate divisor = 1 = 1 Case 3 Cube-Repl × PE-SP divisor = P = 8 Case 3' Cube-TP × PE-SP divisor = C·P = 64 Case 4 Cube-SP × PE-SP ★ divisor = C·P = 64 Per-PE bytes @ S_kv = 1 M = KV/tok · 1 M / divisor: Case 1 → 5.0 GB ✗ (exceeds 4.24 GB headroom) Case 1' → 640 MB ✓ Case 2 → 40.0 GB ✗ Case 3 → 5.0 GB ✗ Case 3' → 640 MB ✓ Case 4 → 640 MB ✓ Max KV context per PE = 4.24 GB · 1024² / (KV/tok ÷ divisor): Case 1, 3 : 889 K tokens Case 1', 3', 4 : 7.11 M tokens ✓ Case 2 : 111 K tokens ================================================================ (2) PER-PE COMMUNICATION (right panel) ================================================================ (A) WEIGHT AllReduces (constant across all cases) Wo AR ≈ 8 KB / layer / PE (Wq replicated → AR partial Y) FFN AR ≈ 8 KB / layer / PE × 80 layers = 1.25 MB / token / PE (the bottom blue stack on every bar) (B) ATTENTION-TIME collective (the differentiator) Cases that compute partial attention locally (PE-SP / PE-repl): Q is replicated and each rank attends to its own complete KV slice (full d_head) → only the small (m, ℓ, O) triple needs AllReducing. Payload (h_q · S_q · d_head · 2) ≈ 4 KB per AR step, hierarchical ≈ 32 KB / layer. CONSTANT in S_kv. Case 1 : inter-cube AR on (m,ℓ,O) → ~32 KB / layer Case 2 : nothing → 0 Case 3 : intra-cube AR on (m,ℓ,O) → ~32 KB / layer Case 4 ★ : intra + inter-cube AR on (m,ℓ,O) → ~64 KB / layer d_head-TP variants (1', 3'): Each rank holds only d_head/divisor dims → Q·K^T produces only partial sums → must AR PARTIAL SCORES before softmax. The score tile per AR is (h_q · S_q · S_kv / slices_for_seq) · 2 bytes, which SCALES with S_kv. Case 1' : intra-cube partial-score AR payload = h_q · S_q · (S_kv/C) · 2 = 2 MB / layer at 1 M Case 3' : inter-cube partial-score AR (UCIe, slower than NoC) payload = h_q · S_q · (S_kv/P) · 2 = 2 MB / layer at 1 M Plus tiny (m,ℓ,O) AR for both 1' and 3' to combine partial attentions across the remaining axis (~32 KB / layer). (C) Total per output token per PE (sum over 80 layers, S_kv = 1 M): Case 1 ≈ 4 MB (1.25 MB Wo+FFN + 2.5 MB inter-cube AR) Case 1' ≈ 166 MB (1.25 MB + 160 MB partial-score AR + ~5 MB other) Case 2 ≈ 1.2 MB (Wo+FFN only) Case 3 ≈ 4 MB (1.25 MB + 2.5 MB intra-cube AR) Case 3' ≈ 166 MB (1.25 MB + 160 MB partial-score AR + ~5 MB other) Case 4 ★ ≈ 6 MB (1.25 MB + 5 MB 2-phase AR) ================================================================ KEY TAKEAWAYS ================================================================ - Memory winners (640 MB / PE @ 1M): Cases 1', 3', 4 (any 64-way sharding). - Comm winner among those three: Case 4 (6 MB). Cases 1' and 3' both pay ~160 MB because the d_head-sharded variant ARs a 2 MB partial- score tile every layer — the score tensor scales with S_kv while the (m, ℓ, O) triple does not. - Cases 1 and 3 (8-way sharding) don't fit 1M context (5 GB vs 4.24 GB headroom). - Case 2 has the cheapest comm but the most memory (40 GB / PE). - Case 4 ★ is the Pareto-best: fits 1M context AND lowest comm among memory-feasible options. Output PNG: src/kernbench/benches/1H_milestone_output/gqa/long_ctx/ gqa_4cases_summary.png """ from __future__ import annotations import json from pathlib import Path import matplotlib.patches as mpatches import matplotlib.pyplot as plt # ── System constants (LLaMA-3.1-70B single-KV-head group, decode) ─── _N_LAYERS = 80 _H_KV = 1 _H_Q = 8 _D_HEAD = 128 _D_MODEL = 8192 _BYTES_PER_ELEM = 2 # FP16 _C = 8 # cubes per KV group _P = 8 # PEs per cube _B = 1 _S_Q = 1 # decode _HBM_PER_PE_GB = 6.0 _WEIGHTS_PER_PE_GB = 1.76 _HEADROOM_GB = _HBM_PER_PE_GB - _WEIGHTS_PER_PE_GB # 4.24 GB _HEADLINE_S_KV = 1 << 20 # 1 Mi tokens # Per-token KV bytes (single KV group, all 80 layers). _KV_PER_TOK_BYTES = ( 2 * _H_KV * _D_HEAD * _BYTES_PER_ELEM * _N_LAYERS ) # 40 KB # Per-token Wo + FFN AR (constant across cases). _WO_PER_LAYER_BYTES = 8 * 1024 _FFN_PER_LAYER_BYTES = 8 * 1024 # ── Per-PE attention weight breakdown (single KV-head group) ──────── # # LLaMA-3.1-70B single-KV-head group dimensions: # d_model = 8192 # h_q per group = 8 (8 query heads attend to 1 KV head per group) # h_kv per group = 1 # d_head = 128 # FP16 (2 bytes) # # Per layer, per PE: # Wq shape (d_model, h_q · d_head) REPLICATED across all 64 PEs of group # = 8192 · 8·128 · 2 bytes = 16 MB / layer / PE # Wk shape (d_model, h_kv · d_head) REPLICATED # = 8192 · 1·128 · 2 bytes = 2 MB / layer / PE # Wv shape (d_model, h_kv · d_head) REPLICATED # = 8192 · 1·128 · 2 bytes = 2 MB / layer / PE # Wo shape (h_q · d_head, d_model) ROW-SPLIT across C=8 cubes, # replicated within cube # = (1024/C) · 8192 · 2 bytes = 2 MB / layer / PE # ──────────────────────────────────────────────────────────────── # total attn weights / layer / PE = 22 MB # × 80 layers = 1.76 GB / PE (decimal GB) # # KV headroom per PE = HBM_PER_PE - weights = 6.0 - 1.76 = 4.24 GB. # (FFN weights are accounted for in a separate budget, not in this 4.24 GB.) # Slide-17 convention: per-layer values in binary MB (MiB), totals in # "GB" formed by ×80 layers ÷ 1000 — giving the canonical 1.76 GB total # and 4.24 GB headroom that the slide-17 chart reports. _WQ_MB_PER_LAYER = (_D_MODEL * _H_Q * _D_HEAD * _BYTES_PER_ELEM) / (1024 ** 2) # 16 _WK_MB_PER_LAYER = (_D_MODEL * _H_KV * _D_HEAD * _BYTES_PER_ELEM) / (1024 ** 2) # 2 _WV_MB_PER_LAYER = (_D_MODEL * _H_KV * _D_HEAD * _BYTES_PER_ELEM) / (1024 ** 2) # 2 _WO_MB_PER_LAYER = ((_H_Q * _D_HEAD) // _C * _D_MODEL * _BYTES_PER_ELEM) / (1024 ** 2) # 2 _WQ_GB = _WQ_MB_PER_LAYER * _N_LAYERS / 1000 # 1.28 GB _WK_GB = _WK_MB_PER_LAYER * _N_LAYERS / 1000 # 0.16 GB _WV_GB = _WV_MB_PER_LAYER * _N_LAYERS / 1000 # 0.16 GB _WO_GB = _WO_MB_PER_LAYER * _N_LAYERS / 1000 # 0.16 GB _WEIGHTS_GB = _WQ_GB + _WK_GB + _WV_GB + _WO_GB # 1.76 GB # (m, ℓ, O) AR payload — used by Cases 1, 3, 4 (different axes per case). _MLO_INTRA_BYTES_PER_LAYER = 32 * 1024 # ~32 KB intra-cube AR _MLO_INTER_BYTES_PER_LAYER = 32 * 1024 # ~32 KB inter-cube AR # Cases — renumbered in MEMORY-DESCENDING ORDER (left to right): # Case 1 (40 GB) : no sharding # Cases 2, 3 ( 5 GB) : single-axis sharding (cube OR PE only) # Cases 4, 5, 6 (640 MB) : two-axis sharding (cube AND PE) # — Case 6 ★ is the Pareto-best (lowest comm) _CASES = (1, 2, 3, 4, 5, 6) _DIVISOR = { 1: 1, # Cube-Repl × PE-replicate — no sharding 2: _C, # Cube-SP × PE-replicate — cube-axis only 3: _P, # Cube-Repl × PE-SP — PE-axis (S_kv) only 4: _C * _P, # Cube-SP × PE-TP (d_head) — 64-way (d_head intra) 5: _C * _P, # Cube-TP × PE-SP — 64-way (d_head inter) 6: _C * _P, # Cube-SP × PE-SP — 64-way (S_kv both axes) ★ } _CASE_LABEL = { 1: "Case 1\nCube-Repl\nPE-repl", 2: "Case 2\nCube-SP\nPE-repl", 3: "Case 3\nCube-Repl\nPE-SP", 4: "Case 4\nCube-SP\nPE-TP", 5: "Case 5\nCube-TP\nPE-SP", 6: "Case 6 ★\nCube-SP\nPE-SP", } _CASE_COLOR = { 1: "#C0504D", # red — worst memory (no sharding) 2: "#E0834A", # orange — single-axis sharded (cube) 3: "#EBA854", # tan — single-axis sharded (PE) 4: "#A6C2E0", # light blue — d_head-TP 64-way (PE) 5: "#C7D8A0", # light green — d_head-TP 64-way (cube) 6: "#8064A2", # purple — Pareto winner ★ (S_kv 64-way) } _ATTN_DESC = { 1: "none", 2: "online-softmax\nmerge of\n(m,ℓ,O)\ninter-cube", 3: "online-softmax\nmerge of\n(m,ℓ,O)\nintra-cube", 4: "partial\nscores\n+ (m,ℓ,O)\nmerge\n(d_head-TP)", 5: "partial\nscores\n+ (m,ℓ,O)\nmerge\n(d_head-TP)", 6: "online-softmax\nmerge of\n(m,ℓ,O)\nintra + inter", } _WO_COLOR = "#9EC5E8" _FFN_COLOR = "#4A78B8" _ATTN_COLOR = "#E07A3F" _OUT_DIR = ( Path(__file__).resolve().parents[2] / "src" / "kernbench" / "benches" / "1H_milestone_output" / "gqa" / "long_ctx" ) _MEASURED_JSON = _OUT_DIR / "gqa_3cases_measured_comm.json" def _load_measured() -> dict[int, float] | None: """Load measured per-PE comm bytes (already scaled to S_kv=1M). Returns {case_id: per_token_total_bytes} or None if JSON missing. Produced by scripts/paper/measure_gqa_decode_placement_comm.py. """ if not _MEASURED_JSON.exists(): return None data = json.loads(_MEASURED_JSON.read_text()) return { int(cid): info["per_pe_total_bytes_per_token_at_1M"] for cid, info in data["cases"].items() } # ── Formulae ──────────────────────────────────────────────────────── def kv_per_pe_bytes(case, s_kv: int) -> int: """Per-PE KV bytes at the given S_kv.""" return _KV_PER_TOK_BYTES * s_kv // _DIVISOR[case] def max_s_kv(case) -> int: """Max KV context per PE given the 4.24 GB headroom.""" headroom_bytes = int(_HEADROOM_GB * (1 << 30)) bytes_per_tok = _KV_PER_TOK_BYTES // _DIVISOR[case] return headroom_bytes // bytes_per_tok def _partial_score_bytes_per_layer(s_kv: int, slices_for_seq: int) -> int: """AR payload for partial scores in d_head-TP cases.""" return _H_Q * _S_Q * (s_kv // slices_for_seq) * _BYTES_PER_ELEM def attn_comm_per_layer_bytes(case: int, s_kv: int) -> int: """Per-layer per-PE attention-time comm bytes (decode, B=1, S_q=1). Case numbering follows the memory-descending order defined in _CASES (1=no sharding, 6=Pareto winner). """ if case == 1: # Cube-Repl × PE-repl — no sharding return 0 if case == 2: # Cube-SP × PE-repl — inter-cube AR on (m,ℓ,O) return _MLO_INTER_BYTES_PER_LAYER if case == 3: # Cube-Repl × PE-SP — intra-cube AR on (m,ℓ,O) return _MLO_INTRA_BYTES_PER_LAYER if case == 4: # Cube-SP × PE-TP(d_head) — partial-score AR return (_partial_score_bytes_per_layer(s_kv, _C) + _MLO_INTRA_BYTES_PER_LAYER + _MLO_INTER_BYTES_PER_LAYER) if case == 5: # Cube-TP(d_head) × PE-SP — partial-score AR return (_partial_score_bytes_per_layer(s_kv, _P) + _MLO_INTRA_BYTES_PER_LAYER + _MLO_INTER_BYTES_PER_LAYER) if case == 6: # Cube-SP × PE-SP — 2-phase (m,ℓ,O) AR ★ return _MLO_INTRA_BYTES_PER_LAYER + _MLO_INTER_BYTES_PER_LAYER raise ValueError(f"unknown case {case}") def per_token_bytes(case, s_kv: int) -> tuple[int, int, int]: """(Wo AR, FFN AR, Attn) bytes per output token per PE — × 80 layers.""" wo = _WO_PER_LAYER_BYTES * _N_LAYERS ffn = _FFN_PER_LAYER_BYTES * _N_LAYERS attn = attn_comm_per_layer_bytes(case, s_kv) * _N_LAYERS return wo, ffn, attn # ── Formatters ────────────────────────────────────────────────────── def _fmt_bytes(b: float) -> str: if b >= (1 << 30): return f"{b / (1 << 30):.2f} GB" if b >= (1 << 20): return f"{b / (1 << 20):.1f} MB" if b >= (1 << 10): return f"{b / (1 << 10):.0f} KB" return f"{b:.0f} B" def _fmt_tokens(n: int) -> str: if n >= 1_000_000: return f"{n / 1_000_000:.2f} M" if n >= 1_000: return f"{n / 1_000:.0f} K" return f"{n}" # ── Panels ────────────────────────────────────────────────────────── def _plot_budget(ax) -> None: """HBM budget per PE — stacked weights + KV headroom + HBM ceiling. Wq is the dominant weight slice (~1.28 GB). Wk, Wv, Wo are each small (~0.16 GB) so their slice labels would overlap on the bar — they're shown in the legend only, and only Wq + KV-headroom get on-bar annotations. """ components = [ ("Wq (REPL)", _WQ_GB, "#7B9CC4"), ("Wk (REPL)", _WK_GB, "#A0BBD8"), ("Wv (REPL)", _WV_GB, "#C5D6E8"), ("Wo (cube-split)", _WO_GB, "#E2EAF3"), ("KV cache headroom", _HBM_PER_PE_GB - _WEIGHTS_GB, "#9BBB59"), ] bottom = 0.0 for label, val, color in components: ax.bar(0, val, bottom=bottom, color=color, edgecolor="black", width=0.7, label=f"{label} · {val:.2f} GB") # Only annotate slices thick enough to fit text without overlap. if val >= 0.50: ax.text(0, bottom + val / 2, f"{label}\n{val:.2f} GB", ha="center", va="center", fontsize=9, weight="bold") bottom += val ax.axhline(_HBM_PER_PE_GB, color="red", ls="--", lw=1.4, label=f"HBM = {_HBM_PER_PE_GB} GB") ax.set_xticks([0]) ax.set_xticklabels(["per-PE HBM"], fontsize=10) ax.set_ylabel("GB per PE") ax.set_ylim(0, _HBM_PER_PE_GB * 1.10) ax.set_title( f"Per-PE HBM budget\n" f"weights {_WEIGHTS_GB:.2f} GB + KV = {_HBM_PER_PE_GB} GB", fontsize=10, ) ax.grid(axis="y", ls=":", alpha=0.5) ax.legend(loc="upper right", fontsize=7.5, framealpha=0.92) def _plot_memory(ax) -> None: vals_gb = [kv_per_pe_bytes(c, _HEADLINE_S_KV) / (1 << 30) for c in _CASES] labels = [_CASE_LABEL[c] for c in _CASES] colors = [_CASE_COLOR[c] for c in _CASES] x = list(range(len(_CASES))) bars = ax.bar(x, vals_gb, color=colors, width=0.65) ax.set_xticks(x) ax.set_xticklabels(labels, fontsize=9) ax.set_ylabel("KV bytes per PE (GB, log)") ax.set_yscale("log") ax.set_title( f"Per-PE KV memory at S_kv = {_HEADLINE_S_KV:,} tokens " f"(across {_N_LAYERS} layers, FP16)", fontsize=11, ) ax.grid(axis="y", ls=":", alpha=0.5, which="both") ax.axhline(_HEADROOM_GB, color="red", ls="--", lw=1.4, label=f"KV headroom = {_HEADROOM_GB} GB / PE") for bar, v_gb in zip(bars, vals_gb): v_bytes = v_gb * (1 << 30) fits = v_gb <= _HEADROOM_GB ax.text(bar.get_x() + bar.get_width() / 2, v_gb * 1.10, _fmt_bytes(v_bytes) + (" ✓" if fits else " ✗"), ha="center", va="bottom", fontsize=9, color="green" if fits else "red", weight="bold") ax.legend(loc="upper right", fontsize=9) def _plot_comm(ax, *, mode: str = "analytical") -> None: """Per-PE comm panel. mode = "analytical": single solid bars from per_token_bytes formula. mode = "paired" : analytical (solid) + simulator-measured (hatched) side-by-side per case, when the measurement JSON is available. """ labels = [_CASE_LABEL[c] for c in _CASES] wo_mb_list: list[float] = [] ffn_mb_list: list[float] = [] attn_mb: list[float] = [] for c in _CASES: wo, ffn, attn = per_token_bytes(c, _HEADLINE_S_KV) wo_mb_list.append(wo / (1 << 20)) ffn_mb_list.append(ffn / (1 << 20)) attn_mb.append(attn / (1 << 20)) measured = _load_measured() if mode == "paired" else None paired = measured is not None source_tag = "analytical (solid) vs simulator-measured (hatched)" \ if paired else "analytical" n_cases = len(_CASES) x = list(range(n_cases)) bar_w = 0.36 if paired else 0.65 x_ana = [xi - bar_w / 2 for xi in x] if paired else x x_meas = [xi + bar_w / 2 for xi in x] if paired else None # Analytical bars (solid). ax.bar(x_ana, wo_mb_list, width=bar_w, color=_WO_COLOR, edgecolor="black") ax.bar(x_ana, ffn_mb_list, width=bar_w, bottom=wo_mb_list, color=_FFN_COLOR, edgecolor="black") bottoms_attn = [w + f for w, f in zip(wo_mb_list, ffn_mb_list)] ax.bar(x_ana, attn_mb, width=bar_w, bottom=bottoms_attn, color=_ATTN_COLOR, edgecolor="black") # Measured bars (hatched) — same Wo+FFN base, attn from op_log. meas_attn_mb: list[float] = [] if paired: for i, c in enumerate(_CASES): meas_total = measured.get(c, 0) / (1 << 20) meas_attn_mb.append( max(meas_total - wo_mb_list[i] - ffn_mb_list[i], 0.0)) ax.bar(x_meas, wo_mb_list, width=bar_w, color=_WO_COLOR, edgecolor="black", hatch="///", alpha=0.85) ax.bar(x_meas, ffn_mb_list, width=bar_w, bottom=wo_mb_list, color=_FFN_COLOR, edgecolor="black", hatch="///", alpha=0.85) ax.bar(x_meas, meas_attn_mb, width=bar_w, bottom=bottoms_attn, color=_ATTN_COLOR, edgecolor="black", hatch="///", alpha=0.85) ax.set_xticks(x) ax.set_xticklabels(labels, fontsize=9) ax.set_ylabel("Comm per token per PE (MB, log)") ax.set_yscale("log") title = ( f"Per-PE communication at S_kv = {_HEADLINE_S_KV:,} tokens " f"(decode S_q=1, B=1; {_N_LAYERS} layers) — {source_tag}" ) if paired: title += "\n(simulator measured at S_kv = 8K; " \ "partial-score AR scaled ×128 to S_kv = 1M)" ax.set_title(title, fontsize=10) ax.grid(axis="y", ls=":", alpha=0.5, which="both") totals_ana = [wo_mb_list[i] + ffn_mb_list[i] + attn_mb[i] for i in range(n_cases)] ymax = max(totals_ana) if paired: ymax = max(ymax, max( (measured.get(c, 0) / (1 << 20)) for c in _CASES)) ax.set_ylim(top=ymax * 22) for i, c in enumerate(_CASES): ana_bytes = totals_ana[i] * (1 << 20) if paired: meas_bytes = measured.get(c, 0) top_y = max(totals_ana[i], meas_bytes / (1 << 20)) label = (f"ana: {_fmt_bytes(ana_bytes)}\n" f"sim: {_fmt_bytes(meas_bytes)}") else: top_y = totals_ana[i] label = _fmt_bytes(ana_bytes) ax.text(x[i], top_y * 1.5, label, ha="center", va="bottom", fontsize=8 if paired else 9, weight="bold", linespacing=1.05) # In-bar attention-time AR description (online-softmax (m,ℓ,O) # merge, partial-score AR, intra/inter-cube). Placed on the # analytical bar; with paired mode it lives on the solid bar so # the hatched measured bar stays uncluttered. if attn_mb[i] > 0: mid = bottoms_attn[i] + attn_mb[i] / 2 ax.text(x_ana[i], mid, _ATTN_DESC[c], ha="center", va="center", fontsize=6 if paired else 7, color="black", weight="bold") else: ax.text(x_ana[i], totals_ana[i] * 0.4, _ATTN_DESC[c], ha="center", fontsize=6.5 if paired else 7.5, color="grey", style="italic") legend_handles = [ mpatches.Patch(facecolor=_WO_COLOR, edgecolor="black", label=f"Wo AR (× {_N_LAYERS} layers)"), mpatches.Patch(facecolor=_FFN_COLOR, edgecolor="black", label=f"FFN AR (× {_N_LAYERS} layers)"), mpatches.Patch(facecolor=_ATTN_COLOR, edgecolor="black", label="Attn-time collective"), ] if paired: legend_handles.append( mpatches.Patch(facecolor="white", edgecolor="black", hatch="///", label="simulator-measured")) ax.legend(handles=legend_handles, loc="upper right", fontsize=8, framealpha=0.92) def main() -> Path: _OUT_DIR.mkdir(parents=True, exist_ok=True) # (a) Per-PE HBM budget — standalone PNG. 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" fig_b.savefig(out_b, dpi=150) plt.close(fig_b) print(f"wrote {out_b}") # (b) Combined 3-panel summary — HBM budget + KV memory + comm. fig = plt.figure(figsize=(22.0, 6.5)) gs = fig.add_gridspec(1, 3, width_ratios=[0.7, 1.6, 1.6], wspace=0.22) ax_b2 = fig.add_subplot(gs[0, 0]) ax_m = fig.add_subplot(gs[0, 1]) ax_c = fig.add_subplot(gs[0, 2]) _plot_budget(ax_b2) _plot_memory(ax_m) _plot_comm(ax_c) fig.tight_layout() out = _OUT_DIR / "gqa_4cases_summary.png" fig.savefig(out, dpi=150) plt.close(fig) print(f"wrote {out}") # (c) 2-panel companion (analytical only). fig2 = plt.figure(figsize=(18.0, 6.5)) gs2 = fig2.add_gridspec(1, 2, width_ratios=[1.0, 1.0], wspace=0.18) ax_m2 = fig2.add_subplot(gs2[0, 0]) ax_c2 = fig2.add_subplot(gs2[0, 1]) _plot_memory(ax_m2) _plot_comm(ax_c2, mode="analytical") fig2.tight_layout() out2 = _OUT_DIR / "gqa_4cases_memory_comm_analytical.png" fig2.savefig(out2, dpi=150) plt.close(fig2) print(f"wrote {out2}") # (d) 2-panel companion — analytical vs simulator-measured paired. fig3 = plt.figure(figsize=(19.0, 6.5)) gs3 = fig3.add_gridspec(1, 2, width_ratios=[1.0, 1.0], wspace=0.18) ax_m3 = fig3.add_subplot(gs3[0, 0]) ax_c3 = fig3.add_subplot(gs3[0, 1]) _plot_memory(ax_m3) _plot_comm(ax_c3, mode="paired") fig3.tight_layout() out3 = _OUT_DIR / "gqa_4cases_memory_comm_paired.png" fig3.savefig(out3, dpi=150) plt.close(fig3) print(f"wrote {out3}") # Paper-ready table to stdout. print() print(f" {'Case':<27} {'KV/tok·PE':>12} {'KV @ 1M':>11} " f"{'Max S_kv':>11} {'Comm @ 1M':>12}") print(" " + "-" * 80) for c in _CASES: kv_per_tok_pe = _KV_PER_TOK_BYTES / _DIVISOR[c] mem_1m = kv_per_pe_bytes(c, _HEADLINE_S_KV) max_s = max_s_kv(c) comm_1m = sum(per_token_bytes(c, _HEADLINE_S_KV)) label = _CASE_LABEL[c].replace(chr(10), " ") print(f" {label:<27} " f"{kv_per_tok_pe / 1024:>9.3f} KB " f"{_fmt_bytes(mem_1m):>11} " f"{_fmt_tokens(max_s):>11} " f"{_fmt_bytes(comm_1m):>12}") return out if __name__ == "__main__": main()