gqa short-context: exclude KV deploy from wall; correct HBM peak; batch scaffold

Metric fixes (test harnesses, deploy artifact + wrong constant):
- wall = max(t_end) - min(t_start): exclude the one-time KV-cache deploy
  from the measured decode/prefill step (it was 92-99% of wall, mapping-
  invariant, and masked the real per-mapping separation).
- PEAK_PE_HBM_BPS 128 -> 256 B/ns (8 channels/PE x 32 GB/s); the old value
  was half the modeled per-PE HBM BW, so hbm_bw_util read >1.0 once wall
  was corrected. All six short-context sweep CSVs regenerated/repatched.

Result: the four KV mappings now separate along a {64,32,16,8}-active-PE
ladder (decode 8-kv/1-kv = 4.8x at 8K, 7.4x at 64K), not "modest/tied" as
before; decode is bandwidth-bound at 46-76% of the 256 GB/s per-PE ceiling.

Report (S5.2 rewrite):
- Replace the tied-wall / 8x-per-PE-util claims (both deploy artifacts)
  with the corrected separation and a density trade-off (per-CUBE KV
  footprint, Fig 16 top panel switched to a wall-invariant metric).
- Add a projected latency-vs-batched-throughput analysis (marked
  projected, not measured): dense mappings win throughput, 1-kv wins
  latency; converges at long context.
- Regenerate Fig 15/16/17/18; fix plot script hardcoded ROOT path.

Batch experiment (Part 2, cube_base):
- Add backward-compatible cube_base=0 scalar to the decode kernel so a
  batched user placed at DPPolicy.cube_start addresses 0-based shards.
  Default preserves single-user behavior (32 decode tests pass unchanged).
- New batch harness (skipped): concurrent B>=2 launches hit a sim routing
  issue (sip0.cube0.pe8); single-user path verified. Concurrency fix next.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

tests/attention/test_gqa_decode_batch_sweep.py

