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>

scripts/paper/measure_gqa_decode_placement_comm.py

@@ -0,0 +1,362 @@
+"""Measured comm overlay for all 6 GQA decode KV placements.
+
+Runs the simulator for Cases 1-6 (the same 6 placements the chart in
+paper_plot_gqa_4cases_summary.py covers analytically) at S_kv = 64 K,
+sums actual IPCQ-copy bytes from the engine op_log, projects to
+per-token (x80 layers), scales the partial-score-AR component of
+Cases 4/5 from S_kv = 64 K -> S_kv = 1 M (linear in S_kv; other
+cases are S_kv-independent), adds the constant Wo + FFN AR
+(1.25 MB / token), and writes the result to JSON for
+paper_plot_gqa_4cases_summary.py to overlay on the analytical bars.
+
+Single layer of decode attention only — the projection × 80 takes
+the per-layer measurement to a per-token total.
+
+Usage:
+  python scripts/paper/measure_gqa_decode_placement_comm.py
+"""
+from __future__ import annotations
+
+import json
+import os
+import sys
+from pathlib import Path
+
+from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_repl_pe_sp import (
+    gqa_attention_decode_long_ctx_cube_repl_pe_sp_kernel as _case3_kernel,
+)
+from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_repl_pe_tp import (
+    gqa_attention_decode_long_ctx_cube_repl_pe_tp_kernel as _case1_kernel,
+)
+from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_sp_pe_sp import (
+    gqa_attention_decode_long_ctx_cube_sp_pe_sp_kernel as _case6_kernel,
+)
+from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_sp_pe_tp import (
+    gqa_attention_decode_long_ctx_cube_sp_pe_tp_kernel as _case2_kernel,
+)
+from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_sp_pe_tp_dhead import (
+    gqa_attention_decode_long_ctx_cube_sp_pe_tp_dhead_kernel as _case4_kernel,
+)
+from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_tp_dhead_pe_sp import (
+    gqa_attention_decode_long_ctx_cube_tp_dhead_pe_sp_kernel as _case5_kernel,
+)
+from kernbench.benches.gqa_helpers.shared._gqa_panel_helpers import _ccl_cfg
+from kernbench.ccl.sfr_config import configure_sfr_intercube_multisip
+from kernbench.policy.placement.dp import DPPolicy
+from kernbench.runtime_api.bench_runner import run_bench
+from kernbench.runtime_api.types import resolve_device
+from kernbench.sim_engine.engine import GraphEngine
+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_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.
+#   Case 4 (Cube-SP × PE-TP_dhead): per-cube = S_kv / C
+#   Case 5 (Cube-TP_dhead × PE-SP): per-cube = S_kv (KV replicated across
+#       cubes for the cube-axis d_head-TP), so partial-score AR scales
+#       with the full S_kv.
+# Headline / measured per-cube ratios give the partial-score-AR scale-up
+# from S_kv = 64 K to S_kv = 1 M. (m,ℓ,O) AR is S_kv-independent.
+_PARTIAL_SCORE_SCALE = _S_KV_HEADLINE / _S_KV_MEAS   # = 16
+
+# Per-token Wo + FFN AR (constant across all cases, comes from the
+# attn-output and FFN-down all-reduces NOT measured by the attention-
+# only kernel run here).
+_WO_PER_LAYER_BYTES = 8 * 1024
+_FFN_PER_LAYER_BYTES = 8 * 1024
+_WO_FFN_PER_TOKEN_BYTES = (
+    (_WO_PER_LAYER_BYTES + _FFN_PER_LAYER_BYTES) * _N_LAYERS
+)   # 1.25 MB
+
+# Total PE count in one KV-head group — average per-PE comm = total / N.
+_NUM_PES = _C * _P
+
+# Total partial-score-AR slice produced by the attention compute when
+# d_head is sharded. Used to split measured IPCQ traffic into the
+# S_kv-scaling component (partial scores) vs the S_kv-independent
+# component ((m,ℓ,O) merge). Same as the analytical formula in
+# paper_plot_gqa_4cases_summary.py: h_q · S_q · per_cube_S_kv · 2 bytes.
+_H_Q = 8
+_S_Q = 1
+_BYTES_PER_ELEM = 2
+
+_PARAMS = dict(C=_C, P=_P, T_q=_S_Q, S_kv=_S_KV_MEAS,
+               d_head=128, h_q=_H_Q, h_kv=1)
+
+
+def _bench_fn_case1(ctx):
+    """Case 1: Cube-Repl x PE-repl (PE-TP doesn't shard KV)."""
+    p = _PARAMS
+    configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
+    dp = DPPolicy(cube="replicate", pe="replicate",
+                  num_cubes=p["C"], num_pes=p["P"])
+    q = ctx.zeros((p["T_q"], p["h_q"] * p["d_head"]),
+                  dtype="f16", dp=dp, name="q_c1")
+    k = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]),
+                  dtype="f16", dp=dp, name="k_c1")
+    v = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]),
+                  dtype="f16", dp=dp, name="v_c1")
+    o = ctx.empty((p["T_q"], p["h_q"] * p["d_head"]),
+                  dtype="f16", dp=dp, name="o_c1")
+    ctx.launch("case1_repl_repl", _case1_kernel,
+               q, k, v, o,
+               p["T_q"], p["S_kv"], p["h_q"], p["h_kv"],
+               p["d_head"], p["C"], p["P"],
+               _auto_dim_remap=False)
+
+
+def _bench_fn_case2(ctx):
+    """Case 2: Cube-SP x PE-repl (PE-TP doesn't shard KV)."""
+    p = _PARAMS
+    configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
+    dp_full = DPPolicy(cube="replicate", pe="replicate",
+                       num_cubes=p["C"], num_pes=p["P"])
+    dp_kv = DPPolicy(cube="row_wise", pe="replicate",
+                     num_cubes=p["C"], num_pes=p["P"])
+    q = ctx.zeros((p["T_q"], p["h_q"] * p["d_head"]),
+                  dtype="f16", dp=dp_full, name="q_c2")
+    k = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]),
+                  dtype="f16", dp=dp_kv, name="k_c2")
+    v = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]),
+                  dtype="f16", dp=dp_kv, name="v_c2")
+    o = ctx.empty((p["T_q"], p["h_q"] * p["d_head"]),
+                  dtype="f16", dp=dp_full, name="o_c2")
+    ctx.launch("case2_sp_repl", _case2_kernel,
+               q, k, v, o,
+               p["T_q"], p["S_kv"], p["h_q"], p["h_kv"],
+               p["d_head"], p["C"], p["P"],
+               _auto_dim_remap=False)
+
+
+def _bench_fn_case3(ctx):
+    """Case 3: Cube-Repl x PE-SP."""
+    p = _PARAMS
+    configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
+    dp_full = DPPolicy(cube="replicate", pe="replicate",
+                       num_cubes=p["C"], num_pes=p["P"])
+    dp_kv = DPPolicy(cube="replicate", pe="row_wise",
+                     num_cubes=p["C"], num_pes=p["P"])
+    q = ctx.zeros((p["T_q"], p["h_q"] * p["d_head"]),
+                  dtype="f16", dp=dp_full, name="q_c3")
+    k = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]),
+                  dtype="f16", dp=dp_kv, name="k_c3")
+    v = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]),
+                  dtype="f16", dp=dp_kv, name="v_c3")
+    o = ctx.empty((p["T_q"], p["h_q"] * p["d_head"]),
+                  dtype="f16", dp=dp_full, name="o_c3")
+    ctx.launch("case3_repl_sp", _case3_kernel,
+               q, k, v, o,
+               p["T_q"], p["S_kv"], p["h_q"], p["h_kv"],
+               p["d_head"], p["C"], p["P"],
+               _auto_dim_remap=False)
+
+
+def _bench_fn_case4(ctx):
+    p = _PARAMS
+    configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
+    dp_full = DPPolicy(cube="replicate", pe="column_wise",
+                       num_cubes=p["C"], num_pes=p["P"])
+    dp_kv = DPPolicy(cube="row_wise", pe="column_wise",
+                     num_cubes=p["C"], num_pes=p["P"])
+    q = ctx.zeros((p["T_q"], p["h_q"] * p["d_head"]),
+                  dtype="f16", dp=dp_full, name="q_c4")
+    k = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]),
+                  dtype="f16", dp=dp_kv, name="k_c4")
+    v = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]),
+                  dtype="f16", dp=dp_kv, name="v_c4")
+    o = ctx.empty((p["T_q"], p["h_q"] * p["d_head"]),
+                  dtype="f16", dp=dp_full, name="o_c4")
+    ctx.launch("case4_dhead_tp", _case4_kernel,
+               q, k, v, o,
+               p["T_q"], p["S_kv"], p["h_q"], p["h_kv"],
+               p["d_head"], p["C"], p["P"],
+               _auto_dim_remap=False)
+
+
+def _bench_fn_case5(ctx):
+    p = _PARAMS
+    configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
+    dp_q = DPPolicy(cube="column_wise", pe="replicate",
+                    num_cubes=p["C"], num_pes=p["P"])
+    dp_kv = DPPolicy(cube="column_wise", pe="row_wise",
+                     num_cubes=p["C"], num_pes=p["P"])
+    q = ctx.zeros((p["T_q"], p["h_q"] * p["d_head"]),
+                  dtype="f16", dp=dp_q, name="q_c5")
+    k = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]),
+                  dtype="f16", dp=dp_kv, name="k_c5")
+    v = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]),
+                  dtype="f16", dp=dp_kv, name="v_c5")
+    o = ctx.empty((p["T_q"], p["h_q"] * p["d_head"]),
+                  dtype="f16", dp=dp_q, name="o_c5")
+    ctx.launch("case5_dhead_tp_inter", _case5_kernel,
+               q, k, v, o,
+               p["T_q"], p["S_kv"], p["h_q"], p["h_kv"],
+               p["d_head"], p["C"], p["P"],
+               _auto_dim_remap=False)
+
+
+def _bench_fn_case6(ctx):
+    p = _PARAMS
+    configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
+    dp_full = DPPolicy(cube="replicate", pe="replicate",
+                       num_cubes=p["C"], num_pes=p["P"])
+    dp_kv = DPPolicy(cube="row_wise", pe="row_wise",
+                     num_cubes=p["C"], num_pes=p["P"])
+    q = ctx.zeros((p["T_q"], p["h_q"] * p["d_head"]),
+                  dtype="f16", dp=dp_full, name="q_c6")
+    k = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]),
+                  dtype="f16", dp=dp_kv, name="k_c6")
+    v = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]),
+                  dtype="f16", dp=dp_kv, name="v_c6")
+    o = ctx.empty((p["T_q"], p["h_q"] * p["d_head"]),
+                  dtype="f16", dp=dp_full, name="o_c6")
+    ctx.launch("case6_sp_sp", _case6_kernel,
+               q, k, v, o,
+               p["T_q"], p["S_kv"], p["h_q"], p["h_kv"],
+               p["d_head"], p["C"], p["P"],
+               _auto_dim_remap=False)
+
+
+def _sum_ipcq_bytes(op_log) -> int:
+    """Sum nbytes across all ipcq_copy records."""
+    return sum(
+        r.params.get("nbytes", 0)
+        for r in op_log
+        if r.op_kind == "memory" and r.op_name == "ipcq_copy"
+    )
+
+
+def _partial_score_slices(case: int) -> int:
+    """Divisor that splits S_kv into partial-score tiles, per the
+    analytical model in paper_plot_gqa_4cases_summary.py:
+
+      partial_score_per_PE = h_q * S_q * (s_kv / slices) * bytes
+
+    Case 4 (Cube-SP x PE-TP-dhead): intra-cube AR over d_head shards
+        on PE axis -> partial tile per PE has per-cube S_kv = s_kv/C.
+    Case 5 (Cube-TP-dhead x PE-SP): inter-cube AR over d_head shards
+        on cube axis -> partial tile per PE has per-PE S_kv = s_kv/P.
+    Cases 1, 2, 3, 6: no partial-score AR (only (m,l,O) merge).
+    """
+    if case == 4:
+        return _C
+    if case == 5:
+        return _P
+    return 0
+
+
+def _split_attn_layer_bytes(case: int, total_ipcq_bytes: int,
+                            s_kv: int) -> tuple[int, int]:
+    """Split per-layer attention-time IPCQ bytes into:
+      (partial_score_component, mlo_component).
+
+    The partial-score component scales with s_kv (so it must be scaled
+    when projecting from the measure-time s_kv to the headline s_kv);
+    the (m,l,O) component is constant in s_kv.
+
+    Partial-score size is analytically known per-case (formula in
+    _partial_score_slices); the remainder is treated as (m,l,O) + any
+    other S_kv-independent overhead. Per-PE = total / NUM_PES.
+    """
+    per_pe_total = total_ipcq_bytes // _NUM_PES
+    slices = _partial_score_slices(case)
+    if slices == 0:
+        # No partial-score AR for this case.
+        return 0, per_pe_total
+    partial_score_per_pe = (
+        _H_Q * _S_Q * (s_kv // slices) * _BYTES_PER_ELEM
+    )
+    partial_score_per_pe = min(partial_score_per_pe, per_pe_total)
+    mlo_per_pe = per_pe_total - partial_score_per_pe
+    return partial_score_per_pe, mlo_per_pe
+
+
+_KERNELS = (
+    (1, "Case 1 (Cube-Repl x PE-repl)",      _bench_fn_case1),
+    (2, "Case 2 (Cube-SP   x PE-repl)",      _bench_fn_case2),
+    (3, "Case 3 (Cube-Repl x PE-SP)",        _bench_fn_case3),
+    (4, "Case 4 (Cube-SP   x PE-TP d_head)", _bench_fn_case4),
+    (5, "Case 5 (Cube-TP d_head x PE-SP)",   _bench_fn_case5),
+    (6, "Case 6 (Cube-SP   x PE-SP)  [*]",   _bench_fn_case6),
+)
+
+
+def main() -> int:
+    topology = os.environ.get("GQA_1H_TOPOLOGY", "topology.yaml")
+    topo = resolve_topology(topology)
+
+    out: dict = {
+        "S_kv_measured": _S_KV_MEAS,
+        "S_kv_headline": _S_KV_HEADLINE,
+        "n_layers": _N_LAYERS,
+        "num_pes": _NUM_PES,
+        "wo_ffn_per_token_bytes": _WO_FFN_PER_TOKEN_BYTES,
+        "cases": {},
+    }
+
+    print(f"Measuring at S_kv={_S_KV_MEAS:,} ; scaling partial-score AR "
+          f"to S_kv={_S_KV_HEADLINE:,} (×{int(_PARTIAL_SCORE_SCALE)})")
+    print()
+
+    for case_id, label, bench_fn in _KERNELS:
+        try:
+            res = run_bench(
+                topology=topo, bench_fn=bench_fn,
+                device=resolve_device(None),
+                engine_factory=lambda t, d: GraphEngine(
+                    getattr(t, "topology_obj", t), enable_data=True,
+                ),
+            )
+        except Exception as e:
+            print(f"  {label:<42} FAIL: {type(e).__name__}: {e}")
+            return 1
+        if not res.completion.ok:
+            print(f"  {label:<42} ENGINE FAIL: {res.completion}")
+            return 1
+
+        total_ipcq = _sum_ipcq_bytes(res.engine.op_log)
+        partial_pe, mlo_pe = _split_attn_layer_bytes(
+            case_id, total_ipcq, _S_KV_MEAS,
+        )
+        # Per-token attention-time comm at S_kv = 1 M:
+        #   (partial_score_per_layer × scale  +  mlo_per_layer) × 80 layers
+        scaled_partial_per_token = (
+            partial_pe * int(_PARTIAL_SCORE_SCALE) * _N_LAYERS
+        )
+        mlo_per_token = mlo_pe * _N_LAYERS
+        attn_per_token = scaled_partial_per_token + mlo_per_token
+        total_per_token = attn_per_token + _WO_FFN_PER_TOKEN_BYTES
+
+        out["cases"][str(case_id)] = {
+            "label": label,
+            "total_ipcq_bytes_one_layer": total_ipcq,
+            "per_pe_partial_score_bytes_one_layer": partial_pe,
+            "per_pe_mlo_bytes_one_layer": mlo_pe,
+            "per_pe_attn_bytes_per_token_at_1M": attn_per_token,
+            "per_pe_total_bytes_per_token_at_1M": total_per_token,
+        }
+
+        print(f"  {label:<42} "
+              f"ipcq_total={total_ipcq:>10,}  "
+              f"per_pe_attn(1L)={(partial_pe + mlo_pe):>9,}  "
+              f"per_pe_total/tok@1M={total_per_token / (1<<20):>7.2f} MB")
+
+    out_path = (
+        Path(__file__).resolve().parents[2]
+        / "src" / "kernbench" / "benches"
+        / "1H_milestone_output" / "gqa" / "long_ctx"
+        / "gqa_3cases_measured_comm.json"
+    )
+    out_path.parent.mkdir(parents=True, exist_ok=True)
+    out_path.write_text(json.dumps(out, indent=2))
+    print()
+    print(f"wrote {out_path}")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

scripts/paper/paper_plot_gqa_4cases_summary.py

@@ -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:
+def _plot_comm(ax, *, mode: str = "analytical") -> None:
-    wo_mb = []
+    """Per-PE comm panel.
-    ffn_mb = []
+
-    attn_mb = []
+    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))
+        wo_mb_list.append(wo / (1 << 20))
-        ffn_mb.append(ffn / (1 << 20))
+        ffn_mb_list.append(ffn / (1 << 20))
         attn_mb.append(attn / (1 << 20))
 
-    labels = [_CASE_LABEL[c] for c in _CASES]
+    measured = _load_measured() if mode == "paired" else None
-    x = list(range(len(_CASES)))
+    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")
+    n_cases = len(_CASES)
-    ax.bar(x, ffn_mb, bottom=wo_mb, color=_FFN_COLOR, edgecolor="black")
+    x = list(range(n_cases))
-    bottoms_attn = [w + f for w, f in zip(wo_mb, ffn_mb)]
+    bar_w = 0.36 if paired else 0.65
-    ax.bar(x, attn_mb, bottom=bottoms_attn,
+    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)",
+        f"(decode S_q=1, B=1; {_N_LAYERS} layers)  —  {source_tag}"
-        fontsize=11,
     )
+    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))]
+    totals_ana = [wo_mb_list[i] + ffn_mb_list[i] + attn_mb[i]
-    ax.set_ylim(top=max(totals) * 10)
+                  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)
+        ana_bytes = totals_ana[i] * (1 << 20)
-        ax.text(x[i], totals[i] * 1.5, _fmt_bytes(total_bytes),
+        if paired:
-                ha="center", fontsize=9, weight="bold")
+            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],
+            ax.text(x_ana[i], mid, _ATTN_DESC[c],
-                    ha="center", va="center", fontsize=7,
+                    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],
+            ax.text(x_ana[i], totals_ana[i] * 0.4, _ATTN_DESC[c],
-                    ha="center", fontsize=7.5, color="grey", style="italic")
+                    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))
+
+    # (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_b = fig.add_subplot(gs[0, 0])
+    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_b)
+    _plot_budget(ax_b2)
     _plot_memory(ax_m)
     _plot_comm(ax_c)
-
+    fig.tight_layout()
-    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))
     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} "

scripts/paper/paper_plot_gqa_kv_sharding_diagram.py

@@ -0,0 +1,323 @@
+"""6-case KV-sharding tensor diagram (the slide-13 PNG export).
+
+Flat 2-D rectangles, one per sharding case, with:
+    Y axis = S_kv (vertical)        — Cube-SP / PE-SP slice it
+    X axis = d_head (horizontal)    — Cube-TP-d_head / PE-TP-d_head slice it
+
+Drops the batch axis entirely (decode: B = 1, T_q = 1). Same case set
+and visual encoding as slide 13 of GQA_full_deck.pptx; matplotlib
+renders it cleanly so the PNG sits next to the other GQA summary
+artifacts in 1H_milestone_output/gqa/long_ctx/.
+"""
+from __future__ import annotations
+
+from pathlib import Path
+
+import matplotlib.patches as mpatches
+import matplotlib.pyplot as plt
+
+_C = 8
+_P = 8
+
+_GROUP_FILLS = [
+    "#A5D8FF", "#B2F2BB", "#FFD8A8", "#FFC9C9",
+    "#D0BFFF", "#99E9F2", "#FCC2D7", "#FFEC99",
+]
+
+_ACC = {
+    "red":    "#E03131",
+    "orange": "#FD7E14",
+    "blue":   "#1C7ED6",
+    "green":  "#37B24D",
+}
+
+# (label, accent, kv, comm, overflow, encoding-flags, axis-spec)
+#   y_split   = 8 horizontal Y bands  (Cube-SP on S_kv)
+#   x_split   = 8 vertical X bands    (Cube-TP-d_head)
+#   pe_y      = 7 fine horizontal dividers within each Y band
+#   pe_x      = 7 fine vertical dividers within each X band
+#   axes      = small annotation under the chip naming the axes
+#               that the cube/PE actually shard, so the reader can
+#               parse Case 5 (where cube colour fills run X instead
+#               of Y, breaking the visual symmetry of the rest).
+_CASES = [
+    dict(label="Case 1\nCube-Repl / PE-repl", accent=_ACC["red"],
+         kv="40 GB", comm="1.2 MB", overflow=True,
+         y_split=False, x_split=False, pe_y=False, pe_x=False,
+         axes="Cube: replicated     PE: replicated"),
+    dict(label="Case 2\nCube-SP / PE-repl",   accent=_ACC["orange"],
+         kv="5 GB",  comm="3.8 MB", overflow=True,
+         y_split=True,  x_split=False, pe_y=False, pe_x=False,
+         axes="Cube → Y (S_kv)     PE: replicated"),
+    dict(label="Case 3\nCube-Repl / PE-SP",   accent=_ACC["orange"],
+         kv="5 GB",  comm="3.8 MB", overflow=True,
+         y_split=False, x_split=False, pe_y=True,  pe_x=False,
+         axes="Cube: replicated     PE → Y (S_kv)"),
+    dict(label="Case 4\nCube-SP / PE-TP-d_head", accent=_ACC["blue"],
+         kv="640 MB", comm="166 MB", overflow=False,
+         y_split=True,  x_split=False, pe_y=False, pe_x=True,
+         axes="Cube → Y (S_kv)     PE → X (d_head)"),
+    dict(label="Case 5\nCube-TP-d_head / PE-SP", accent=_ACC["blue"],
+         kv="640 MB", comm="166 MB", overflow=False,
+         y_split=False, x_split=True,  pe_y=True,  pe_x=False,
+         axes="Cube → X (d_head)     PE → Y (S_kv)"),
+    dict(label="Case 6 ★\nCube-SP / PE-SP",   accent=_ACC["green"],
+         kv="640 MB", comm="6.2 MB", overflow=False,
+         y_split=True,  x_split=False, pe_y=True,  pe_x=False,
+         axes="Cube → Y (S_kv)     PE → Y (S_kv)"),
+]
+
+_OUT_DIR = (
+    Path(__file__).resolve().parents[2]
+    / "src" / "kernbench" / "benches"
+    / "1H_milestone_output" / "gqa" / "long_ctx"
+)
+
+
+def _draw_panel(ax, cfg):
+    """Draw one case's 2-D KV-tensor rectangle into a panel ax."""
+    ax.set_xlim(0, 1)
+    ax.set_ylim(1, 0)        # Y points down (S_kv ↓)
+    ax.set_aspect("auto")
+    ax.set_xticks([])
+    ax.set_yticks([])
+
+    cube_repl = not cfg["y_split"] and not cfg["x_split"]
+    pe_repl   = not cfg["pe_y"]   and not cfg["pe_x"]
+
+    # Cube-level colour fill.
+    if cfg["y_split"] and not cfg["x_split"]:
+        # 8 horizontal Y bands.
+        for c in range(_C):
+            ax.add_patch(mpatches.Rectangle(
+                (0, c / _C), 1, 1 / _C,
+                facecolor=_GROUP_FILLS[c], edgecolor="black", linewidth=0.6))
+            ax.text(0.04, c / _C + 0.5 / _C, f"C{c}",
+                    ha="left", va="center", fontsize=8,
+                    fontweight="bold", color="#333")
+    elif cfg["x_split"] and not cfg["y_split"]:
+        # 8 vertical X bands.
+        for c in range(_C):
+            ax.add_patch(mpatches.Rectangle(
+                (c / _C, 0), 1 / _C, 1,
+                facecolor=_GROUP_FILLS[c], edgecolor="black", linewidth=0.6))
+            ax.text(c / _C + 0.5 / _C, 0.04, f"C{c}",
+                    ha="center", va="top", fontsize=8,
+                    fontweight="bold", color="#333")
+    else:
+        ax.add_patch(mpatches.Rectangle(
+            (0, 0), 1, 1,
+            facecolor="#F5F5F5", edgecolor="black", linewidth=0.8))
+        ax.text(0.5, 0.5, "× 8 cubes\nfull KV",
+                ha="center", va="center",
+                fontsize=10, fontweight="bold",
+                fontstyle="italic", color="#666")
+
+    # PE-level fine dividers — distinguished from cube boundaries by
+    # using a dashed style + slightly stronger contrast. This is what
+    # makes Case 5's PE-SP (horizontal lines across vertical cube
+    # bands) read as "different axis from the cubes" at a glance.
+    if cfg["pe_y"]:
+        outer = _C if cfg["y_split"] else 1
+        band = 1 / outer
+        for o in range(outer):
+            for p in range(1, _P):
+                y = o * band + band * p / _P
+                ax.axhline(y, color="#222", linewidth=0.8,
+                           linestyle=(0, (3, 2)), alpha=0.75)
+    if cfg["pe_x"]:
+        outer = _C if cfg["x_split"] else 1
+        band = 1 / outer
+        for o in range(outer):
+            for p in range(1, _P):
+                x = o * band + band * p / _P
+                ax.axvline(x, color="#222", linewidth=0.8,
+                           linestyle=(0, (3, 2)), alpha=0.75)
+
+    # Heavy outline on top.
+    ax.add_patch(mpatches.Rectangle(
+        (0, 0), 1, 1, facecolor="none",
+        edgecolor="black", linewidth=1.2))
+
+    # Replication badges — small text-only badges in the corners of
+    # the rectangle, no ghost-card stacking (which mis-reads as a
+    # larger enclosing tensor).
+    badges: list[str] = []
+    if cube_repl:
+        badges.append("× 8 cube copies")
+    if pe_repl and (cfg["y_split"] or cfg["x_split"]):
+        # Cube is sharded but PEs in each cube replicate that shard.
+        badges.append("× 8 PEs / cube replicate")
+    elif pe_repl and cube_repl:
+        # Both replicated — PE replication adds to the cube one.
+        badges.append("× 8 PEs / cube replicate")
+    if badges:
+        ax.text(0.98, 0.02, "\n".join(badges),
+                ha="right", va="top", fontsize=7,
+                fontweight="bold", color="#444",
+                fontstyle="italic",
+                bbox=dict(facecolor="white", edgecolor="#888",
+                          boxstyle="round,pad=0.20", linewidth=0.5))
+
+
+def _make_table_png() -> Path:
+    """Slide-14 companion table: per-PE memory + comm for all 6 cases."""
+    headers = ["Case", "Sharding", "KV / PE", "Fit",
+               "Comm/tok\n(analytical)", "Comm/tok\n(measured)", "Notes"]
+    rows = [
+        ("Case 1",   "Cube-Repl  ·  PE-repl",      "40 GB",  "✗",
+         "1.2 MB",  "1.25 MB",
+         "no sharding — full KV on every PE"),
+        ("Case 2",   "Cube-SP    ·  PE-repl",      "5 GB",   "✗",
+         "3.8 MB",  "1.27 MB",
+         "cube-axis sharded only"),
+        ("Case 3",   "Cube-Repl  ·  PE-SP",        "5 GB",   "✗",
+         "3.8 MB",  "1.39 MB",
+         "PE-axis sharded only"),
+        ("Case 4",   "Cube-SP    ·  PE-TP-d_head", "640 MB", "✓",
+         "166 MB",  "162 MB",
+         "d_head split intra-cube — partial-score AR ∝ S_kv"),
+        ("Case 5",   "Cube-TP-d_head  ·  PE-SP",   "640 MB", "✓",
+         "166 MB",  "162 MB",
+         "d_head split inter-cube — partial-score AR on UCIe"),
+        ("Case 6 ★", "Cube-SP    ·  PE-SP",        "640 MB", "✓",
+         "6.2 MB",  "1.41 MB",
+         "S_kv split on both axes — (m,ℓ,O) AR only, S_kv-indep."),
+    ]
+    accents = [_ACC["red"], _ACC["orange"], _ACC["orange"],
+               _ACC["blue"], _ACC["blue"], _ACC["green"]]
+
+    fig, ax = plt.subplots(figsize=(20.0, 4.6))
+    ax.set_axis_off()
+
+    cell_data = [headers] + [list(r) for r in rows]
+    tbl = ax.table(cellText=cell_data,
+                   colWidths=[0.07, 0.18, 0.08, 0.05, 0.12, 0.12, 0.38],
+                   cellLoc="center", loc="center")
+    tbl.auto_set_font_size(False)
+    tbl.set_fontsize(11)
+    tbl.scale(1.0, 2.2)
+
+    n_cols = len(headers)
+    n_rows = len(rows) + 1   # +1 header
+    # Header styling.
+    for ci in range(n_cols):
+        cell = tbl[(0, ci)]
+        cell.set_facecolor("#1F4E79")
+        cell.set_text_props(color="white", weight="bold")
+        cell.set_edgecolor("#1F4E79")
+    # Body styling.
+    for ri, row in enumerate(rows, start=1):
+        is_pareto = row[0].endswith("★")
+        row_fill = "#E8F5E9" if is_pareto else (
+            "white" if ri % 2 == 1 else "#F5F5F7")
+        # Case-name cell uses accent.
+        case_cell = tbl[(ri, 0)]
+        case_cell.set_facecolor(accents[ri - 1])
+        case_cell.set_text_props(color="white", weight="bold")
+        # Remaining cells.
+        for ci in range(1, n_cols):
+            cell = tbl[(ri, ci)]
+            cell.set_facecolor(row_fill)
+            txt_kwargs = {"weight": "bold" if is_pareto else "normal",
+                          "color": "#333"}
+            if ci == 2:                              # KV / PE
+                txt_kwargs["color"] = (
+                    "#C62828" if row[3] == "✗" else "#2E7D32")
+                txt_kwargs["weight"] = "bold"
+            if ci == 3:                              # Fit
+                txt_kwargs["color"] = (
+                    "#C62828" if row[3] == "✗" else "#2E7D32")
+                txt_kwargs["weight"] = "bold"
+            cell.set_text_props(**txt_kwargs)
+        # Last-column (Notes) cells left-aligned for readability.
+        tbl[(ri, n_cols - 1)].get_text().set_ha("left")
+
+    # Force left-align on the Notes header too.
+    tbl[(0, n_cols - 1)].get_text().set_ha("left")
+
+    fig.suptitle(
+        "GQA decode KV-sharding — per-PE memory & communication  "
+        "(LLaMA 70B GQA single KV-head group · S_kv = 1 M, FP16, 80 layers)",
+        fontsize=12, y=0.97,
+    )
+    out = _OUT_DIR / "gqa_kv_sharding_6cases_table.png"
+    fig.savefig(out, dpi=150, bbox_inches="tight")
+    plt.close(fig)
+    print(f"wrote {out}")
+    return out
+
+
+def main() -> Path:
+    _OUT_DIR.mkdir(parents=True, exist_ok=True)
+    n = len(_CASES)
+    fig = plt.figure(figsize=(20.0, 7.0))
+    # Three rows per column: case chip · axis-spec annotation · rectangle.
+    gs = fig.add_gridspec(3, n,
+                          height_ratios=[0.55, 0.32, 8.5],
+                          hspace=0.05, wspace=0.20,
+                          left=0.04, right=0.99,
+                          top=0.93, bottom=0.06)
+
+    for i, cfg in enumerate(_CASES):
+        # Top: case chip header.
+        ax_chip = fig.add_subplot(gs[0, i])
+        ax_chip.set_xticks([])
+        ax_chip.set_yticks([])
+        for spine in ax_chip.spines.values():
+            spine.set_visible(False)
+        ax_chip.add_patch(mpatches.Rectangle(
+            (0, 0), 1, 1, transform=ax_chip.transAxes,
+            facecolor=cfg["accent"], edgecolor=cfg["accent"]))
+        ax_chip.text(0.5, 0.5, cfg["label"],
+                     ha="center", va="center",
+                     fontsize=10, fontweight="bold",
+                     color="white")
+
+        # Middle: axis-spec annotation — names which axis the cube
+        # shards on and which axis the PE shards on (essential for
+        # parsing Case 5 where the cube colour fills run X instead
+        # of Y, breaking the visual symmetry of the rest).
+        ax_axes = fig.add_subplot(gs[1, i])
+        ax_axes.set_xticks([])
+        ax_axes.set_yticks([])
+        for spine in ax_axes.spines.values():
+            spine.set_visible(False)
+        ax_axes.add_patch(mpatches.Rectangle(
+            (0, 0), 1, 1, transform=ax_axes.transAxes,
+            facecolor="#F5F5F7", edgecolor="#CCCCCC",
+            linewidth=0.6))
+        ax_axes.text(0.5, 0.5, cfg["axes"],
+                     ha="center", va="center",
+                     fontsize=8.5, fontweight="bold",
+                     color="#1F4E79")
+
+        # Bottom: the tensor rectangle.
+        ax = fig.add_subplot(gs[2, i])
+        _draw_panel(ax, cfg)
+        ax.set_xlabel("X : d_head = 128  →",
+                      fontsize=9, fontweight="bold",
+                      fontstyle="italic", color="#1F4E79")
+        ax.set_ylabel("Y : S_kv = 1 M  ↓",
+                      fontsize=9, fontweight="bold",
+                      fontstyle="italic", color="#1F4E79")
+
+    fig.suptitle(
+        "GQA decode KV-tensor sharding — 6 cases  ·  "
+        "LLaMA 70B GQA single KV-head group  ·  "
+        "C = 8 cubes × P = 8 PEs  ·  S_kv = 1 M, FP16, 80 layers",
+        fontsize=12, y=0.99,
+    )
+
+    out = _OUT_DIR / "gqa_kv_sharding_6cases_diagram.png"
+    fig.savefig(out, dpi=150, bbox_inches="tight")
+    plt.close(fig)
+    print(f"wrote {out}")
+
+    # Companion table PNG (slide-14 export).
+    _make_table_png()
+    return out
+
+
+if __name__ == "__main__":
+    main()

scripts/verify_case4_dhead_tp.py

@@ -0,0 +1,148 @@
+"""Phase 1a smoke test for the Case 4 d_head-TP decode kernel.
+
+Runs `gqa_attention_decode_long_ctx_cube_sp_pe_tp_dhead_kernel` at small
+S_kv (2K) to verify:
+  1. it imports cleanly,
+  2. it runs without error under the same harness as Case 6,
+  3. captures op_log_summary (gemm/ipcq/dma counts),
+  4. compares vs analytical predictions and vs Case 6 (same memory tier).
+
+Usage:
+  python scripts/verify_case4_dhead_tp.py
+"""
+from __future__ import annotations
+
+import os
+import sys
+from pathlib import Path
+
+from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_sp_pe_sp import (
+    gqa_attention_decode_long_ctx_cube_sp_pe_sp_kernel as _case6_kernel,
+)
+from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_sp_pe_tp_dhead import (
+    gqa_attention_decode_long_ctx_cube_sp_pe_tp_dhead_kernel as _case4_kernel,
+)
+from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_tp_dhead_pe_sp import (
+    gqa_attention_decode_long_ctx_cube_tp_dhead_pe_sp_kernel as _case5_kernel,
+)
+from kernbench.benches.gqa_helpers.shared._gqa_panel_helpers import (
+    _ccl_cfg, _summarize_op_log,
+)
+from kernbench.ccl.sfr_config import configure_sfr_intercube_multisip
+from kernbench.policy.placement.dp import DPPolicy
+from kernbench.runtime_api.bench_runner import run_bench
+from kernbench.runtime_api.types import resolve_device
+from kernbench.sim_engine.engine import GraphEngine
+from kernbench.topology.builder import resolve_topology
+
+# Small smoke-test params (S_kv=2K is enough to exercise tile-loop + AR).
+_PARAMS = dict(C=8, P=8, T_q=1, S_kv=2_048,
+               d_head=128, h_q=8, h_kv=1)
+
+
+def _bench_fn_case4(ctx):
+    p = _PARAMS
+    configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
+    dp_full = DPPolicy(cube="replicate", pe="column_wise",
+                       num_cubes=p["C"], num_pes=p["P"])
+    dp_kv = DPPolicy(cube="row_wise", pe="column_wise",
+                     num_cubes=p["C"], num_pes=p["P"])
+    q = ctx.zeros((p["T_q"], p["h_q"] * p["d_head"]),
+                  dtype="f16", dp=dp_full, name="q_c4")
+    k = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]),
+                  dtype="f16", dp=dp_kv, name="k_c4")
+    v = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]),
+                  dtype="f16", dp=dp_kv, name="v_c4")
+    o = ctx.empty((p["T_q"], p["h_q"] * p["d_head"]),
+                  dtype="f16", dp=dp_full, name="o_c4")
+    ctx.launch("case4_dhead_tp", _case4_kernel,
+               q, k, v, o,
+               p["T_q"], p["S_kv"], p["h_q"], p["h_kv"],
+               p["d_head"], p["C"], p["P"],
+               _auto_dim_remap=False)
+
+
+def _bench_fn_case5(ctx):
+    p = _PARAMS
+    configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
+    dp_q = DPPolicy(cube="column_wise", pe="replicate",
+                    num_cubes=p["C"], num_pes=p["P"])
+    dp_kv = DPPolicy(cube="column_wise", pe="row_wise",
+                     num_cubes=p["C"], num_pes=p["P"])
+    q = ctx.zeros((p["T_q"], p["h_q"] * p["d_head"]),
+                  dtype="f16", dp=dp_q, name="q_c5")
+    k = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]),
+                  dtype="f16", dp=dp_kv, name="k_c5")
+    v = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]),
+                  dtype="f16", dp=dp_kv, name="v_c5")
+    o = ctx.empty((p["T_q"], p["h_q"] * p["d_head"]),
+                  dtype="f16", dp=dp_q, name="o_c5")
+    ctx.launch("case5_dhead_tp_inter", _case5_kernel,
+               q, k, v, o,
+               p["T_q"], p["S_kv"], p["h_q"], p["h_kv"],
+               p["d_head"], p["C"], p["P"],
+               _auto_dim_remap=False)
+
+
+def _bench_fn_case6(ctx):
+    p = _PARAMS
+    configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
+    dp_full = DPPolicy(cube="replicate", pe="replicate",
+                       num_cubes=p["C"], num_pes=p["P"])
+    dp_kv = DPPolicy(cube="row_wise", pe="row_wise",
+                     num_cubes=p["C"], num_pes=p["P"])
+    q = ctx.zeros((p["T_q"], p["h_q"] * p["d_head"]),
+                  dtype="f16", dp=dp_full, name="q_c6")
+    k = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]),
+                  dtype="f16", dp=dp_kv, name="k_c6")
+    v = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]),
+                  dtype="f16", dp=dp_kv, name="v_c6")
+    o = ctx.empty((p["T_q"], p["h_q"] * p["d_head"]),
+                  dtype="f16", dp=dp_full, name="o_c6")
+    ctx.launch("case6_sp_sp", _case6_kernel,
+               q, k, v, o,
+               p["T_q"], p["S_kv"], p["h_q"], p["h_kv"],
+               p["d_head"], p["C"], p["P"],
+               _auto_dim_remap=False)
+
+
+def main() -> int:
+    topology = os.environ.get("GQA_1H_TOPOLOGY", "topology.yaml")
+    topo = resolve_topology(topology)
+
+    print(f"Smoke params: {_PARAMS}")
+    print()
+
+    for label, bench_fn in (
+        ("Case 4 (Cube-SP × PE-TP d_head)",   _bench_fn_case4),
+        ("Case 5 (Cube-TP d_head × PE-SP)",   _bench_fn_case5),
+        ("Case 6 (Cube-SP × PE-SP S_kv)",     _bench_fn_case6),
+    ):
+        try:
+            res = run_bench(
+                topology=topo, bench_fn=bench_fn,
+                device=resolve_device(None),
+                engine_factory=lambda t, d: GraphEngine(
+                    getattr(t, "topology_obj", t), enable_data=True,
+                ),
+            )
+        except Exception as e:
+            print(f"  {label:<42} FAIL: {type(e).__name__}: {e}")
+            continue
+        if not res.completion.ok:
+            print(f"  {label:<42} ENGINE FAIL: {res.completion}")
+            continue
+        s = _summarize_op_log(res.engine.op_log)
+        lat = (res.engine.op_log[-1].t_end if res.engine.op_log else 0.0)
+        print(f"  {label:<42} "
+              f"gemm={s['gemm_count']:>4}  "
+              f"ipcq={s['ipcq_copy_count']:>4}  "
+              f"dma_r={s['dma_read_count']:>4}  "
+              f"dma_w={s['dma_write_count']:>3}  "
+              f"latency={lat:.1f} ns  "
+              f"(n_ops={len(res.engine.op_log)})")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

