gqa(decode): add d_head-TP kernels + measurement runner + figure generators
Two new long-ctx decode attention kernels for the d_head-TP sharding
variants the 6-case chart predicts:
· _gqa_attention_decode_long_ctx_cube_sp_pe_tp_dhead.py
Cube-SP × PE-TP-d_head (Case 4 in chart). Per cube holds
S_kv/C tokens of full d_head; per PE holds same tokens but
only d_head/P dims. Partial Q·Kᵀ scores reduced intra-cube
before softmax; outer (m,ℓ,O) merge two-phase (intra+inter).
· _gqa_attention_decode_long_ctx_cube_tp_dhead_pe_sp.py
Cube-TP-d_head × PE-SP (Case 5). Per cube holds full S_kv
with only d_head/C dims; per PE holds S_kv/P of those dims.
Partial scores reduced inter-cube (UCIe) before softmax.
Sweep dispatch (gqa_decode_long_ctx_4cases.py) extended with two
new panels so the milestone-1h-gqa sweep covers all 6 cases.
Smoke test scripts/verify_case4_dhead_tp.py runs Cases 4/5/6 at
S_kv=2K to validate the kernels load and execute end-to-end.
Plus the figure-generation toolchain that produced the committed
PNGs in the prior commit (dd3337f):
· paper_plot_gqa_4cases_summary.py - 3-panel summary +
2-panel (analytical / paired-measured) chart generator.
_plot_comm now takes mode="analytical" | "paired".
· paper_plot_gqa_kv_sharding_diagram.py - 6-case 2-D KV-tensor
diagram + companion comparison-table PNG.
· measure_gqa_decode_placement_comm.py - runs all 6 kernels at
S_kv=8K, sums actual IPCQ-copy bytes from engine.op_log,
scales partial-score AR ×128 to S_kv=1M, writes
gqa_3cases_measured_comm.json (committed).
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
@@ -114,6 +114,7 @@ Output PNG:
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
import matplotlib.patches as mpatches
|
||||
@@ -240,6 +241,22 @@ _OUT_DIR = (
|
||||
/ "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 ────────────────────────────────────────────────────────
|
||||
@@ -388,51 +405,114 @@ def _plot_memory(ax) -> None:
|
||||
ax.legend(loc="upper right", fontsize=9)
|
||||
|
||||
|
||||
def _plot_comm(ax) -> None:
|
||||
wo_mb = []
|
||||
ffn_mb = []
|
||||
attn_mb = []
|
||||
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.append(wo / (1 << 20))
|
||||
ffn_mb.append(ffn / (1 << 20))
|
||||
wo_mb_list.append(wo / (1 << 20))
|
||||
ffn_mb_list.append(ffn / (1 << 20))
|
||||
attn_mb.append(attn / (1 << 20))
|
||||
|
||||
labels = [_CASE_LABEL[c] for c in _CASES]
|
||||
x = list(range(len(_CASES)))
|
||||
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"
|
||||
|
||||
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,
|
||||
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")
|
||||
ax.set_title(
|
||||
title = (
|
||||
f"Per-PE communication at S_kv = {_HEADLINE_S_KV:,} tokens "
|
||||
f"(decode S_q=1, B=1; {_N_LAYERS} layers)",
|
||||
fontsize=11,
|
||||
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 = [wo_mb[i] + ffn_mb[i] + attn_mb[i] for i in range(len(_CASES))]
|
||||
ax.set_ylim(top=max(totals) * 10)
|
||||
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):
|
||||
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")
|
||||
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[i], mid, _ATTN_DESC[c],
|
||||
ha="center", va="center", fontsize=7,
|
||||
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[i], totals[i] * 0.4, _ATTN_DESC[c],
|
||||
ha="center", fontsize=7.5, color="grey", style="italic")
|
||||
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",
|
||||
@@ -442,40 +522,67 @@ def _plot_comm(ax) -> None:
|
||||
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)
|
||||
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])
|
||||
|
||||
# (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.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))
|
||||
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} "
|
||||
|
||||
Reference in New Issue
Block a user