mukesh
694e0cc9b9
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>
500 lines
20 KiB
Python
500 lines
20 KiB
Python
"""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()
|