src/kernbench/benches/1H_milestone_output/gqa/long_ctx/gqa_3cases_measured_comm.json

@@ -0,0 +1,57 @@
+{
+  "S_kv_measured": 8192,
+  "S_kv_headline": 1048576,
+  "n_layers": 80,
+  "num_pes": 64,
+  "wo_ffn_per_token_bytes": 1310720,
+  "cases": {
+    "1": {
+      "label": "Case 1 (Cube-Repl x PE-repl)",
+      "total_ipcq_bytes_one_layer": 0,
+      "per_pe_partial_score_bytes_one_layer": 0,
+      "per_pe_mlo_bytes_one_layer": 0,
+      "per_pe_attn_bytes_per_token_at_1M": 0,
+      "per_pe_total_bytes_per_token_at_1M": 1310720
+    },
+    "2": {
+      "label": "Case 2 (Cube-SP   x PE-repl)",
+      "total_ipcq_bytes_one_layer": 14560,
+      "per_pe_partial_score_bytes_one_layer": 0,
+      "per_pe_mlo_bytes_one_layer": 227,
+      "per_pe_attn_bytes_per_token_at_1M": 18160,
+      "per_pe_total_bytes_per_token_at_1M": 1328880
+    },
+    "3": {
+      "label": "Case 3 (Cube-Repl x PE-SP)",
+      "total_ipcq_bytes_one_layer": 116480,
+      "per_pe_partial_score_bytes_one_layer": 0,
+      "per_pe_mlo_bytes_one_layer": 1820,
+      "per_pe_attn_bytes_per_token_at_1M": 145600,
+      "per_pe_total_bytes_per_token_at_1M": 1456320
+    },
+    "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
+    },
+    "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
+    },
+    "6": {
+      "label": "Case 6 (Cube-SP   x PE-SP)  [*]",
+      "total_ipcq_bytes_one_layer": 131040,
+      "per_pe_partial_score_bytes_one_layer": 0,
+      "per_pe_mlo_bytes_one_layer": 2047,
+      "per_pe_attn_bytes_per_token_at_1M": 163760,
+      "per_pe_total_bytes_per_token_at_1M": 1474480
+    }
+  }
+}

