gqa: ADR-0060/0062/0063/0064 unified GQA kernels + CPU cost model

Land the new GQA fused-attention kernels (ADR-0060) for prefill/decode
across long and short context, the TL discipline primitives they depend
on (ADR-0062 lazy load, ADR-0063 scratch_scope + copy_to), and the
per-op-type CPU issue cost model (ADR-0064). Remove the pre-ADR-0060
mesh-attention baseline now that the unified kernels supersede it.

ADR-0060 (long context)
- _gqa_decode.py: M-fold + 2-level chain reduce-to-root (Level-2
  intra-CUBE row-then-col + Level-1 inter-CUBE) — root-only output.
- _gqa_prefill.py: head-parallel + Ring KV rotation around C CUBEs,
  online-softmax merge per ring step, per-CUBE distributed output.
- Each merge stage wraps in scratch_scope() and persists running
  (m, l, O) via copy_to() to lift the 1 MiB scratch ceiling.

ADR-0060 §B.split.2 (short context, kv_per_cube in {1,2,4,8})
- _gqa_decode_short.py / _gqa_prefill_short.py: no cube-SP; each CUBE
  owns whole KV heads; PE-parallel heads with intra-group chain
  reduce. Prefill has no Ring KV (each head fully resident).

ADR-0062 (lazy tl.load): future-bearing TensorHandle, auto-wait at
first consuming op (dot/MATH/store/send/copy_to/composite).

ADR-0063 (tl.scratch_scope + tl.copy_to): scoped per-tile arena with
copy_to writeback primitive for persistent running state.

ADR-0064 (CPU issue cost model)
- common/cpu_issue_cost.py: per-op-type table (composite=40 ns,
  primitives=5 ns); ratios are load-bearing per D1.
- TLContext: issue_cost_table param; _emit_dispatch_overhead(kind)
  consults table with dispatch_cycles fallback (ADR-0046 §D6
  back-compat).
- Live PE_CPU paths (greenlet + legacy) construct TLContext with
  DEFAULT_CPU_ISSUE_COST so saturation lever (ADR-0060 §1) is
  measurable end-to-end.

P7 headline bench: milestone-gqa-headline writes per-panel
op_log_summary to 1H_milestone_output/gqa_headline/sweep.json. No
figure renderers yet (deferred).

Removals (pre-ADR-0060 baseline now superseded):
- benches: _attention_mesh_kv.py, _attention_mesh_mlo.py,
  _attention_mesh_mlo_2d.py, milestone_gqa_llama70b.py
- tests: test_attention_*, test_mesh_*, test_milestone_gqa_llama70b
- topology: llama70b_4sip.yaml (only consumer was the deleted diag)
- artifacts: 1H_milestone_output/gqa/ (sweep.json + 5 PNGs)
- tests/gqa/ plot helper + test (broken on Windows Tcl/Tkinter)
- ADR-0060/0061 references to deleted file paths cleaned up
  (EN + KO kept in sync).

Tests: 124/124 focused regression green (attention + Phase E + TL
discipline + triton_emu + pe_components). Full regression: 764 pass,
2 pre-existing test_bench_registry failures (stale EXPECTED_NAMES
across multiple benches, not introduced here).

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

tests/test_phase_e_cpu_issue_cost.py

