mukesh
65c365f858
The legacy panel names suggested batched serving semantics they never had — all four modeled a single user with KV sharded differently (C=1 single-cube; C=4 multi-cube Cube-SP), at toy dims (T_q=4, S_kv≤128). The single-KV-group C=8 panel + the new milestone-gqa-decode-4cases bench cover the meaningful comparisons; pytest regression already covers C=1/C=4 configurations end-to-end at richer scale. Changes: - milestone_gqa_headline.py: drop the 4 legacy panels; _PANELS now contains only single_kv_group_prefill_gqa_c8_p8. Update docstring. - tests/attention/test_milestone_gqa_headline.py: drop the 3 legacy- panel architectural tests (Ring-KV traffic, root-only decode write, per-CUBE distributed output) and test_decode_panels_use_real_gqa (no decode panels in this bench anymore). Equivalent properties are asserted in test_milestone_gqa_single_kv_group_prefill_panel.py (64 dma_writes, 896 ipcq_copy) and test_milestone_gqa_decode_4cases.py (1 dma_write at cube 6, 189 ipcq_copy). - tests/attention/test_milestone_gqa_single_kv_group_prefill_panel.py: drop test_existing_prefill_panel_runner_backward_compat (it exercised multi_user_prefill_gqa which no longer exists). - scripts/paper/paper_plot_gqa.py: replace the 4 legacy _LABELS entries with the single single_kv_group_prefill_gqa_c8_p8 label. - Regenerate 1H_milestone_output/gqa_headline/sweep.json from the new panel set. Verification: 9/9 tests pass across the 3 affected test files. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
107 lines
3.5 KiB
Python
107 lines
3.5 KiB
Python
"""Report harness: render GQA figures from ``gqa_latency.json``.
|
|
|
|
Isolated under ``scripts/paper/`` for the 1H codesign report only.
|
|
Reads the JSON emitted by ``paper_gqa_latency.py`` and writes two PNGs into
|
|
``docs/report/1H-codesign-paper/figures/``:
|
|
|
|
gqa_latency_by_panel.png end-to-end latency per panel
|
|
gqa_op_engine_breakdown.png op-counts + engine occupancy per panel
|
|
|
|
Run (after paper_gqa_latency.py):
|
|
python scripts/paper/paper_plot_gqa.py
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
from pathlib import Path
|
|
|
|
import matplotlib
|
|
|
|
matplotlib.use("Agg")
|
|
import matplotlib.pyplot as plt # noqa: E402
|
|
|
|
_FIG_DIR = Path(__file__).resolve().parents[2] / "docs" / "report" / "1H-codesign-paper" / "figures"
|
|
_IN_JSON = _FIG_DIR / "gqa_latency.json"
|
|
|
|
_LABELS = {
|
|
"single_kv_group_prefill_gqa_c8_p8":
|
|
"prefill\nC=8, P=8\n(single KV group)",
|
|
}
|
|
|
|
|
|
def _load() -> list[dict]:
|
|
return json.loads(_IN_JSON.read_text())["rows"]
|
|
|
|
|
|
def _plot_latency(rows: list[dict]) -> Path:
|
|
labels = [_LABELS[r["panel"]] for r in rows]
|
|
lat = [r["latency_ns"] for r in rows]
|
|
fig, ax = plt.subplots(figsize=(7.0, 4.0))
|
|
bars = ax.bar(labels, lat, color="#3b6ea5", width=0.6)
|
|
ax.set_ylabel("end-to-end latency (ns)")
|
|
ax.set_title("Fused GQA — end-to-end latency per panel")
|
|
ax.bar_label(bars, fmt="%.0f", padding=3, fontsize=9)
|
|
ax.grid(axis="y", ls=":", alpha=0.5)
|
|
ax.set_ylim(0, max(lat) * 1.15)
|
|
fig.tight_layout()
|
|
out = _FIG_DIR / "gqa_latency_by_panel.png"
|
|
fig.savefig(out, dpi=150)
|
|
plt.close(fig)
|
|
return out
|
|
|
|
|
|
def _plot_breakdown(rows: list[dict]) -> Path:
|
|
labels = [_LABELS[r["panel"]] for r in rows]
|
|
x = range(len(rows))
|
|
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(11.0, 4.2))
|
|
|
|
# Left: op-count breakdown (grouped bars).
|
|
keys = ["gemm_count", "ipcq_copy_count", "dma_read_count", "dma_write_count"]
|
|
disp = ["GEMM", "IPCQ copy", "DMA read", "DMA write"]
|
|
colors = ["#3b6ea5", "#c0504d", "#9bbb59", "#8064a2"]
|
|
w = 0.2
|
|
for i, (k, d, c) in enumerate(zip(keys, disp, colors)):
|
|
vals = [r["op_log_summary"][k] for r in rows]
|
|
ax1.bar([xi + (i - 1.5) * w for xi in x], vals, width=w, label=d, color=c)
|
|
ax1.set_xticks(list(x))
|
|
ax1.set_xticklabels(labels, fontsize=8)
|
|
ax1.set_ylabel("op count")
|
|
ax1.set_title("Per-panel op-count breakdown")
|
|
ax1.legend(fontsize=8)
|
|
ax1.grid(axis="y", ls=":", alpha=0.5)
|
|
|
|
# Right: engine occupancy (log scale — compute vs data movement).
|
|
engs = ["pe_gemm", "pe_math", "pe_dma"]
|
|
edisp = ["GEMM engine", "MATH engine", "DMA engine"]
|
|
ecolors = ["#3b6ea5", "#e8a33d", "#c0504d"]
|
|
w2 = 0.25
|
|
for i, (e, d, c) in enumerate(zip(engs, edisp, ecolors)):
|
|
vals = [max(r["engine_occupancy_ns"][e], 0.1) for r in rows]
|
|
ax2.bar([xi + (i - 1) * w2 for xi in x], vals, width=w2, label=d, color=c)
|
|
ax2.set_yscale("log")
|
|
ax2.set_xticks(list(x))
|
|
ax2.set_xticklabels(labels, fontsize=8)
|
|
ax2.set_ylabel("summed engine occupancy (ns, log)")
|
|
ax2.set_title("Compute vs data-movement occupancy")
|
|
ax2.legend(fontsize=8)
|
|
ax2.grid(axis="y", ls=":", alpha=0.5)
|
|
|
|
fig.suptitle("Fused GQA — where the work goes", fontsize=12)
|
|
fig.tight_layout()
|
|
out = _FIG_DIR / "gqa_op_engine_breakdown.png"
|
|
fig.savefig(out, dpi=150)
|
|
plt.close(fig)
|
|
return out
|
|
|
|
|
|
def main() -> None:
|
|
rows = _load()
|
|
p1 = _plot_latency(rows)
|
|
p2 = _plot_breakdown(rows)
|
|
print(f"wrote {p1}")
|
|
print(f"wrote {p2}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|