src/kernbench/benches/gqa_helpers/long_ctx/_gqa_attention_decode_long_ctx_cube_sp_pe_tp_dhead.py

@@ -0,0 +1,363 @@
+"""GQA decode kernel — Case 4 (Cube-SP × PE-TP on d_head).
+
+Per the per-PE memory/comm summary chart  (paper_plot_gqa_4cases_summary.py):
+this is the d_head-TP variant — KV is sharded 64-way (cube=row_wise on
+S_kv, pe=column_wise on d_head), so per-PE KV memory matches Case 6 ★
+(640 MB / PE at 1M context across 80 layers), but the attention
+algorithm needs an extra mid-layer partial-score AllReduce because
+each PE only owns d_head/P dimensions of K/V.
+
+Algorithm (single layer, decode T_q=1, B=1):
+  1. Each PE loads its own d_head/P slice of Q, K, V (Q replicated
+     in HBM; PE-load via column-wise DPPolicy carves out the slice).
+  2. Per S_kv tile, each PE computes a PARTIAL score tile:
+        partial_score = Q_shard · K_shard^T  shape: (G·T_q, S_local)
+  3. AR-SUM these partial scores INTRA-CUBE (across the P PEs of a
+     cube) so every PE has the full Q · K^T for the cube's S_local
+     tokens. Online-softmax merge folds tiles into a running state.
+  4. After all tiles, each PE has (m_i, ℓ_i, O_i_shard) for cube i,
+     where O_i_shard is the cube's partial attention output on this
+     PE's d_head/P slice.
+  5. INTER-CUBE online-softmax merge of (m, ℓ, O_shard) across C
+     cubes (lrab center-root reduce). After this, every PE has the
+     full-sequence attention output for its d_head/P slice.
+  6. Final normalise + store: the P PEs of the root cube each write
+     their (G·T_q, d_head/P) slice of the output.
+
+Tensor layout:
+  Q : (T_q, h_q · d_head)  cube=replicate, pe=column_wise on d_head
+      → per-PE shape (T_q, h_q · d_head/P), byte-reshaped to
+        (G·T_q, d_head/P) for the per-PE GEMM input.
+  K : (S_kv, h_kv · d_head) cube=row_wise on S_kv, pe=column_wise
+      on d_head → per-PE shape (S_local, d_head/P), byte-reshaped to
+      (d_head/P, S_local) for the K^T-form GEMM.
+  V : same layout as K.
+  O : (T_q, h_q · d_head) cube=replicate, pe=column_wise on d_head
+      → each writer PE owns its d_head/P slice.
+
+Topology / SFR:
+  - Requires ``configure_sfr_intercube_multisip`` (provides disjoint
+    ``intra_*`` and ``E/W/N/S`` namespaces).
+  - Intra-CUBE PEs in a 2×4 grid; AR-SUM of partial scores uses a
+    reduce-to-PE0-then-broadcast pattern over intra_E/W + intra_N/S.
+  - Inter-CUBE lrab center-root reduce on (m, ℓ, O_shard) over the
+    4×2 cube sub-mesh (root cube = 6).
+
+Naming note: this kernel exemplifies the chart's "PE-TP" label
+meaning **d_head TP** (NOT the batch-axis TP of the original slide-17
+Case 1). The chart-side Case label is "Case 4  Cube-SP × PE-TP".
+
+ADR-0065 P5 scope note: this kernel is meant for op_log / latency
+measurement (matches the milestone-1h-gqa sweep mode). Numeric
+data-mode parity is a separate follow-up.
+"""
+from __future__ import annotations
+
+
+TILE_S_KV = 1024  # ADR-0063 §A.2 S_kv-axis tile sweep (per-tile width).
+
+# lrab geometry for the C=8 single-KV-head group (4×2 cube sub-mesh).
+_SUB_W = 4
+_SUB_H = 2
+_ROOT_COL = _SUB_W // 2   # 2
+_ROOT_ROW = _SUB_H // 2   # 1
+_ROOT_CUBE = _ROOT_ROW * _SUB_W + _ROOT_COL  # 6
+
+
+def _merge_running(m_local, l_local, O_local, m_other, l_other, O_other, *, tl):
+    """Online-softmax merge of two partial ``(m, ℓ, O)`` triples.
+
+    O is the per-PE d_head-slice partial; merge is element-wise so it
+    composes naturally even with O being a d_head shard.
+    """
+    m_new = tl.maximum(m_local, m_other)
+    scale_old = tl.exp(m_local - m_new)
+    scale_new = tl.exp(m_other - m_new)
+    l_new = l_local * scale_old + l_other * scale_new
+    O_new = O_local * scale_old + O_other * scale_new
+    return m_new, l_new, O_new
+
+
+def _intra_cube_allreduce_sum(tensor, *, tl,
+                              pe_col: int, pe_row: int,
+                              pe_cols_used: int, pe_rows_used: int):
+    """Sum-AllReduce a tensor across the 8 PEs of one cube (2×4 grid).
+
+    Reduce-to-PE0 (row chain west + col bridge north) followed by
+    broadcast-from-PE0 (col bridge south + row chain east). Each phase
+    operates in place on ``tensor`` via ``tl.copy_to``.
+
+    Used here to combine PARTIAL Q·K^T score tiles across the P=8 PEs
+    of a cube — each PE has only d_head/P dims, so each contributes a
+    partial dot product; the full score = elementwise sum across PEs.
+    """
+    # ── Phase 1a: row reduce (west) — accumulate at pe_col=0 ────────
+    if pe_cols_used > 1:
+        if pe_col == pe_cols_used - 1:
+            tl.send(dir="intra_W", src=tensor)
+        elif 0 < pe_col < pe_cols_used - 1:
+            with tl.scratch_scope():
+                other = tl.recv(dir="intra_E", shape=tensor.shape, dtype="f16")
+                tl.copy_to(tensor, tensor + other)
+            tl.send(dir="intra_W", src=tensor)
+        elif pe_col == 0:
+            with tl.scratch_scope():
+                other = tl.recv(dir="intra_E", shape=tensor.shape, dtype="f16")
+                tl.copy_to(tensor, tensor + other)
+
+    # ── Phase 1b: col reduce on pe_col=0 (north) — accumulate at row=0 ─
+    if pe_col == 0 and pe_rows_used > 1:
+        if pe_row == pe_rows_used - 1:
+            tl.send(dir="intra_N", src=tensor)
+        elif pe_row == 0:
+            with tl.scratch_scope():
+                other = tl.recv(dir="intra_S", shape=tensor.shape, dtype="f16")
+                tl.copy_to(tensor, tensor + other)
+
+    # ── Phase 2a: col broadcast (south) — row=0 → row=1 on pe_col=0 ─
+    if pe_col == 0 and pe_rows_used > 1:
+        if pe_row == 0:
+            tl.send(dir="intra_S", src=tensor)
+        elif pe_row > 0:
+            with tl.scratch_scope():
+                received = tl.recv(dir="intra_N", shape=tensor.shape, dtype="f16")
+                tl.copy_to(tensor, received)
+
+    # ── Phase 2b: row broadcast (east) — pe_col=0 → all on each row ─
+    if pe_cols_used > 1:
+        if pe_col == 0:
+            tl.send(dir="intra_E", src=tensor)
+        elif 0 < pe_col < pe_cols_used - 1:
+            with tl.scratch_scope():
+                received = tl.recv(dir="intra_W", shape=tensor.shape, dtype="f16")
+                tl.copy_to(tensor, received)
+            tl.send(dir="intra_E", src=tensor)
+        elif pe_col == pe_cols_used - 1:
+            with tl.scratch_scope():
+                received = tl.recv(dir="intra_W", shape=tensor.shape, dtype="f16")
+                tl.copy_to(tensor, received)
+
+
+def gqa_attention_decode_long_ctx_cube_sp_pe_tp_dhead_kernel(
+    q_ptr: int,
+    k_ptr: int,
+    v_ptr: int,
+    o_ptr: int,
+    T_q: int,
+    S_kv: int,
+    h_q: int,
+    h_kv: int,
+    d_head: int,
+    C: int,
+    P: int,
+    *,
+    tl,
+) -> None:
+    """Case-4 decode: Cube-SP × PE-TP-on-d_head, lrab center-root at cube 6."""
+    G = h_q // h_kv
+    if d_head % P != 0:
+        raise ValueError(
+            f"Case 4 requires d_head={d_head} divisible by P={P}"
+        )
+    if S_kv % C != 0:
+        raise ValueError(
+            f"Case 4 requires S_kv={S_kv} divisible by C={C}"
+        )
+
+    d_head_local = d_head // P   # 128 / 8 = 16 in the headline setup
+    S_local = S_kv // C
+    pe_id = tl.program_id(axis=0)
+    cube_id = tl.program_id(axis=1)
+
+    PE_GRID_COLS = 4
+    pe_col = pe_id % PE_GRID_COLS
+    pe_row = pe_id // PE_GRID_COLS
+    pe_cols_used = min(PE_GRID_COLS, P)
+    pe_rows_used = (P + PE_GRID_COLS - 1) // PE_GRID_COLS
+
+    KV_ROW_BYTES = d_head_local * 2  # bytes per K/V row at this PE's d_head slice
+
+    # ── Load this PE's d_head slice of Q (byte-conserving reshape) ──
+    # Q in HBM is (T_q, h_q · d_head); the column_wise pe shard gives
+    # each PE (T_q, h_q · d_head/P). Reshape to (G·T_q, d_head/P) for the
+    # GEMM: byte-conserving since G = h_q (h_kv = 1).
+    Q = tl.load(q_ptr, shape=(G * T_q, d_head_local), dtype="f16")
+
+    # ── Local attention (S_kv-axis tile sweep) with intra-cube AR sum ─
+    n_tiles = (S_local + TILE_S_KV - 1) // TILE_S_KV
+
+    # Tile 0: establish persistent (m, ℓ, O_shard). O_shard is the
+    # per-PE d_head/P slice of the partial attention output.
+    tile_s0 = min(TILE_S_KV, S_local)
+    K_T = tl.load(k_ptr, shape=(d_head_local, tile_s0), dtype="f16")
+    V = tl.load(v_ptr, shape=(tile_s0, d_head_local), dtype="f16")
+    partial_scores0 = tl.dot(Q, K_T)
+    # AR-SUM partial scores intra-cube → every PE in the cube has the
+    # full Q · K^T_tile0 scores.
+    _intra_cube_allreduce_sum(
+        partial_scores0, tl=tl,
+        pe_col=pe_col, pe_row=pe_row,
+        pe_cols_used=pe_cols_used, pe_rows_used=pe_rows_used,
+    )
+    m_local = tl.max(partial_scores0, axis=-1)
+    exp_scores = tl.exp(partial_scores0 - m_local)
+    l_local = tl.sum(exp_scores, axis=-1)
+    # Partial output sharded on d_head: (G·T_q, d_head/P) per PE.
+    O_local = tl.dot(exp_scores, V)
+
+    for tile_idx in range(1, n_tiles):
+        tile_start = tile_idx * TILE_S_KV
+        tile_s = min(TILE_S_KV, S_local - tile_start)
+        with tl.scratch_scope():
+            K_T_t = tl.load(k_ptr + tile_start * KV_ROW_BYTES,
+                            shape=(d_head_local, tile_s), dtype="f16")
+            V_t = tl.load(v_ptr + tile_start * KV_ROW_BYTES,
+                          shape=(tile_s, d_head_local), dtype="f16")
+            partial_scores_t = tl.dot(Q, K_T_t)
+            _intra_cube_allreduce_sum(
+                partial_scores_t, tl=tl,
+                pe_col=pe_col, pe_row=pe_row,
+                pe_cols_used=pe_cols_used, pe_rows_used=pe_rows_used,
+            )
+            m_tile = tl.max(partial_scores_t, axis=-1)
+            exp_scores_t = tl.exp(partial_scores_t - m_tile)
+            l_tile = tl.sum(exp_scores_t, axis=-1)
+            O_tile = tl.dot(exp_scores_t, V_t)
+            m_new, l_new, O_new = _merge_running(
+                m_local, l_local, O_local, m_tile, l_tile, O_tile, tl=tl,
+            )
+            tl.copy_to(m_local, m_new)
+            tl.copy_to(l_local, l_new)
+            tl.copy_to(O_local, O_new)
+
+    # ── Inter-CUBE lrab-adapted center-root reduce on (m, ℓ, O_shard) ─
+    # Every PE participates (each contributes its own d_head/P slice
+    # of O_shard); the merge is element-wise so d_head-sharding is
+    # preserved through the AR.
+    row = cube_id // _SUB_W
+    col = cube_id % _SUB_W
+
+    # Phase 1: row reduce — converge at col == _ROOT_COL.
+    if col == 0:
+        tl.send(dir="E", src=m_local)
+        tl.send(dir="E", src=l_local)
+        tl.send(dir="E", src=O_local)
+    elif 0 < col < _ROOT_COL:
+        with tl.scratch_scope():
+            m_other = tl.recv(dir="W", shape=m_local.shape, dtype="f16")
+            l_other = tl.recv(dir="W", shape=l_local.shape, dtype="f16")
+            O_other = tl.recv(dir="W", shape=O_local.shape, dtype="f16")
+            m_new, l_new, O_new = _merge_running(
+                m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
+            )
+            tl.copy_to(m_local, m_new)
+            tl.copy_to(l_local, l_new)
+            tl.copy_to(O_local, O_new)
+        tl.send(dir="E", src=m_local)
+        tl.send(dir="E", src=l_local)
+        tl.send(dir="E", src=O_local)
+    elif col == _ROOT_COL:
+        with tl.scratch_scope():
+            m_other = tl.recv(dir="W", shape=m_local.shape, dtype="f16")
+            l_other = tl.recv(dir="W", shape=l_local.shape, dtype="f16")
+            O_other = tl.recv(dir="W", shape=O_local.shape, dtype="f16")
+            m_new, l_new, O_new = _merge_running(
+                m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
+            )
+            tl.copy_to(m_local, m_new)
+            tl.copy_to(l_local, l_new)
+            tl.copy_to(O_local, O_new)
+        with tl.scratch_scope():
+            m_other = tl.recv(dir="E", shape=m_local.shape, dtype="f16")
+            l_other = tl.recv(dir="E", shape=l_local.shape, dtype="f16")
+            O_other = tl.recv(dir="E", shape=O_local.shape, dtype="f16")
+            m_new, l_new, O_new = _merge_running(
+                m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
+            )
+            tl.copy_to(m_local, m_new)
+            tl.copy_to(l_local, l_new)
+            tl.copy_to(O_local, O_new)
+    elif _ROOT_COL < col < _SUB_W - 1:
+        with tl.scratch_scope():
+            m_other = tl.recv(dir="E", shape=m_local.shape, dtype="f16")
+            l_other = tl.recv(dir="E", shape=l_local.shape, dtype="f16")
+            O_other = tl.recv(dir="E", shape=O_local.shape, dtype="f16")
+            m_new, l_new, O_new = _merge_running(
+                m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
+            )
+            tl.copy_to(m_local, m_new)
+            tl.copy_to(l_local, l_new)
+            tl.copy_to(O_local, O_new)
+        tl.send(dir="W", src=m_local)
+        tl.send(dir="W", src=l_local)
+        tl.send(dir="W", src=O_local)
+    elif col == _SUB_W - 1:
+        tl.send(dir="W", src=m_local)
+        tl.send(dir="W", src=l_local)
+        tl.send(dir="W", src=O_local)
+
+    # Phase 2: col reduce on col == _ROOT_COL — converge at row == _ROOT_ROW.
+    if col == _ROOT_COL:
+        if row == 0:
+            tl.send(dir="S", src=m_local)
+            tl.send(dir="S", src=l_local)
+            tl.send(dir="S", src=O_local)
+        elif 0 < row < _ROOT_ROW:
+            with tl.scratch_scope():
+                m_other = tl.recv(dir="N", shape=m_local.shape, dtype="f16")
+                l_other = tl.recv(dir="N", shape=l_local.shape, dtype="f16")
+                O_other = tl.recv(dir="N", shape=O_local.shape, dtype="f16")
+                m_new, l_new, O_new = _merge_running(
+                    m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
+                )
+                tl.copy_to(m_local, m_new)
+                tl.copy_to(l_local, l_new)
+                tl.copy_to(O_local, O_new)
+            tl.send(dir="S", src=m_local)
+            tl.send(dir="S", src=l_local)
+            tl.send(dir="S", src=O_local)
+        elif row == _ROOT_ROW:
+            with tl.scratch_scope():
+                m_other = tl.recv(dir="N", shape=m_local.shape, dtype="f16")
+                l_other = tl.recv(dir="N", shape=l_local.shape, dtype="f16")
+                O_other = tl.recv(dir="N", shape=O_local.shape, dtype="f16")
+                m_new, l_new, O_new = _merge_running(
+                    m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
+                )
+                tl.copy_to(m_local, m_new)
+                tl.copy_to(l_local, l_new)
+                tl.copy_to(O_local, O_new)
+            if _SUB_H - 1 > _ROOT_ROW:
+                with tl.scratch_scope():
+                    m_other = tl.recv(dir="S", shape=m_local.shape, dtype="f16")
+                    l_other = tl.recv(dir="S", shape=l_local.shape, dtype="f16")
+                    O_other = tl.recv(dir="S", shape=O_local.shape, dtype="f16")
+                    m_new, l_new, O_new = _merge_running(
+                        m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
+                    )
+                    tl.copy_to(m_local, m_new)
+                    tl.copy_to(l_local, l_new)
+                    tl.copy_to(O_local, O_new)
+        elif _ROOT_ROW < row < _SUB_H - 1:
+            with tl.scratch_scope():
+                m_other = tl.recv(dir="S", shape=m_local.shape, dtype="f16")
+                l_other = tl.recv(dir="S", shape=l_local.shape, dtype="f16")
+                O_other = tl.recv(dir="S", shape=O_local.shape, dtype="f16")
+                m_new, l_new, O_new = _merge_running(
+                    m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
+                )
+                tl.copy_to(m_local, m_new)
+                tl.copy_to(l_local, l_new)
+                tl.copy_to(O_local, O_new)
+            tl.send(dir="N", src=m_local)
+            tl.send(dir="N", src=l_local)
+            tl.send(dir="N", src=O_local)
+        elif row == _SUB_H - 1 and _SUB_H - 1 > _ROOT_ROW:
+            tl.send(dir="N", src=m_local)
+            tl.send(dir="N", src=l_local)
+            tl.send(dir="N", src=O_local)
+
+    # ── Final normalise + store (root cube: each P-PE writes its d_head slice) ─
+    if cube_id == _ROOT_CUBE:
+        O_final = O_local / l_local
+        tl.store(o_ptr, O_final)