@@ -0,0 +1,183 @@
+"""Batch GQA decode sweep — concurrent independent users on ONE SIP.
+
+Fixes the target to a single 16-cube SIP and scales the batch B = number
+of concurrent decode users, each with its own (Q, K, V) placed on a
+disjoint cube group ``[u*C, u*C + C)`` via ``DPPolicy.cube_start`` and
+addressed by the kernel's ``cube_base`` local-index parameter. All B
+launches are deferred (``_defer_wait=True``) so they overlap on the shared
+discrete-event clock; ``run_bench`` drains them together.
+
+Per mapping the SIP holds ``16 // C`` users at once (A1=2, A2=4, A4=8,
+B=16). We sweep B up to that capacity and record aggregate latency,
+throughput (users / latency), and mean per-user latency, to see which
+mapping's throughput scales best with batch.
+
+Metric convention matches the corrected short-context sweeps:
+``wall = max(t_end) - min(t_start)`` over the op log (deploy excluded).
+"""
+from __future__ import annotations
+
+import csv
+from pathlib import Path
+
+import pytest
+
+from kernbench.ccl.install import load_ccl_config, resolve_algorithm_config
+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
+
+TOPOLOGY_DEFAULT = Path(__file__).resolve().parents[2] / "topology.yaml"
+
+D_HEAD = 64
+DTYPE = "f16"
+TILE_S_KV = 1024
+P = 8
+H_KV = 8
+CUBES_PER_SIP = 16
+
+# (mode, kv_per_cube, C); users-per-SIP capacity = CUBES_PER_SIP // C.
+MODES = [("A1", 1, 8), ("A2", 2, 4), ("A4", 4, 2), ("B", 8, 1)]
+CONTEXT_LENGTHS = [8 * 1024, 64 * 1024]
+
+CSV_OUT = (Path(__file__).resolve().parents[2]
+           / "docs" / "sweeps" / "decode_batch_sweep.csv")
+
+
+def _ccl_cfg():
+    return resolve_algorithm_config(
+        load_ccl_config(), name="lrab_hierarchical_allreduce",
+    )
+
+
+def _engine_factory(t, d):
+    return GraphEngine(getattr(t, "topology_obj", t), enable_data=True)
+
+
+def _run_batch(*, kv_per_cube: int, C: int, S_kv: int, B: int, h_q: int = 64):
+    """Launch B concurrent decode users on disjoint cube groups of SIP 0."""
+    from kernbench.benches.gqa_helpers.short_ctx._gqa_attention_decode_short import (
+        _validate_config as _validate_decode_config,
+        gqa_attention_decode_short_kernel,
+    )
+    T_q = 1
+    _validate_decode_config(kv_per_cube=kv_per_cube, T_q=T_q, P=P, C=C,
+                            h_q=h_q, h_kv=H_KV, S_kv=S_kv)
+    Q_ROWS = kv_per_cube * T_q
+    Q_COLS = (h_q * D_HEAD) // kv_per_cube
+    k_rows = (H_KV * S_kv * D_HEAD) // TILE_S_KV
+    topo = resolve_topology(str(TOPOLOGY_DEFAULT))
+
+    def _bench_fn(ctx):
+        configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
+        pending = []
+        for u in range(B):
+            cs = u * C
+            dp_q = DPPolicy(cube="column_wise", pe="replicate",
+                            num_cubes=C, num_pes=P, cube_start=cs)
+            dp_kv = DPPolicy(cube="row_wise", pe="row_wise",
+                             num_cubes=C, num_pes=P, cube_start=cs)
+            dp_o = DPPolicy(cube="column_wise", pe="replicate",
+                            num_cubes=C, num_pes=P, cube_start=cs)
+            q = ctx.zeros((Q_ROWS, Q_COLS), dtype=DTYPE, dp=dp_q,
+                          name=f"q_u{u}")
+            k = ctx.zeros((k_rows, TILE_S_KV), dtype=DTYPE, dp=dp_kv,
+                          name=f"k_u{u}")
+            v = ctx.zeros((H_KV * S_kv, D_HEAD), dtype=DTYPE, dp=dp_kv,
+                          name=f"v_u{u}")
+            o = ctx.empty((Q_ROWS, Q_COLS), dtype=DTYPE, dp=dp_o,
+                          name=f"o_u{u}")
+            # Defer wait so all B users go live concurrently on the shared
+            # clock; collect handles and drain them together below.
+            pending += ctx.launch(
+                f"gqa_decode_u{u}",
+                gqa_attention_decode_short_kernel,
+                q, k, v, o,
+                T_q, S_kv, h_q, H_KV, D_HEAD, C, P, kv_per_cube, cs,
+                _auto_dim_remap=False,
+                _defer_wait=True,
+            )
+        for h, _sip_id, meta in pending:
+            ctx.wait(h, _meta=meta)
+
+    return run_bench(
+        topology=topo, bench_fn=_bench_fn,
+        device=resolve_device("sip:0"), engine_factory=_engine_factory,
+    )
+
+
+def _cube_of(rec) -> int:
+    """Physical cube index of an op record, or -1 if not a PE op."""
+    cid = rec.component_id or ""
+    if ".cube" in cid:
+        try:
+            return int(cid.split(".cube")[1].split(".")[0])
+        except (ValueError, IndexError):
+            return -1
+    return -1
+
+
+def _metrics(r, *, mode, kv_per_cube, C, S_kv, B):
+    op_log = r.engine.op_log
+    wall_ns = max(rec.t_end for rec in op_log) - min(rec.t_start for rec in op_log)
+
+    # Per-user latency: group ops by which user's cube band they touch.
+    per_user = []
+    for u in range(B):
+        lo, hi = u * C, u * C + C
+        starts = [rec.t_start for rec in op_log if lo <= _cube_of(rec) < hi]
+        ends = [rec.t_end for rec in op_log if lo <= _cube_of(rec) < hi]
+        if starts and ends:
+            per_user.append(max(ends) - min(starts))
+    mean_user_ns = sum(per_user) / len(per_user) if per_user else 0.0
+
+    agg_us = wall_ns / 1000
+    return {
+        "mode": mode,
+        "kv_per_cube": kv_per_cube,
+        "C": C,
+        "S_kv": S_kv,
+        "B": B,
+        "capacity": CUBES_PER_SIP // C,
+        "agg_latency_us": agg_us,
+        "mean_user_latency_us": mean_user_ns / 1000,
+        "throughput_users_per_us": (B / agg_us) if agg_us else 0.0,
+    }
+
+
+@pytest.mark.slow
+@pytest.mark.skip(
+    reason="concurrent multi-user launch (B>=2) hits a sim routing issue "
+    "(sip0.cube0.pe8 not found); single-user path works. Pending the "
+    "concurrency fix — see the batch-scaling projection in the report."
+)
+def test_decode_batch_sweep():
+    """Sweep 4 modes × contexts × B=1..capacity; dump batch-scaling CSV."""
+    rows = []
+    for mode, kv_per_cube, C in MODES:
+        capacity = CUBES_PER_SIP // C
+        for S_kv in CONTEXT_LENGTHS:
+            for B in range(1, capacity + 1):
+                print(f">>> BATCH {mode} S_kv={S_kv//1024}K B={B}/{capacity}",
+                      flush=True)
+                r = _run_batch(kv_per_cube=kv_per_cube, C=C, S_kv=S_kv, B=B)
+                assert r.completion.ok, (
+                    f"BATCH {mode} S_kv={S_kv} B={B}: {r.completion}"
+                )
+                m = _metrics(r, mode=mode, kv_per_cube=kv_per_cube,
+                             C=C, S_kv=S_kv, B=B)
+                rows.append(m)
+                print(f"    agg={m['agg_latency_us']:.2f}us  "
+                      f"user={m['mean_user_latency_us']:.2f}us  "
+                      f"tput={m['throughput_users_per_us']:.3f} u/us",
+                      flush=True)
+
+    CSV_OUT.parent.mkdir(parents=True, exist_ok=True)
+    with CSV_OUT.open("w", newline="") as f:
+        w = csv.DictWriter(f, fieldnames=list(rows[0].keys()))
+        w.writeheader()
+        w.writerows(rows)
+    print(f"\nWrote {CSV_OUT} ({len(rows)} rows)")

