paper(gqa): derive (m,l,O) AR per-PE cost from kernel topology

Replaces the unsourced 32 KB/layer placeholder in the analytical chart
with T*(N-1)/N where T = h_q * S_q * (d_head*2 + 8) ~ 2.1 KB
(per-PE (m, l, O) payload) and N is the participant count of the
reduce-only chain per stage. Cases 2/3/6 analytical bars now match
the simulator-measured numbers within ~10% (was 6-30x over).

Also:
- Bumps the measurement S_kv from 8 K to 64 K to verify Cases 1, 2, 3, 6
  are genuinely S_kv-independent and Cases 4, 5 scale linearly.
- In-bar descriptor on the paired chart is now wrapped to fit the bar
  width, black, no background box, and centered between the analytical
  and measured bars so the label applies to both.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

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.

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",

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)  [*]",