gqa(prefill-4cases): add long-context prefill 4-cases comparative study
Mirrors the existing decode 4-cases study for prefill on the LLaMA-3.1-70B
single-KV-head group (h_q=8, h_kv=1, d_head=128) across 8 cubes × 8 PEs:
Case 1 Cube-SP × PE-TP → KV S_kv-Ring across cubes; Q/O T_q-split 64-way
Case 2 Cube-Repl × PE-TP → full KV per cube; only CUBE 0's P PEs split T_q
Case 3 Cube-Repl × PE-SP → full KV per cube; PEs SP on S_kv (intra-cube AR)
Case 4 Cube-SP × PE-SP → KV split 64-way + Q T_q-split across cubes
(Ring KV + per-cube intra-CUBE AR) ★ optimal
Cases 2 and 3 use Q-axis tiling (TILE_Q) so the per-PE scratch is bounded
by TILE_Q × TILE_S_KV regardless of T_q (the full Q tensor would be
G·T_q·d_head·2 bytes which scales linearly with T_q and blew the 1 MB
budget at T_q≥128 without tiling).
Per-panel try/except in run_sweep tolerates per-panel failures so
sweep.json always lands with whichever cases succeed plus a failures
list. Case 4 uses configure_sfr_intercube_ring with the snake submesh
(2,4) which installs both the Ring E/W lanes and the intra_* lanes
needed for the intra-CUBE reduce — single SFR config for all 4 cases.
The plot script generates 4 PNGs (latency, traffic, memory, parallelism)
into 1H_milestone_output/gqa/long_ctx/. The parallelism chart shows
total compute work (active_PE × T_q_per_PE × S_kv_processed) — exposes
Case 3's 8× redundancy vs. Cases 1/2/4 which all do unique work.
gqa_prefill_long_ctx_4cases.py exposes run_sweep() rather than a @bench
decorator — the umbrella bench in the next commit invokes it.
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,253 @@
|
||||
"""Comparative figures for milestone-gqa-prefill-long-ctx-4cases.
|
||||
|
||||
Mirror of paper_plot_gqa_decode_long_ctx_4cases but for the prefill
|
||||
variant. Reads sweep.json (emitted by ``kernbench run --bench
|
||||
milestone-gqa-prefill-long-ctx-4cases``) and writes four PNGs into
|
||||
``docs/report/1H-codesign-paper/figures/``:
|
||||
|
||||
gqa_prefill_long_ctx_4cases_latency.png end-to-end latency per case
|
||||
gqa_prefill_long_ctx_4cases_traffic.png ipcq/dma op-count breakdown
|
||||
gqa_prefill_long_ctx_4cases_memory.png per-PE KV bytes per case
|
||||
gqa_prefill_long_ctx_4cases_parallelism.png active-PE × S_local work load
|
||||
|
||||
Run (after the bench):
|
||||
GQA_PREFILL_LONG_CTX_4CASES_RUN=1 python -m kernbench.cli.main run \\
|
||||
--bench milestone-gqa-prefill-long-ctx-4cases --topology topology.yaml
|
||||
python scripts/paper/paper_plot_gqa_prefill_long_ctx_4cases.py
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
import matplotlib
|
||||
|
||||
matplotlib.use("Agg")
|
||||
import matplotlib.pyplot as plt # noqa: E402
|
||||
|
||||
_REPO_ROOT = Path(__file__).resolve().parents[2]
|
||||
# Sweep JSON + PNGs live together under the bench output dir.
|
||||
_FIG_DIR = (
|
||||
_REPO_ROOT / "src" / "kernbench" / "benches"
|
||||
/ "1H_milestone_output" / "gqa" / "long_ctx"
|
||||
)
|
||||
_SWEEP_JSON = _FIG_DIR / "sweep_prefill.json"
|
||||
|
||||
# Panel name → (short label, case ordinal for left-to-right plot order).
|
||||
_CASE_INFO = {
|
||||
"single_kv_group_prefill_long_ctx_gqa_cube_sp_pe_tp": (
|
||||
"Case 1\nCube-SP × PE-TP", 1),
|
||||
"single_kv_group_prefill_long_ctx_gqa_cube_repl_pe_tp": (
|
||||
"Case 2\nCube-Repl × PE-TP", 2),
|
||||
"single_kv_group_prefill_long_ctx_gqa_cube_repl_pe_sp": (
|
||||
"Case 3\nCube-Repl × PE-SP", 3),
|
||||
"single_kv_group_prefill_long_ctx_gqa_cube_sp_pe_sp": (
|
||||
"Case 4 ★\nCube-SP × PE-SP", 4),
|
||||
}
|
||||
|
||||
|
||||
def _load() -> list[dict]:
|
||||
return json.loads(_SWEEP_JSON.read_text())["rows"]
|
||||
|
||||
|
||||
def _sorted_by_case(rows: list[dict]) -> list[dict]:
|
||||
return sorted(rows, key=lambda r: _CASE_INFO[r["panel"]][1])
|
||||
|
||||
|
||||
def _plot_latency(rows: list[dict]) -> Path:
|
||||
rows = _sorted_by_case(rows)
|
||||
labels = [_CASE_INFO[r["panel"]][0] for r in rows]
|
||||
lat_us = [r["latency_ns"] / 1e3 for r in rows]
|
||||
colors = ["#888", "#888", "#888", "#3b6ea5"] # Case 4 highlighted
|
||||
fig, ax = plt.subplots(figsize=(8.0, 4.5))
|
||||
bars = ax.bar(labels, lat_us, color=colors, width=0.6)
|
||||
ax.set_ylabel("end-to-end latency (µs)")
|
||||
ax.set_title(
|
||||
"Long-context prefill 4-cases — end-to-end latency per case\n"
|
||||
"LLaMA-3.1-70B single-KV-head group (8 cubes × 8 PEs)"
|
||||
)
|
||||
ax.bar_label(bars, fmt="%.1f", padding=3, fontsize=9)
|
||||
ax.grid(axis="y", ls=":", alpha=0.5)
|
||||
ax.set_ylim(0, max(lat_us) * 1.15)
|
||||
fig.tight_layout()
|
||||
out = _FIG_DIR / "gqa_prefill_long_ctx_4cases_latency.png"
|
||||
fig.savefig(out, dpi=150)
|
||||
plt.close(fig)
|
||||
return out
|
||||
|
||||
|
||||
def _plot_traffic(rows: list[dict]) -> Path:
|
||||
rows = _sorted_by_case(rows)
|
||||
labels = [_CASE_INFO[r["panel"]][0] for r in rows]
|
||||
x = list(range(len(rows)))
|
||||
keys = ["ipcq_copy_count", "dma_read_count", "dma_write_count"]
|
||||
disp = ["IPCQ copy", "DMA read", "DMA write"]
|
||||
colors = ["#c0504d", "#9bbb59", "#8064a2"]
|
||||
w = 0.25
|
||||
fig, ax = plt.subplots(figsize=(9.0, 4.5))
|
||||
for i, (k, d, c) in enumerate(zip(keys, disp, colors)):
|
||||
vals = [r["op_log_summary"][k] for r in rows]
|
||||
ax.bar([xi + (i - 1) * w for xi in x], vals, width=w, label=d, color=c)
|
||||
ax.set_xticks(list(x))
|
||||
ax.set_xticklabels(labels, fontsize=9)
|
||||
ax.set_ylabel("op count")
|
||||
ax.set_title("Long-context prefill 4-cases — op-count breakdown per case")
|
||||
ax.legend(fontsize=9)
|
||||
ax.grid(axis="y", ls=":", alpha=0.5)
|
||||
fig.tight_layout()
|
||||
out = _FIG_DIR / "gqa_prefill_long_ctx_4cases_traffic.png"
|
||||
fig.savefig(out, dpi=150)
|
||||
plt.close(fig)
|
||||
return out
|
||||
|
||||
|
||||
def _s_local_per_pe(panel: str, *, S_kv: int, C: int, P: int) -> int:
|
||||
"""S_local (token count) each PE attends over locally.
|
||||
|
||||
Encodes the cube/pe sharding axes from the panel name:
|
||||
cube_sp_pe_tp (Case 1): S_kv / C (pe=replicate within cube)
|
||||
cube_repl_pe_tp (Case 2): S_kv (full S_kv per active PE)
|
||||
cube_repl_pe_sp (Case 3): S_kv / P (pe=row_wise within cube)
|
||||
cube_sp_pe_sp (Case 4): S_kv / (C·P) (★ 64-way split)
|
||||
"""
|
||||
S_per_cube = S_kv if "cube_repl" in panel else S_kv // C
|
||||
return S_per_cube // P if "pe_sp" in panel else S_per_cube
|
||||
|
||||
|
||||
def _active_pe_count(panel: str, *, C: int, P: int) -> int:
|
||||
"""Number of PEs doing non-idle attention work.
|
||||
|
||||
Prefill T_q≫1 means PE-TP is *useful* (not wasted as in decode):
|
||||
cube_sp_pe_tp (Case 1): C·P (T_q sharded across all 64 ranks)
|
||||
cube_repl_pe_tp (Case 2): P (CUBE 0 only; T_q sharded across its P PEs)
|
||||
cube_repl_pe_sp (Case 3): C·P (all PEs busy but cubes redundant)
|
||||
cube_sp_pe_sp (Case 4): C·P (all 64 PEs doing unique work)
|
||||
"""
|
||||
if "cube_repl" in panel and "pe_tp" in panel:
|
||||
return P
|
||||
return C * P
|
||||
|
||||
|
||||
def _kv_bytes_per_pe(panel: str, *, S_kv: int, h_kv: int,
|
||||
d_head: int, C: int, P: int) -> int:
|
||||
"""KV bytes a single PE references (K + V, f16, 2 B/elem)."""
|
||||
s_local = _s_local_per_pe(panel, S_kv=S_kv, C=C, P=P)
|
||||
return 2 * s_local * h_kv * d_head * 2
|
||||
|
||||
|
||||
def _plot_memory(rows: list[dict]) -> Path:
|
||||
"""Per-PE KV bytes — Case 4 wins (64-way split)."""
|
||||
rows = _sorted_by_case(rows)
|
||||
labels = [_CASE_INFO[r["panel"]][0] for r in rows]
|
||||
mib_per_pe = [
|
||||
_kv_bytes_per_pe(
|
||||
r["panel"], S_kv=r["S_kv"], h_kv=r["h_kv"],
|
||||
d_head=r["d_head"], C=r["C"], P=r["P"],
|
||||
) / (1024 * 1024)
|
||||
for r in rows
|
||||
]
|
||||
colors = ["#888", "#c0504d", "#888", "#3b6ea5"]
|
||||
fig, ax = plt.subplots(figsize=(8.0, 4.5))
|
||||
bars = ax.bar(labels, mib_per_pe, color=colors, width=0.6)
|
||||
ax.set_ylabel("KV bytes per PE (MiB, K + V, f16)")
|
||||
ax.set_title(
|
||||
"Long-context prefill 4-cases — KV memory per PE\n"
|
||||
"(one KV-head group; per-layer, full S_kv state)"
|
||||
)
|
||||
ax.bar_label(bars, fmt="%.3f", padding=3, fontsize=9)
|
||||
ax.grid(axis="y", ls=":", alpha=0.5)
|
||||
ax.set_ylim(0, max(mib_per_pe) * 1.15)
|
||||
fig.tight_layout()
|
||||
out = _FIG_DIR / "gqa_prefill_long_ctx_4cases_memory.png"
|
||||
fig.savefig(out, dpi=150)
|
||||
plt.close(fig)
|
||||
return out
|
||||
|
||||
|
||||
def _t_q_per_pe(panel: str, *, T_q: int, C: int, P: int) -> int:
|
||||
"""T_q row count each active PE computes attention for.
|
||||
|
||||
cube_sp_pe_tp (Case 1): T_q / (C·P) (T_q sharded across all 64 ranks)
|
||||
cube_repl_pe_tp (Case 2): T_q / P (T_q sharded across CUBE 0's P PEs)
|
||||
cube_repl_pe_sp (Case 3): T_q (Q replicated on every PE)
|
||||
cube_sp_pe_sp (Case 4): T_q / C (Q sharded by cube, replicated within)
|
||||
"""
|
||||
if "cube_sp" in panel and "pe_tp" in panel:
|
||||
return T_q // (C * P)
|
||||
if "cube_repl" in panel and "pe_tp" in panel:
|
||||
return T_q // P
|
||||
if "cube_repl" in panel and "pe_sp" in panel:
|
||||
return T_q
|
||||
return T_q // C # cube_sp_pe_sp
|
||||
|
||||
|
||||
def _s_kv_processed_per_pe(panel: str, *, S_kv: int, C: int, P: int) -> int:
|
||||
"""S_kv tokens each PE actually processes attention over.
|
||||
|
||||
Differs from ``_s_local_per_pe`` (OWNED KV bytes): for cases with
|
||||
a Ring (Case 1, 4) each PE sees C ring steps so processes more
|
||||
tokens than it locally owns.
|
||||
|
||||
cube_sp_pe_tp (Case 1): S_kv (Ring + pe=replicate within cube)
|
||||
cube_repl_pe_tp (Case 2): S_kv (full KV per PE)
|
||||
cube_repl_pe_sp (Case 3): S_kv / P (pe=row_wise; no Ring)
|
||||
cube_sp_pe_sp (Case 4): S_kv / P (Ring restores full S_kv/P per PE)
|
||||
"""
|
||||
if "cube_sp" in panel and "pe_tp" in panel:
|
||||
return S_kv
|
||||
if "cube_repl" in panel and "pe_tp" in panel:
|
||||
return S_kv
|
||||
return S_kv // P # both pe_sp cases
|
||||
|
||||
|
||||
def _plot_parallelism(rows: list[dict]) -> Path:
|
||||
"""Total compute work (PE × T_q × S_kv token-pairs) — exposes Case 3's
|
||||
redundancy. Cases 1, 2, 4 all do the same total work (correct
|
||||
attention over T_q × S_kv); Case 3 does C× more (cubes redundantly
|
||||
repeat the same compute because K/V is cube-replicated).
|
||||
"""
|
||||
rows = _sorted_by_case(rows)
|
||||
labels = [_CASE_INFO[r["panel"]][0] for r in rows]
|
||||
total_work = [
|
||||
_active_pe_count(r["panel"], C=r["C"], P=r["P"])
|
||||
* _t_q_per_pe(r["panel"], T_q=r["T_q"], C=r["C"], P=r["P"])
|
||||
* _s_kv_processed_per_pe(
|
||||
r["panel"], S_kv=r["S_kv"], C=r["C"], P=r["P"],
|
||||
)
|
||||
for r in rows
|
||||
]
|
||||
colors = ["#888", "#888", "#c0504d", "#3b6ea5"] # Case 3 red, Case 4 highlighted
|
||||
fig, ax = plt.subplots(figsize=(8.0, 4.5))
|
||||
bars = ax.bar(labels, total_work, color=colors, width=0.6)
|
||||
ax.set_ylabel(
|
||||
"total compute (PE × T_q × S_kv token-pairs; lower ⇒ less wasted work)"
|
||||
)
|
||||
ax.set_title(
|
||||
"Long-context prefill 4-cases — total compute work across active PEs\n"
|
||||
"(Case 3 replicates K/V across 8 cubes ⇒ 8× redundant compute)"
|
||||
)
|
||||
ax.bar_label(bars, fmt="%d", padding=3, fontsize=9)
|
||||
ax.grid(axis="y", ls=":", alpha=0.5)
|
||||
ax.set_ylim(0, max(total_work) * 1.15)
|
||||
fig.tight_layout()
|
||||
out = _FIG_DIR / "gqa_prefill_long_ctx_4cases_parallelism.png"
|
||||
fig.savefig(out, dpi=150)
|
||||
plt.close(fig)
|
||||
return out
|
||||
|
||||
|
||||
def main() -> None:
|
||||
rows = _load()
|
||||
_FIG_DIR.mkdir(parents=True, exist_ok=True)
|
||||
p1 = _plot_latency(rows)
|
||||
p2 = _plot_traffic(rows)
|
||||
p3 = _plot_memory(rows)
|
||||
p4 = _plot_parallelism(rows)
|
||||
print(f"wrote {p1}")
|
||||
print(f"wrote {p2}")
|
||||
print(f"wrote {p3}")
|
||||
print(f"wrote {p4}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user