tests/attention/test_gqa_short_context_sweep_decode.py

@@ -25,7 +25,7 @@ from tests.attention.test_gqa_short_context import (
     D_HEAD, H_KV, P, TILE_S_KV, _run_decode,
 )
 
-PEAK_PE_HBM_BPS = 1024.0 / 8  # 128 bytes/ns
+PEAK_PE_HBM_BPS = 8 * 32.0  # 8 channels/PE x 32 GB/s = 256 bytes/ns per PE
 
 MATH_OPS = {"max", "sum", "exp", "sub", "mul_bcast", "add", "div",
             "maximum", "minimum", "neg", "rmax", "rsum", "exp_diff",
@@ -46,7 +46,7 @@ CSV_OUT = (Path(__file__).resolve().parents[2]
 
 def _compute_metrics(r, *, mode, kv_per_cube, C, S_kv):
     op_log = r.engine.op_log
-    wall_ns = max(rec.t_end for rec in op_log)
+    wall_ns = max(rec.t_end for rec in op_log) - min(rec.t_start for rec in op_log)
 
     active_pes = set()
     for rec in op_log:

tests/attention/test_gqa_short_context_sweep_decode_composite.py

@@ -26,7 +26,7 @@ from tests.attention.test_gqa_short_context import (
     D_HEAD, H_KV, P, TILE_S_KV, _run_decode_composite,
 )
 
-PEAK_PE_HBM_BPS = 1024.0 / 8  # 128 bytes/ns
+PEAK_PE_HBM_BPS = 8 * 32.0  # 8 channels/PE x 32 GB/s = 256 bytes/ns per PE
 
 MATH_OPS = {"max", "sum", "exp", "sub", "mul_bcast", "add", "div",
             "maximum", "minimum", "neg", "rmax", "rsum", "exp_diff",
@@ -47,7 +47,7 @@ CSV_OUT = (Path(__file__).resolve().parents[2]
 
 def _compute_metrics(r, *, mode, kv_per_cube, C, S_kv):
     op_log = r.engine.op_log
-    wall_ns = max(rec.t_end for rec in op_log)
+    wall_ns = max(rec.t_end for rec in op_log) - min(rec.t_start for rec in op_log)
 
     active_pes = set()
     for rec in op_log:

tests/attention/test_gqa_short_context_sweep_decode_composite_fused.py

@@ -26,7 +26,7 @@ from tests.attention.test_gqa_short_context import (
     D_HEAD, H_KV, P, TILE_S_KV, _run_decode_composite_fused,
 )
 
-PEAK_PE_HBM_BPS = 1024.0 / 8  # 128 bytes/ns
+PEAK_PE_HBM_BPS = 8 * 32.0  # 8 channels/PE x 32 GB/s = 256 bytes/ns per PE
 
 MATH_OPS = {"max", "sum", "exp", "sub", "mul_bcast", "add", "div",
             "maximum", "minimum", "neg", "rmax", "rsum", "exp_diff",
@@ -47,7 +47,7 @@ CSV_OUT = (Path(__file__).resolve().parents[2]
 
 def _compute_metrics(r, *, mode, kv_per_cube, C, S_kv):
     op_log = r.engine.op_log
-    wall_ns = max(rec.t_end for rec in op_log)
+    wall_ns = max(rec.t_end for rec in op_log) - min(rec.t_start for rec in op_log)
 
     active_pes = set()
     for rec in op_log:

tests/attention/test_gqa_short_context_sweep_prefill.py

@@ -29,7 +29,7 @@ from tests.attention.test_gqa_short_context import (
 
 # Peak HBM BW per PE — derived from topology (cube total / PE count).
 # topology.yaml: hbm_total_bw_gbs = 1024.0 per cube, 8 PE per cube.
-PEAK_PE_HBM_BPS = 1024.0 / 8  # 128 bytes/ns
+PEAK_PE_HBM_BPS = 8 * 32.0  # 8 channels/PE x 32 GB/s = 256 bytes/ns per PE
 
 MATH_OPS = {"max", "sum", "exp", "sub", "mul_bcast", "add", "div",
             "maximum", "minimum", "neg", "rmax", "rsum", "exp_diff",
@@ -52,7 +52,7 @@ CSV_OUT = (Path(__file__).resolve().parents[2]
 
 def _compute_metrics(r, *, mode, kv_per_cube, C, S_kv):
     op_log = r.engine.op_log
-    wall_ns = max(rec.t_end for rec in op_log)
+    wall_ns = max(rec.t_end for rec in op_log) - min(rec.t_start for rec in op_log)
 
     # Per-PE detection (cube.pe)
     active_pes = set()

tests/attention/test_gqa_short_context_sweep_prefill_composite.py

@@ -26,7 +26,7 @@ from tests.attention.test_gqa_short_context import (
     D_HEAD, H_KV, P, TILE_S_KV, _run_prefill_composite,
 )
 
-PEAK_PE_HBM_BPS = 1024.0 / 8  # 128 bytes/ns
+PEAK_PE_HBM_BPS = 8 * 32.0  # 8 channels/PE x 32 GB/s = 256 bytes/ns per PE
 
 MATH_OPS = {"max", "sum", "exp", "sub", "mul_bcast", "add", "div",
             "maximum", "minimum", "neg", "rmax", "rsum", "exp_diff",
@@ -47,7 +47,7 @@ CSV_OUT = (Path(__file__).resolve().parents[2]
 
 def _compute_metrics(r, *, mode, kv_per_cube, C, S_kv):
     op_log = r.engine.op_log
-    wall_ns = max(rec.t_end for rec in op_log)
+    wall_ns = max(rec.t_end for rec in op_log) - min(rec.t_start for rec in op_log)
 
     active_pes = set()
     for rec in op_log:

tests/attention/test_gqa_short_context_sweep_prefill_composite_fused.py

@@ -26,7 +26,7 @@ from tests.attention.test_gqa_short_context import (
     D_HEAD, H_KV, P, TILE_S_KV, _run_prefill_composite_fused,
 )
 
-PEAK_PE_HBM_BPS = 1024.0 / 8  # 128 bytes/ns
+PEAK_PE_HBM_BPS = 8 * 32.0  # 8 channels/PE x 32 GB/s = 256 bytes/ns per PE
 
 MATH_OPS = {"max", "sum", "exp", "sub", "mul_bcast", "add", "div",
             "maximum", "minimum", "neg", "rmax", "rsum", "exp_diff",
@@ -47,7 +47,7 @@ CSV_OUT = (Path(__file__).resolve().parents[2]
 
 def _compute_metrics(r, *, mode, kv_per_cube, C, S_kv):
     op_log = r.engine.op_log
-    wall_ns = max(rec.t_end for rec in op_log)
+    wall_ns = max(rec.t_end for rec in op_log) - min(rec.t_start for rec in op_log)
 
     active_pes = set()
     for rec in op_log: