paper(gqa): add 6-case per-PE memory + comm summary chart

Analytical plot generator + PNG covering the 4 sharding cases from
slide 17 plus the 2 d_head-TP variants (Cases 4, 5). Three panels:
per-PE HBM budget breakdown (Wq + Wk + Wv + Wo + 4.24 GB KV
headroom), per-PE KV memory at S_kv=1M, and per-PE communication
per output token (decode). Numbers match slide 17 (max KV context
~7.1M for the 64-way sharded cases).

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
2026-06-16 16:10:49 -07:00
parent 6b6e29968a
commit 694e0cc9b9
2 changed files with 499 additions and 0 deletions
@@ -0,0 +1,499 @@
"""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
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"
)
# ── 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) -> None:
wo_mb = []
ffn_mb = []
attn_mb = []
for c in _CASES:
wo, ffn, attn = per_token_bytes(c, _HEADLINE_S_KV)
wo_mb.append(wo / (1 << 20))
ffn_mb.append(ffn / (1 << 20))
attn_mb.append(attn / (1 << 20))
labels = [_CASE_LABEL[c] for c in _CASES]
x = list(range(len(_CASES)))
ax.bar(x, wo_mb, color=_WO_COLOR, edgecolor="black")
ax.bar(x, ffn_mb, bottom=wo_mb, color=_FFN_COLOR, edgecolor="black")
bottoms_attn = [w + f for w, f in zip(wo_mb, ffn_mb)]
ax.bar(x, attn_mb, bottom=bottoms_attn,
color=_ATTN_COLOR, edgecolor="black")
ax.set_xticks(x)
ax.set_xticklabels(labels, fontsize=9)
ax.set_ylabel("Comm per token per PE (MB, log)")
ax.set_yscale("log")
ax.set_title(
f"Per-PE communication at S_kv = {_HEADLINE_S_KV:,} tokens "
f"(decode S_q=1, B=1; {_N_LAYERS} layers)",
fontsize=11,
)
ax.grid(axis="y", ls=":", alpha=0.5, which="both")
totals = [wo_mb[i] + ffn_mb[i] + attn_mb[i] for i in range(len(_CASES))]
ax.set_ylim(top=max(totals) * 10)
for i, c in enumerate(_CASES):
total_bytes = (wo_mb[i] + ffn_mb[i] + attn_mb[i]) * (1 << 20)
ax.text(x[i], totals[i] * 1.5, _fmt_bytes(total_bytes),
ha="center", fontsize=9, weight="bold")
if attn_mb[i] > 0:
mid = bottoms_attn[i] + attn_mb[i] / 2
ax.text(x[i], mid, _ATTN_DESC[c],
ha="center", va="center", fontsize=7,
color="black", weight="bold")
else:
ax.text(x[i], totals[i] * 0.4, _ATTN_DESC[c],
ha="center", fontsize=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"),
]
ax.legend(handles=legend_handles, loc="upper right", fontsize=8,
framealpha=0.92)
def main() -> Path:
_OUT_DIR.mkdir(parents=True, exist_ok=True)
fig = plt.figure(figsize=(21.0, 6.0))
gs = fig.add_gridspec(1, 3, width_ratios=[0.7, 1.6, 1.6], wspace=0.22)
ax_b = 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_b)
_plot_memory(ax_m)
_plot_comm(ax_c)
fig.suptitle(
f"GQA per-PE memory + communication — LLaMA-3.1-70B "
f"single-KV-head group (C={_C}, P={_P}, {_N_LAYERS} layers, "
f"S_kv = 1 M, FP16)",
fontsize=12, y=0.985,
)
fig.text(
0.5, 0.94,
"Cases ordered by memory · Case 1 (40 GB, no sharding) → "
"Cases 2-3 (5 GB, 1-axis sharded) → Cases 4-6 (640 MB, 2-axis sharded) "
"· Case 6 ★ = lowest comm among memory-feasible cases",
ha="center", fontsize=9.5, color="#444",
)
fig.tight_layout(rect=(0, 0, 1, 0.92))
out = _OUT_DIR / "gqa_4cases_summary.png"
fig.savefig(out, dpi=150)
plt.close(fig)
print(f"wrote {out}")
# 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()