@@ -0,0 +1,382 @@
+"""Phase 1 spec tests for ADR-0064 (per-op-type CPU issue cost model).
+
+Phase E lands the cost-table machinery and turns it on by default so the
+hybrid's CPU-saturation lever (ADR-0060 §1) becomes measurable instead of
+modelled away. Today every ``tl.*`` call goes through
+``_emit_dispatch_overhead()`` with a single uniform ``dispatch_cycles``
+scalar that is hard-coded to 0 on the live PE_CPU paths
+(``pe_cpu.py:_execute_legacy`` and ``kernel_runner.py:run``).
+
+These tests assume the post-Phase-2 surface:
+
+- ``kernbench.common.cpu_issue_cost`` exports ``OpKind``,
+  ``DEFAULT_CPU_ISSUE_COST`` (the ADR-0064 D1 table) and
+  ``get_issue_cost(kind, table=None) -> int``.
+- ``TLContext.__init__`` accepts ``issue_cost_table: dict[str, int] | None``.
+  When provided, ``_emit_dispatch_overhead(kind)`` looks up the per-kind
+  cost; when absent, it falls back to the uniform ``dispatch_cycles``
+  (back-compat with the existing ADR-0046 §D6 contract).
+- Live PE_CPU paths (greenlet + legacy replay) construct TLContext with
+  ``issue_cost_table=DEFAULT_CPU_ISSUE_COST``.
+
+Phase 1 (this commit): tests only. All tests FAIL until Phase 2.
+"""
+from __future__ import annotations
+
+from pathlib import Path
+
+import pytest
+
+from kernbench.common.pe_commands import (
+    CompositeCmd,
+    DmaReadCmd,
+    DmaWriteCmd,
+    GemmCmd,
+    MathCmd,
+    PeCpuOverheadCmd,
+)
+from kernbench.policy.address.phyaddr import PhysAddr
+from kernbench.runtime_api.kernel import KernelLaunchMsg, KernelRef
+from kernbench.sim_engine.engine import GraphEngine
+from kernbench.sim_engine.transaction import Transaction
+from kernbench.topology.builder import load_topology
+from kernbench.triton_emu.registry import clear_registry, register_kernel
+from kernbench.triton_emu.tl_context import TLContext, run_kernel
+
+TOPOLOGY_PATH = Path(__file__).parent.parent / "topology.yaml"
+
+
+def _engine():
+    return GraphEngine(load_topology(TOPOLOGY_PATH))
+
+
+def _hbm_pa(sip: int = 0, cube: int = 0, pe_id: int = 0) -> int:
+    slice_bytes = 48 * (1 << 30) // 8
+    pa = PhysAddr.pe_hbm_addr(
+        sip_id=sip, die_id=cube, pe_id=pe_id,
+        pe_local_hbm_offset=0x1000, slice_size_bytes=slice_bytes,
+    )
+    return pa.encode()
+
+
+# ── T1: default table shape ──────────────────────────────────────
+
+
+def test_default_cost_table_has_expected_keys():
+    """ADR-0064 D1 table — 8 keys with composite ≫ primitive ratio."""
+    from kernbench.common.cpu_issue_cost import DEFAULT_CPU_ISSUE_COST
+
+    expected_keys = {
+        "composite", "load", "store", "dot", "math",
+        "ipcq_send", "ipcq_recv", "copy_to",
+    }
+    assert set(DEFAULT_CPU_ISSUE_COST.keys()) == expected_keys, (
+        f"DEFAULT_CPU_ISSUE_COST keys must match ADR-0064 D1; "
+        f"got {set(DEFAULT_CPU_ISSUE_COST.keys())}"
+    )
+    # Composite is the lever — ratio against primitives must be ≥ 4×.
+    assert DEFAULT_CPU_ISSUE_COST["composite"] == 40
+    for primitive in ("load", "store", "dot", "math",
+                      "ipcq_send", "ipcq_recv", "copy_to"):
+        assert DEFAULT_CPU_ISSUE_COST[primitive] == 5, (
+            f"primitive {primitive!r} default cost must be 5 ns"
+        )
+
+
+# ── T2: get_issue_cost lookup ────────────────────────────────────
+
+
+def test_get_issue_cost_lookup():
+    """Helper returns table value; unknown kind returns 0 (no charge)."""
+    from kernbench.common.cpu_issue_cost import (
+        DEFAULT_CPU_ISSUE_COST,
+        get_issue_cost,
+    )
+
+    assert get_issue_cost("composite") == 40
+    assert get_issue_cost("load") == 5
+    assert get_issue_cost("unknown_kind") == 0
+    # Custom table override
+    custom = {"composite": 100, "load": 1}
+    assert get_issue_cost("composite", table=custom) == 100
+    assert get_issue_cost("load", table=custom) == 1
+    assert get_issue_cost("store", table=custom) == 0
+
+
+# ── T3: TLContext consumes a passed cost table ───────────────────
+
+
+def test_tlcontext_accepts_issue_cost_table():
+    """TLContext(issue_cost_table=...) → tl.load emits the per-kind cycles."""
+    tl = TLContext(
+        pe_id=0, num_programs=1,
+        dispatch_cycles=0,
+        issue_cost_table={"load": 7, "store": 3, "dot": 2, "math": 2},
+    )
+    tl.load(0x1000, shape=(4, 4), dtype="f16")
+    overheads = [c for c in tl.commands if isinstance(c, PeCpuOverheadCmd)]
+    assert len(overheads) == 1, (
+        f"expected exactly one PeCpuOverheadCmd before the DmaReadCmd; "
+        f"got {len(overheads)} (cmds={[type(c).__name__ for c in tl.commands]})"
+    )
+    assert overheads[0].cycles == 7, (
+        f"load issue cost from table must be 7; got {overheads[0].cycles}"
+    )
+
+
+# ── T4: composite ≫ primitive issue cost differential ────────────
+
+
+def test_tlcontext_composite_vs_primitive_charges_differ():
+    """Composite kernel charges once (40); primitive sequence charges per-op."""
+    from kernbench.common.cpu_issue_cost import DEFAULT_CPU_ISSUE_COST
+
+    # Composite kernel: 1 load + 1 composite = 5 + 40 = 45 cycles of issue cost.
+    tl_comp = TLContext(
+        pe_id=0, num_programs=1,
+        dispatch_cycles=0,
+        issue_cost_table=DEFAULT_CPU_ISSUE_COST,
+    )
+    a = tl_comp.load(0x1000, shape=(8, 16), dtype="f16")
+    b_ref = tl_comp.ref(0x2000, shape=(16, 8), dtype="f16")
+    tl_comp.composite("gemm", a, b_ref, out_ptr=0x3000)
+    comp_cycles = sum(
+        c.cycles for c in tl_comp.commands if isinstance(c, PeCpuOverheadCmd)
+    )
+
+    # Primitive kernel: 2 loads + 1 dot + 1 math (exp) = 5 + 5 + 5 + 5 = 20 cycles.
+    tl_prim = TLContext(
+        pe_id=0, num_programs=1,
+        dispatch_cycles=0,
+        issue_cost_table=DEFAULT_CPU_ISSUE_COST,
+    )
+    a = tl_prim.load(0x1000, shape=(8, 16), dtype="f16")
+    b = tl_prim.load(0x2000, shape=(16, 8), dtype="f16")
+    c_out = tl_prim.dot(a, b)
+    tl_prim.exp(c_out)
+    prim_cycles = sum(
+        c.cycles for c in tl_prim.commands if isinstance(c, PeCpuOverheadCmd)
+    )
+
+    # ADR-0064 ratio: composite issue cost >> primitive issue cost per op.
+    # For these specific kernels: composite=45 (5+40), primitive=20 (4×5).
+    assert comp_cycles == 45, f"composite kernel: expected 45, got {comp_cycles}"
+    assert prim_cycles == 20, f"primitive kernel: expected 20, got {prim_cycles}"
+    # Headline assertion: a single composite charges more than all 3
+    # post-load primitives combined (40 > 3×5) — the ADR-0060 §1 lever.
+    assert comp_cycles - 5 > 3 * 5, (
+        f"composite issue charge {comp_cycles - 5} must exceed "
+        f"3× primitive issue charge {3 * 5} (ADR-0064 ratio)"
+    )
+
+
+# ── T5: cost table is purely additive (Q2 — no double-count) ─────
+
+
+def test_cost_table_is_additive_only():
+    """Q2 invariant: the cost table is additive on PE_CPU, NOT folded into
+    DMA/GEMM/MATH command shapes. Two TLContexts running the same kernel
+    with different cost tables must produce identical non-overhead command
+    sequences (same DmaReadCmd, GemmCmd, MathCmd, addrs, shapes, dtypes).
+    Only the PeCpuOverheadCmd ``cycles`` field is allowed to differ.
+    """
+    from kernbench.common.cpu_issue_cost import DEFAULT_CPU_ISSUE_COST
+
+    def kernel(tl):
+        a = tl.load(0x1000, shape=(8, 16), dtype="f16")
+        b = tl.load(0x2000, shape=(16, 8), dtype="f16")
+        c = tl.dot(a, b)
+        d = tl.exp(c)
+        tl.store(0x3000, d)
+
+    tl_zero = TLContext(
+        pe_id=0, num_programs=1,
+        dispatch_cycles=0,
+        issue_cost_table={},  # empty table → every kind = 0 cost
+    )
+    run_kernel(kernel, tl_zero)
+
+    tl_default = TLContext(
+        pe_id=0, num_programs=1,
+        dispatch_cycles=0,
+        issue_cost_table=DEFAULT_CPU_ISSUE_COST,
+    )
+    run_kernel(kernel, tl_default)
+
+    # Strip PeCpuOverheadCmd from both streams; what remains must match.
+    non_overhead_zero = [
+        c for c in tl_zero.commands if not isinstance(c, PeCpuOverheadCmd)
+    ]
+    non_overhead_default = [
+        c for c in tl_default.commands if not isinstance(c, PeCpuOverheadCmd)
+    ]
+    assert len(non_overhead_zero) == len(non_overhead_default), (
+        f"non-overhead command count differs: "
+        f"{len(non_overhead_zero)} vs {len(non_overhead_default)}"
+    )
+    for a_cmd, b_cmd in zip(non_overhead_zero, non_overhead_default):
+        assert type(a_cmd) is type(b_cmd), (
+            f"non-overhead command type changed under cost table: "
+            f"{type(a_cmd).__name__} vs {type(b_cmd).__name__}"
+        )
+    # Total overhead under zero-table must be 0; under default must be > 0.
+    cycles_zero = sum(
+        c.cycles for c in tl_zero.commands if isinstance(c, PeCpuOverheadCmd)
+    )
+    cycles_default = sum(
+        c.cycles for c in tl_default.commands if isinstance(c, PeCpuOverheadCmd)
+    )
+    assert cycles_zero == 0, f"empty table must add 0 cycles; got {cycles_zero}"
+    assert cycles_default > 0, (
+        f"default table must add > 0 cycles; got {cycles_default}"
+    )
+
+
+# ── T6: greenlet vs legacy replay use same cost table (review #6) ─
+
+
+def test_greenlet_and_legacy_path_parity():
+    """Both PE_CPU execution paths read the same cost table.
+
+    The greenlet path (kernel_runner.py:run) and the legacy replay path
+    (pe_cpu.py:_execute_legacy) must construct TLContext with the same
+    default cost table so the same kernel produces identical PE_CPU
+    overhead cycles via either route. This is ADR-0064 review item #6.
+    """
+    from kernbench.common.cpu_issue_cost import DEFAULT_CPU_ISSUE_COST
+
+    # Build the command list for a representative kernel via TLContext
+    # using the default table — this is what both live paths should see.
+    def kernel(tl):
+        a = tl.load(0x1000, shape=(4, 4), dtype="f16")
+        b = tl.load(0x2000, shape=(4, 4), dtype="f16")
+        c = tl.dot(a, b)
+        tl.store(0x3000, c)
+
+    tl1 = TLContext(
+        pe_id=0, num_programs=1,
+        dispatch_cycles=0,
+        issue_cost_table=DEFAULT_CPU_ISSUE_COST,
+    )
+    run_kernel(kernel, tl1)
+    cycles1 = sum(
+        c.cycles for c in tl1.commands if isinstance(c, PeCpuOverheadCmd)
+    )
+
+    tl2 = TLContext(
+        pe_id=0, num_programs=1,
+        dispatch_cycles=0,
+        issue_cost_table=DEFAULT_CPU_ISSUE_COST,
+    )
+    run_kernel(kernel, tl2)
+    cycles2 = sum(
+        c.cycles for c in tl2.commands if isinstance(c, PeCpuOverheadCmd)
+    )
+
+    assert cycles1 == cycles2, (
+        f"same kernel via same default table must produce identical "
+        f"overhead cycles; got {cycles1} vs {cycles2}"
+    )
+    # Concrete expected for this kernel: 2×load(5) + 1×dot(5) + 1×store(5) = 20.
+    assert cycles1 == 20, (
+        f"expected 20 cycles total (2 load + 1 dot + 1 store at 5 ns); "
+        f"got {cycles1}"
+    )
+
+
+# ── T7: back-compat — no table → uniform dispatch_cycles ─────────
+
+
+def test_back_compat_no_table_uses_dispatch_cycles():
+    """ADR-0046 §D6 contract preserved when no issue_cost_table is passed.
+
+    Existing call sites doing ``TLContext(dispatch_cycles=0)`` must
+    continue to emit zero overhead. Existing call sites doing
+    ``TLContext(dispatch_cycles=1)`` must continue to emit uniform 1.
+    """
+    # Zero path (most existing tests use this).
+    tl_zero = TLContext(pe_id=0, num_programs=1, dispatch_cycles=0)
+    tl_zero.load(0x1000, shape=(4, 4), dtype="f16")
+    overheads = [c for c in tl_zero.commands if isinstance(c, PeCpuOverheadCmd)]
+    assert overheads == [], (
+        f"dispatch_cycles=0 with no table must emit no overhead; "
+        f"got {[c.cycles for c in overheads]}"
+    )
+
+    # Uniform path (test_dispatch_overhead_inserted relies on this).
+    tl_one = TLContext(pe_id=0, num_programs=1, dispatch_cycles=1)
+    tl_one.load(0x1000, shape=(4, 4), dtype="f16")
+    overheads = [c for c in tl_one.commands if isinstance(c, PeCpuOverheadCmd)]
+    assert overheads == [PeCpuOverheadCmd(cycles=1)], (
+        f"dispatch_cycles=1 with no table must emit cycles=1; "
+        f"got {[c.cycles for c in overheads]}"
+    )
+
+
+# ── T8: live PE_CPU constructs TLContext with the default table ──
+
+
+def test_live_pe_cpu_uses_default_table(monkeypatch):
+    """End-to-end: live PE_CPU greenlet path constructs TLContext with
+    ``issue_cost_table=DEFAULT_CPU_ISSUE_COST``. This is the wiring assertion
+    for ADR-0064 D4 "active by default" + review item #6 (path parity).
+
+    We patch ``TLContext.__init__`` to record its kwargs and run a
+    single-load kernel through the live PE_CPU. The recorded
+    ``issue_cost_table`` must equal ``DEFAULT_CPU_ISSUE_COST``.
+    """
+    from kernbench.common.cpu_issue_cost import DEFAULT_CPU_ISSUE_COST
+    from kernbench.triton_emu import tl_context as _tlc
+
+    captured: list[dict] = []
+    real_init = _tlc.TLContext.__init__
+
+    def spy_init(self, *args, **kwargs):
+        captured.append(dict(kwargs))
+        real_init(self, *args, **kwargs)
+
+    monkeypatch.setattr(_tlc.TLContext, "__init__", spy_init)
+
+    clear_registry()
+    hbm_pa = _hbm_pa(sip=0, cube=0, pe_id=0)
+
+    def single_load_kernel(tl):
+        tl.load(hbm_pa, shape=(4, 4), dtype="f16")
+
+    register_kernel("test_active_default_table", single_load_kernel)
+
+    engine = _engine()
+    pe_cpu_id = "sip0.cube0.pe0.pe_cpu"
+    done = engine._env.event()
+    txn = Transaction(
+        request=KernelLaunchMsg(
+            correlation_id="t", request_id="r",
+            kernel_ref=KernelRef(name="test_active_default_table", kind="builtin"),
+            args=(),
+        ),
+        path=[pe_cpu_id], step=0, nbytes=0, done=done,
+    )
+
+    def inject():
+        yield engine._components[pe_cpu_id]._inbox.put(txn)
+        yield done
+
+    engine._env.process(inject())
+    engine._env.run()
+    clear_registry()
+
+    # Live PE_CPU constructs TLContext on either the greenlet path
+    # (kernel_runner.py:run) or the legacy replay path
+    # (pe_cpu.py:_execute_legacy) — whichever is chosen depends on whether
+    # MemoryStore is wired. Both must pass DEFAULT_CPU_ISSUE_COST.
+    pe_ctx_calls = [c for c in captured if c.get("pe_id") == 0
+                    and "issue_cost_table" in c]
+    assert len(pe_ctx_calls) >= 1, (
+        f"expected at least one TLContext constructed by a live PE_CPU "
+        f"path with an issue_cost_table kwarg; got {len(pe_ctx_calls)} "
+        f"(all captures: {captured})"
+    )
+    last = pe_ctx_calls[-1]
+    assert last["issue_cost_table"] == DEFAULT_CPU_ISSUE_COST, (
+        f"live PE_CPU must use DEFAULT_CPU_ISSUE_COST; got {last['issue_cost_table']}"
+    )