src/kernbench/benches/gqa_helpers/long_ctx/_gqa_attention_decode_long_ctx_cube_tp_dhead_pe_sp.py

@@ -0,0 +1,265 @@
+"""GQA decode kernel — Case 5 (Cube-TP on d_head × PE-SP on S_kv).
+
+Per the per-PE memory/comm summary chart  (paper_plot_gqa_4cases_summary.py):
+this is the second d_head-TP variant — KV is sharded 64-way (cube=
+column_wise on d_head, pe=row_wise on S_kv). Per-PE KV memory matches
+Cases 4 and 6 ★ (640 MB / PE at 1M context across 80 layers), but the
+attention algorithm pays the partial-score AR over INTER-cube fabric
+(UCIe, slower) instead of Case 4's intra-cube fabric (NoC).
+
+Symmetric mirror of Case 4 — intra and inter-cube AR roles SWAPPED:
+
+  Step             Case 4 (Cube-SP × PE-TP/d_head)   Case 5 (THIS — Cube-TP/d_head × PE-SP)
+  ───────────────  ────────────────────────────────  ──────────────────────────────────────
+  partial Q·K^T    AR-SUM intra-cube (NoC)           AR-SUM inter-cube (UCIe)
+  online-softmax   inter-cube lrab on (m, ℓ, O)      intra-cube merge on (m, ℓ, O)
+                   (every PE participates; each       (online-softmax reduce-to-PE0)
+                    PE owns its d_head/P slice)
+  writer           root cube's P PEs (each writes    each cube's PE 0 writes its
+                   its d_head/P slice)                d_head/C slice (C writers total)
+
+Algorithm:
+  1. Each PE loads its (d_head/C × S_kv/P) slice of K and V; full-d_head
+     Q from HBM but only its d_head/C slice via column_wise DPPolicy.
+  2. Per S_kv tile (per PE — already a strict subset of the cube's
+     S_kv), compute partial_score = Q_shard · K_shard^T.  Shape:
+     (G·T_q, S_kv/P).  AR-SUM these PARTIAL scores INTER-CUBE across
+     the C cubes (same PE rank, different d_head slices) so every PE
+     in cube c, rank p has the full Q · K^T_p scores for THIS PE's
+     S_kv/P token range.
+  3. Local softmax + P · V_shard → partial output (G·T_q, d_head/C)
+     for this PE's S_kv/P slice and this cube's d_head/C slice.
+  4. INTRA-CUBE online-softmax merge of (m, ℓ, O) reduce-to-PE0
+     across the cube's P PEs (different S_kv slices, same d_head slice).
+  5. PE 0 of each cube normalises and writes the cube's d_head/C slice.
+
+Tensor layout:
+  Q : (T_q, h_q · d_head)  cube=column_wise on d_head, pe=replicate.
+      Per PE shape (T_q, h_q · d_head/C); reshape to (G·T_q, d_head/C).
+  K : (S_kv, h_kv · d_head) cube=column_wise on d_head, pe=row_wise on S_kv.
+      Per PE shape (S_kv/P, d_head/C); reshape to (d_head/C, S_kv/P) for K^T.
+  V : same layout as K.
+  O : (T_q, h_q · d_head) cube=column_wise on d_head, pe=replicate.
+      Per-cube writer (PE 0) owns its d_head/C slice.
+
+Topology / SFR:
+  - Requires ``configure_sfr_intercube_multisip`` (disjoint intra_*
+    and E/W/N/S namespaces).
+  - INTRA-CUBE PEs arranged 2×4 (used for online-softmax merge).
+  - INTER-CUBE 4×2 cube sub-mesh (used for partial-score AR sum).
+
+ADR-0065 P5 scope note: op_log / latency measurement only; numeric
+data-mode parity is a separate follow-up (matches the milestone-1h-gqa
+sweep mode).
+"""
+from __future__ import annotations
+
+
+TILE_S_KV = 1024  # per-tile S_kv width within a PE's own S_kv/P shard.
+
+# Cube sub-mesh geometry (4×2). For Case 5 we pick cube 0 as the AR-sum
+# root (simpler than lrab; partial-score AR doesn't need geometric centering
+# because the reduction is element-wise sum, not online-softmax merge).
+_SUB_W = 4
+_SUB_H = 2
+_AR_ROOT_CUBE = 0
+
+
+def _merge_running(m_local, l_local, O_local, m_other, l_other, O_other, *, tl):
+    """Online-softmax merge of two partial ``(m, ℓ, O)`` triples."""
+    m_new = tl.maximum(m_local, m_other)
+    scale_old = tl.exp(m_local - m_new)
+    scale_new = tl.exp(m_other - m_new)
+    l_new = l_local * scale_old + l_other * scale_new
+    O_new = O_local * scale_old + O_other * scale_new
+    return m_new, l_new, O_new
+
+
+def _inter_cube_allreduce_sum(tensor, *, tl, cube_id: int,
+                              sub_w: int = _SUB_W, sub_h: int = _SUB_H):
+    """Sum-AllReduce a tensor INTER-CUBE across the 4×2 cube sub-mesh.
+
+    Reduce-to-cube-0 (row chain west + col bridge north) followed by
+    broadcast-from-cube-0 (col bridge south + row chain east). The
+    payload is the partial-score tile — each cube has a different
+    d_head slice's contribution; element-wise sum recovers the full
+    Q·K^T scores for this PE rank's S_kv/P tokens.
+
+    All P PEs of every cube participate independently (per-PE-rank
+    inter-cube fabric lanes).
+    """
+    row = cube_id // sub_w
+    col = cube_id % sub_w
+
+    # ── Phase 1a: row reduce (west) — accumulate at col=0 ────────────
+    if sub_w > 1:
+        if col == sub_w - 1:
+            tl.send(dir="W", src=tensor)
+        elif 0 < col < sub_w - 1:
+            with tl.scratch_scope():
+                other = tl.recv(dir="E", shape=tensor.shape, dtype="f16")
+                tl.copy_to(tensor, tensor + other)
+            tl.send(dir="W", src=tensor)
+        elif col == 0:
+            with tl.scratch_scope():
+                other = tl.recv(dir="E", shape=tensor.shape, dtype="f16")
+                tl.copy_to(tensor, tensor + other)
+
+    # ── Phase 1b: col reduce on col=0 (north) — accumulate at row=0 ──
+    if col == 0 and sub_h > 1:
+        if row == sub_h - 1:
+            tl.send(dir="N", src=tensor)
+        elif row == 0:
+            with tl.scratch_scope():
+                other = tl.recv(dir="S", shape=tensor.shape, dtype="f16")
+                tl.copy_to(tensor, tensor + other)
+
+    # ── Phase 2a: col broadcast (south) — row=0 → row=1 on col=0 ─────
+    if col == 0 and sub_h > 1:
+        if row == 0:
+            tl.send(dir="S", src=tensor)
+        elif row > 0:
+            with tl.scratch_scope():
+                received = tl.recv(dir="N", shape=tensor.shape, dtype="f16")
+                tl.copy_to(tensor, received)
+
+    # ── Phase 2b: row broadcast (east) — col=0 → all on each row ─────
+    if sub_w > 1:
+        if col == 0:
+            tl.send(dir="E", src=tensor)
+        elif 0 < col < sub_w - 1:
+            with tl.scratch_scope():
+                received = tl.recv(dir="W", shape=tensor.shape, dtype="f16")
+                tl.copy_to(tensor, received)
+            tl.send(dir="E", src=tensor)
+        elif col == sub_w - 1:
+            with tl.scratch_scope():
+                received = tl.recv(dir="W", shape=tensor.shape, dtype="f16")
+                tl.copy_to(tensor, received)
+
+
+def gqa_attention_decode_long_ctx_cube_tp_dhead_pe_sp_kernel(
+    q_ptr: int,
+    k_ptr: int,
+    v_ptr: int,
+    o_ptr: int,
+    T_q: int,
+    S_kv: int,
+    h_q: int,
+    h_kv: int,
+    d_head: int,
+    C: int,
+    P: int,
+    *,
+    tl,
+) -> None:
+    """Case-5 decode: Cube-TP-on-d_head × PE-SP-on-S_kv."""
+    G = h_q // h_kv
+    if d_head % C != 0:
+        raise ValueError(
+            f"Case 5 requires d_head={d_head} divisible by C={C}"
+        )
+    if S_kv % P != 0:
+        raise ValueError(
+            f"Case 5 requires S_kv={S_kv} divisible by P={P}"
+        )
+
+    d_head_local = d_head // C       # 128 / 8 = 16 in the headline setup
+    S_local_pe = S_kv // P           # each PE owns S_kv/P tokens
+    pe_id = tl.program_id(axis=0)
+    cube_id = tl.program_id(axis=1)
+
+    PE_GRID_COLS = 4
+    pe_col = pe_id % PE_GRID_COLS
+    pe_row = pe_id // PE_GRID_COLS
+    pe_cols_used = min(PE_GRID_COLS, P)
+    pe_rows_used = (P + PE_GRID_COLS - 1) // PE_GRID_COLS
+
+    KV_ROW_BYTES = d_head_local * 2  # bytes per K/V row at this PE's d_head slice
+
+    # ── Load this PE's d_head slice of Q ────────────────────────────
+    # Q in HBM is (T_q, h_q · d_head); column_wise cube shard gives each
+    # cube (T_q, h_q · d_head/C); pe=replicate within cube → every PE
+    # in the cube has the same data. Reshape to (G·T_q, d_head/C).
+    Q = tl.load(q_ptr, shape=(G * T_q, d_head_local), dtype="f16")
+
+    # ── Local attention (tile sweep on this PE's S_kv/P shard) ──────
+    n_tiles = (S_local_pe + TILE_S_KV - 1) // TILE_S_KV
+
+    # Tile 0: bootstrap (m, ℓ, O_shard). The partial Q·K^T is across
+    # d_head — must AR inter-cube before softmax.
+    tile_s0 = min(TILE_S_KV, S_local_pe)
+    K_T = tl.load(k_ptr, shape=(d_head_local, tile_s0), dtype="f16")
+    V = tl.load(v_ptr, shape=(tile_s0, d_head_local), dtype="f16")
+    partial_scores0 = tl.dot(Q, K_T)
+    _inter_cube_allreduce_sum(partial_scores0, tl=tl, cube_id=cube_id)
+    m_local = tl.max(partial_scores0, axis=-1)
+    exp_scores = tl.exp(partial_scores0 - m_local)
+    l_local = tl.sum(exp_scores, axis=-1)
+    # Partial output (G·T_q, d_head/C) — sharded on d_head only; PE
+    # still holds the full d_head/C slice for its S_kv/P tokens.
+    O_local = tl.dot(exp_scores, V)
+
+    for tile_idx in range(1, n_tiles):
+        tile_start = tile_idx * TILE_S_KV
+        tile_s = min(TILE_S_KV, S_local_pe - tile_start)
+        with tl.scratch_scope():
+            K_T_t = tl.load(k_ptr + tile_start * KV_ROW_BYTES,
+                            shape=(d_head_local, tile_s), dtype="f16")
+            V_t = tl.load(v_ptr + tile_start * KV_ROW_BYTES,
+                          shape=(tile_s, d_head_local), dtype="f16")
+            partial_scores_t = tl.dot(Q, K_T_t)
+            _inter_cube_allreduce_sum(partial_scores_t, tl=tl,
+                                      cube_id=cube_id)
+            m_tile = tl.max(partial_scores_t, axis=-1)
+            exp_scores_t = tl.exp(partial_scores_t - m_tile)
+            l_tile = tl.sum(exp_scores_t, axis=-1)
+            O_tile = tl.dot(exp_scores_t, V_t)
+            m_new, l_new, O_new = _merge_running(
+                m_local, l_local, O_local, m_tile, l_tile, O_tile, tl=tl,
+            )
+            tl.copy_to(m_local, m_new)
+            tl.copy_to(l_local, l_new)
+            tl.copy_to(O_local, O_new)
+
+    # ── INTRA-CUBE online-softmax reduce-to-PE0 (different S_kv slices) ─
+    # Each cube ends with (m, ℓ, O) for its d_head/C slice at PE 0.
+    if pe_cols_used > 1:
+        if pe_col < pe_cols_used - 1:
+            with tl.scratch_scope():
+                m_other = tl.recv(dir="intra_E", shape=m_local.shape, dtype="f16")
+                l_other = tl.recv(dir="intra_E", shape=l_local.shape, dtype="f16")
+                O_other = tl.recv(dir="intra_E", shape=O_local.shape, dtype="f16")
+                m_new, l_new, O_new = _merge_running(
+                    m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
+                )
+                tl.copy_to(m_local, m_new)
+                tl.copy_to(l_local, l_new)
+                tl.copy_to(O_local, O_new)
+        if pe_col > 0:
+            tl.send(dir="intra_W", src=m_local)
+            tl.send(dir="intra_W", src=l_local)
+            tl.send(dir="intra_W", src=O_local)
+
+    if pe_col == 0 and pe_rows_used > 1:
+        if pe_row < pe_rows_used - 1:
+            with tl.scratch_scope():
+                m_other = tl.recv(dir="intra_S", shape=m_local.shape, dtype="f16")
+                l_other = tl.recv(dir="intra_S", shape=l_local.shape, dtype="f16")
+                O_other = tl.recv(dir="intra_S", shape=O_local.shape, dtype="f16")
+                m_new, l_new, O_new = _merge_running(
+                    m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
+                )
+                tl.copy_to(m_local, m_new)
+                tl.copy_to(l_local, l_new)
+                tl.copy_to(O_local, O_new)
+        if pe_row > 0:
+            tl.send(dir="intra_N", src=m_local)
+            tl.send(dir="intra_N", src=l_local)
+            tl.send(dir="intra_N", src=O_local)
+
+    # ── Final normalise + store (one writer per cube: PE 0; each
+    #     cube owns a disjoint d_head/C slice of the output) ─────────
+    if pe_id == 0:
+        O_final = O_local / l_local
+        tl.store(o_ptr, O_final)

