diff --git a/docs/report/1H-codesign-paper/figures/gqa_long_ctx_6cases_memory_comm_analytical.png b/docs/report/1H-codesign-paper/figures/gqa_long_ctx_6cases_memory_comm_analytical.png index 3241151..41ff5db 100644 Binary files a/docs/report/1H-codesign-paper/figures/gqa_long_ctx_6cases_memory_comm_analytical.png and b/docs/report/1H-codesign-paper/figures/gqa_long_ctx_6cases_memory_comm_analytical.png differ diff --git a/docs/report/1H-codesign-paper/figures/gqa_long_ctx_6cases_memory_comm_paired.png b/docs/report/1H-codesign-paper/figures/gqa_long_ctx_6cases_memory_comm_paired.png index 3e16e1f..8e378c0 100644 Binary files a/docs/report/1H-codesign-paper/figures/gqa_long_ctx_6cases_memory_comm_paired.png and b/docs/report/1H-codesign-paper/figures/gqa_long_ctx_6cases_memory_comm_paired.png differ diff --git a/docs/report/1H-codesign-paper/figures/gqa_long_ctx_6cases_summary.png b/docs/report/1H-codesign-paper/figures/gqa_long_ctx_6cases_summary.png index 2a1779b..7e9893e 100644 Binary files a/docs/report/1H-codesign-paper/figures/gqa_long_ctx_6cases_summary.png and b/docs/report/1H-codesign-paper/figures/gqa_long_ctx_6cases_summary.png differ diff --git a/scripts/paper/measure_gqa_decode_placement_comm.py b/scripts/paper/measure_gqa_decode_placement_comm.py index 2a9ac87..60a82a5 100644 --- a/scripts/paper/measure_gqa_decode_placement_comm.py +++ b/scripts/paper/measure_gqa_decode_placement_comm.py @@ -51,7 +51,7 @@ from kernbench.topology.builder import resolve_topology _C = 8 _P = 8 _N_LAYERS = 80 -_S_KV_MEAS = 8 * 1024 # per-run simulator S_kv (small — 1/128th of headline) +_S_KV_MEAS = 64 * 1024 # per-run simulator S_kv (1/16th of headline) _S_KV_HEADLINE = 1 << 20 # 1 Mi tokens, the chart's headline S_kv # Per-cube S_kv share for the d_head-TP partial-score AR cost. diff --git a/scripts/paper/paper_plot_gqa_4cases_summary.py b/scripts/paper/paper_plot_gqa_4cases_summary.py index b5b9a37..584ea46 100644 --- a/scripts/paper/paper_plot_gqa_4cases_summary.py +++ b/scripts/paper/paper_plot_gqa_4cases_summary.py @@ -115,6 +115,7 @@ Output PNG: from __future__ import annotations import json +import textwrap from pathlib import Path import matplotlib.patches as mpatches @@ -187,9 +188,36 @@ _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 +# (m, ℓ, O) AR cost per layer, derived from the kernel topology. +# +# Per-PE payload T for one (m, ℓ, O) merge step: +# O — h_q · S_q · d_head · 2 bytes (FP16) +# m — h_q · S_q · 4 bytes (FP32) +# ℓ — h_q · S_q · 4 bytes (FP32) +# T = h_q · S_q · (d_head · 2 + 8) ≈ 2.1 KB +# +# Reduce algorithm: the decode kernels use hierarchical reduce-only +# (chain / tree) rather than ring all-reduce, because the merged result +# only needs to land on the cube that runs the downstream Wo gemm — +# not on every PE. For a chain of N participants the total traffic is +# (N-1)·T and the per-PE average is (N-1)/N · T. +# +# The previous version used a hard-coded 32 KB / layer placeholder which +# overestimated the per-PE cost by ~6× (and ~3× even under a hypothetical +# ring-AR assumption). See `gqa_long_ctx_6cases_measured_comm.json` for +# the measured numbers this matches against. +_MLO_PAYLOAD_BYTES = _H_Q * _S_Q * (_D_HEAD * _BYTES_PER_ELEM + 8) + + +def _reduce_chain_per_pe_bytes(n_participants: int) -> int: + """Per-PE average bytes for a single-stage chain/tree reduce of T.""" + if n_participants <= 1: + return 0 + return _MLO_PAYLOAD_BYTES * (n_participants - 1) // n_participants + + +_MLO_INTRA_BYTES_PER_LAYER = _reduce_chain_per_pe_bytes(_P) # PE-axis +_MLO_INTER_BYTES_PER_LAYER = _reduce_chain_per_pe_bytes(_C) # cube-axis # Cases — renumbered in MEMORY-DESCENDING ORDER (left to right): # Case 1 (40 GB) : no sharding @@ -225,11 +253,11 @@ _CASE_COLOR = { } _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", + 2: "online-softmax (m,ℓ,O) — inter-cube", + 3: "online-softmax (m,ℓ,O) — intra-cube", + 4: "partial scores + (m,ℓ,O) merge (d_head-TP)", + 5: "partial scores + (m,ℓ,O) merge (d_head-TP)", + 6: "online-softmax (m,ℓ,O) — intra + inter", } _WO_COLOR = "#9EC5E8" @@ -498,21 +526,25 @@ def _plot_comm(ax, *, mode: str = "analytical") -> None: 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. + # Attention-time AR descriptor — placed at the attention-segment + # midpoint, centered between the analytical and measured bars so + # it visually spans both. Coloured the same as the attention bar + # (no background box) so it lives within the case's bar zone. + # Wrapped narrow so each line fits inside the bar-pair width. + wrapped = textwrap.fill(_ATTN_DESC[c], width=14) if attn_mb[i] > 0: mid = bottoms_attn[i] + attn_mb[i] / 2 - ax.text(x_ana[i], mid, _ATTN_DESC[c], + ax.text(x[i], mid, wrapped, ha="center", va="center", - fontsize=6 if paired else 7, - color="black", weight="bold") + fontsize=7 if paired else 8, + color="black", weight="bold", + linespacing=1.0) else: - ax.text(x_ana[i], totals_ana[i] * 0.4, _ATTN_DESC[c], + ax.text(x[i], totals_ana[i] * 0.4, wrapped, ha="center", - fontsize=6.5 if paired else 7.5, - color="grey", style="italic") + fontsize=7 if paired else 8, + color="grey", style="italic", + linespacing=1.0) legend_handles = [ mpatches.Patch(facecolor=_WO_COLOR, edgecolor="black", diff --git a/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/gqa_long_ctx_6cases_measured_comm.json b/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/gqa_long_ctx_6cases_measured_comm.json index e527aea..c37d1d6 100644 --- a/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/gqa_long_ctx_6cases_measured_comm.json +++ b/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/gqa_long_ctx_6cases_measured_comm.json @@ -1,5 +1,5 @@ { - "S_kv_measured": 8192, + "S_kv_measured": 65536, "S_kv_headline": 1048576, "n_layers": 80, "num_pes": 64, @@ -31,19 +31,19 @@ }, "4": { "label": "Case 4 (Cube-SP x PE-TP d_head)", - "total_ipcq_bytes_one_layer": 1851136, - "per_pe_partial_score_bytes_one_layer": 16384, - "per_pe_mlo_bytes_one_layer": 12540, - "per_pe_attn_bytes_per_token_at_1M": 168775360, - "per_pe_total_bytes_per_token_at_1M": 170086080 + "total_ipcq_bytes_one_layer": 14696192, + "per_pe_partial_score_bytes_one_layer": 131072, + "per_pe_mlo_bytes_one_layer": 98556, + "per_pe_attn_bytes_per_token_at_1M": 175656640, + "per_pe_total_bytes_per_token_at_1M": 176967360 }, "5": { "label": "Case 5 (Cube-TP d_head x PE-SP)", - "total_ipcq_bytes_one_layer": 1851136, - "per_pe_partial_score_bytes_one_layer": 16384, - "per_pe_mlo_bytes_one_layer": 12540, - "per_pe_attn_bytes_per_token_at_1M": 168775360, - "per_pe_total_bytes_per_token_at_1M": 170086080 + "total_ipcq_bytes_one_layer": 14696192, + "per_pe_partial_score_bytes_one_layer": 131072, + "per_pe_mlo_bytes_one_layer": 98556, + "per_pe_attn_bytes_per_token_at_1M": 175656640, + "per_pe_total_bytes_per_token_at_1M": 176967360 }, "6": { "label": "Case 6 (Cube-SP x PE-SP) [*]", diff --git a/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/gqa_long_ctx_6cases_memory_comm_analytical.png b/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/gqa_long_ctx_6cases_memory_comm_analytical.png index 3241151..41ff5db 100644 Binary files a/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/gqa_long_ctx_6cases_memory_comm_analytical.png and b/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/gqa_long_ctx_6cases_memory_comm_analytical.png differ diff --git a/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/gqa_long_ctx_6cases_memory_comm_paired.png b/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/gqa_long_ctx_6cases_memory_comm_paired.png index 3e16e1f..8e378c0 100644 Binary files a/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/gqa_long_ctx_6cases_memory_comm_paired.png and b/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/gqa_long_ctx_6cases_memory_comm_paired.png differ diff --git a/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/gqa_long_ctx_6cases_summary.png b/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/gqa_long_ctx_6cases_summary.png index 2a1779b..7e9893e 100644 Binary files a/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/gqa_long_ctx_6cases_summary.png and b/src/kernbench/benches/1H_milestone_output/gqa/long_ctx/gqa_long_ctx_6cases_summary.png differ