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/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())