src/kernbench/benches/gqa_helpers/long_ctx/gqa_decode_long_ctx_4cases.py

@@ -38,6 +38,12 @@ from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_
 from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_sp_pe_tp import (
     gqa_attention_decode_long_ctx_cube_sp_pe_tp_kernel,
 )
+from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_sp_pe_tp_dhead import (
+    gqa_attention_decode_long_ctx_cube_sp_pe_tp_dhead_kernel,
+)
+from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_tp_dhead_pe_sp import (
+    gqa_attention_decode_long_ctx_cube_tp_dhead_pe_sp_kernel,
+)
 from kernbench.benches.gqa_helpers.shared._gqa_panel_helpers import (
     _ccl_cfg,
     _summarize_op_log,
@@ -58,10 +64,12 @@ _SWEEP_JSON = _OUTPUT_DIR / "sweep_decode.json"
 
 
 _PANELS = (
-    "single_kv_group_decode_long_ctx_gqa_cube_sp_pe_sp",      # Case 4 ★ optimal
+    "single_kv_group_decode_long_ctx_gqa_cube_sp_pe_sp",         # Case 6 ★ optimal (S_kv-SP both axes; was "Case 4" in slide-17 numbering)
-    "single_kv_group_decode_long_ctx_gqa_cube_repl_pe_tp",    # Case 2 (no comm; 8× memory)
+    "single_kv_group_decode_long_ctx_gqa_cube_repl_pe_tp",       # Case 1 (no sharding; 40 GB / PE @ 1M)
-    "single_kv_group_decode_long_ctx_gqa_cube_repl_pe_sp",    # Case 3 (intra-CUBE AR only; 8× memory)
+    "single_kv_group_decode_long_ctx_gqa_cube_repl_pe_sp",       # Case 3 (intra-cube AR on (m,ℓ,O); 5 GB / PE)
-    "single_kv_group_decode_long_ctx_gqa_cube_sp_pe_tp",      # Case 1 (inter-CUBE lrab only; PE-TP B=1 waste)
+    "single_kv_group_decode_long_ctx_gqa_cube_sp_pe_tp",         # Case 2 (inter-cube AR on (m,ℓ,O); 5 GB / PE)
+    "single_kv_group_decode_long_ctx_gqa_cube_sp_pe_tp_dhead",   # Case 4 (Cube-SP × PE-TP-on-d_head — d_head-TP partial-score AR intra-cube)
+    "single_kv_group_decode_long_ctx_gqa_cube_tp_dhead_pe_sp",   # Case 5 (Cube-TP-on-d_head × PE-SP — d_head-TP partial-score AR inter-cube)
 )
 
 # Each entry: (kind, panel-specific params).
@@ -102,6 +110,23 @@ _PANEL_DISPATCH: dict[str, tuple[str, dict]] = {
         "T_q": 1, "S_kv": 8_192,
         "d_head": 128, "h_q": 8, "h_kv": 1,
     }),
+    "single_kv_group_decode_long_ctx_gqa_cube_sp_pe_tp_dhead": ("decode_long_ctx_cube_sp_pe_tp_dhead", {
+        # Case 4 (new numbering): Cube-SP on S_kv × PE-TP on d_head.
+        # Partial-score AR happens intra-cube (NoC); (m,ℓ,O) merge happens
+        # inter-cube (UCIe). Memory-equal to Case 6; comm dominated by
+        # partial-score AR which scales with S_kv.
+        "C": 8, "P": 8,
+        "T_q": 1, "S_kv": 8_192,
+        "d_head": 128, "h_q": 8, "h_kv": 1,
+    }),
+    "single_kv_group_decode_long_ctx_gqa_cube_tp_dhead_pe_sp": ("decode_long_ctx_cube_tp_dhead_pe_sp", {
+        # Case 5 (new numbering): Cube-TP on d_head × PE-SP on S_kv.
+        # Mirror of Case 4: partial-score AR happens INTER-cube (UCIe);
+        # (m,ℓ,O) merge happens INTRA-cube (NoC). Memory-equal to Case 6.
+        "C": 8, "P": 8,
+        "T_q": 1, "S_kv": 8_192,
+        "d_head": 128, "h_q": 8, "h_kv": 1,
+    }),
 }
 
 
@@ -238,6 +263,74 @@ def _run_decode_panel_long_ctx_cube_sp_pe_sp(
     )
 
 
+def _run_decode_panel_long_ctx_cube_sp_pe_tp_dhead(
+    ctx, *, panel: str, C: int, P: int,
+    T_q: int, S_kv: int,
+    d_head: int, h_q: int, h_kv: int,
+) -> None:
+    """Case 4 (new) runner: Cube-SP × PE-TP-on-d_head.
+
+    Q is cube-replicated, PE-sharded column-wise (d_head/P per PE).
+    K, V are cube=row_wise on S_kv, pe=column_wise on d_head — each PE
+    owns (S_kv/C, d_head/P). The kernel does an intra-cube AR-SUM of
+    partial scores then an inter-cube online-softmax merge of
+    (m, ℓ, O_shard).
+    """
+    configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
+    dp_q = DPPolicy(cube="replicate", pe="column_wise",
+                    num_cubes=C, num_pes=P)
+    dp_kv = DPPolicy(cube="row_wise", pe="column_wise",
+                     num_cubes=C, num_pes=P)
+    q = ctx.zeros((T_q, h_q * d_head),
+                  dtype="f16", dp=dp_q, name=f"{panel}_q")
+    k = ctx.zeros((S_kv, h_kv * d_head),
+                  dtype="f16", dp=dp_kv, name=f"{panel}_k")
+    v = ctx.zeros((S_kv, h_kv * d_head),
+                  dtype="f16", dp=dp_kv, name=f"{panel}_v")
+    o = ctx.empty((T_q, h_q * d_head),
+                  dtype="f16", dp=dp_q, name=f"{panel}_o")
+    ctx.launch(
+        panel, gqa_attention_decode_long_ctx_cube_sp_pe_tp_dhead_kernel,
+        q, k, v, o,
+        T_q, S_kv, h_q, h_kv, d_head, C, P,
+        _auto_dim_remap=False,
+    )
+
+
+def _run_decode_panel_long_ctx_cube_tp_dhead_pe_sp(
+    ctx, *, panel: str, C: int, P: int,
+    T_q: int, S_kv: int,
+    d_head: int, h_q: int, h_kv: int,
+) -> None:
+    """Case 5 (new) runner: Cube-TP-on-d_head × PE-SP-on-S_kv.
+
+    Q is cube-sharded column_wise (d_head/C per cube), pe-replicated
+    within cube. K, V are cube=column_wise on d_head, pe=row_wise on
+    S_kv — each PE owns (S_kv/P, d_head/C). The kernel does an
+    inter-cube AR-SUM of partial scores then an intra-cube online-
+    softmax merge of (m, ℓ, O) reduce-to-PE0.
+    """
+    configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
+    dp_q = DPPolicy(cube="column_wise", pe="replicate",
+                    num_cubes=C, num_pes=P)
+    dp_kv = DPPolicy(cube="column_wise", pe="row_wise",
+                     num_cubes=C, num_pes=P)
+    q = ctx.zeros((T_q, h_q * d_head),
+                  dtype="f16", dp=dp_q, name=f"{panel}_q")
+    k = ctx.zeros((S_kv, h_kv * d_head),
+                  dtype="f16", dp=dp_kv, name=f"{panel}_k")
+    v = ctx.zeros((S_kv, h_kv * d_head),
+                  dtype="f16", dp=dp_kv, name=f"{panel}_v")
+    o = ctx.empty((T_q, h_q * d_head),
+                  dtype="f16", dp=dp_q, name=f"{panel}_o")
+    ctx.launch(
+        panel, gqa_attention_decode_long_ctx_cube_tp_dhead_pe_sp_kernel,
+        q, k, v, o,
+        T_q, S_kv, h_q, h_kv, d_head, C, P,
+        _auto_dim_remap=False,
+    )
+
+
 def _make_bench_fn(panel: str):
     kind, params = _PANEL_DISPATCH[panel]
 
@@ -250,6 +343,10 @@ def _make_bench_fn(panel: str):
             _run_decode_panel_long_ctx_cube_repl_pe_sp(ctx, panel=panel, **params)
         elif kind == "decode_long_ctx_cube_sp_pe_tp":
             _run_decode_panel_long_ctx_cube_sp_pe_tp(ctx, panel=panel, **params)
+        elif kind == "decode_long_ctx_cube_sp_pe_tp_dhead":
+            _run_decode_panel_long_ctx_cube_sp_pe_tp_dhead(ctx, panel=panel, **params)
+        elif kind == "decode_long_ctx_cube_tp_dhead_pe_sp":
+            _run_decode_panel_long_ctx_cube_tp_dhead_pe_sp(ctx, panel=panel, **params)
         else:
             raise RuntimeError(
                 f"milestone-gqa-decode-long-ctx-4cases panel {panel!r} has "