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>
2026-06-09 18:15:59 -07:00
parent b3730a33eb
commit d282144339
52 changed files with 3952 additions and 2526 deletions
@@ -0,0 +1,153 @@
+"""Phase 1 spec test for P2a GQA decode SP (chain reduce-to-root, Level-2 only).
+
+P2a is the first half of DDD-0060 P2: the kernel becomes multi-PE within
+one CUBE and reduces to root (PE 0) using a chain over the 1D intra-cube
+ring (W direction). This **replaces the baseline's bidirectional O(N)
+fan-out** where every rank ends with O — ADR-0060 §A.2's headline.
+
+Deviation from DDD-0060 §7 P2 gate: the gate text asks for
+``⌈log₂ P⌉`` reduce rounds. The intra-cube SFR install
+(``configure_sfr_intracube_pe_ring``) wires only a 1D E/W ring, so a
+true tree would require either multi-hop forwarding or a new SFR install
+(future ADR). P2a uses a **chain reduce-to-root**: ``P-1`` rounds along
+W. The architectural property the ADR cares about
+(root-only output vs every-rank-has-O) is preserved; the logarithmic
+collective is deferred.
+
+P2b (deferred) covers Level-1 inter-CUBE center-mesh reduce (C>1).
+
+Phase 1 (this commit): tests only — production code lands in Phase 2.
+"""
+from __future__ import annotations
+
+from pathlib import Path
+
+from kernbench.benches._gqa_decode import gqa_decode_kernel  # noqa: F401
+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"
+
+T_Q = 1
+D_HEAD = 64
+DTYPE = "f16"
+
+
+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_decode_sp(*, h_q: int, h_kv: int, P: int, S_kv: int):
+    """Single-CUBE SP decode: P PEs share the work along the intra-cube ring."""
+    topo = resolve_topology(str(TOPOLOGY_DEFAULT))
+
+    def _bench_fn(ctx):
+        configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
+        dp_full = DPPolicy(cube="replicate", pe="replicate",
+                           num_cubes=1, num_pes=P)
+        dp_kv = DPPolicy(cube="replicate", pe="row_wise",
+                         num_cubes=1, num_pes=P)
+        q = ctx.zeros((T_Q, h_q * D_HEAD),
+                      dtype=DTYPE, dp=dp_full, name=f"q_h{h_q}_kv{h_kv}_p{P}")
+        # KV: total S_kv split across P PEs along axis 0 (row_wise sharding).
+        # Each PE sees (S_kv/P, h_kv·D_HEAD).
+        k = ctx.zeros((S_kv, h_kv * D_HEAD),
+                      dtype=DTYPE, dp=dp_kv, name=f"k_h{h_q}_kv{h_kv}_p{P}")
+        v = ctx.zeros((S_kv, h_kv * D_HEAD),
+                      dtype=DTYPE, dp=dp_kv, name=f"v_h{h_q}_kv{h_kv}_p{P}")
+        o = ctx.empty((T_Q, h_q * D_HEAD),
+                      dtype=DTYPE, dp=dp_full, name=f"o_h{h_q}_kv{h_kv}_p{P}")
+        ctx.launch(
+            f"gqa_decode_sp_h{h_q}_kv{h_kv}_p{P}",
+            gqa_decode_kernel,
+            q, k, v, o,
+            T_Q, S_kv, h_q, h_kv, D_HEAD,
+            1, P,                                 # C=1, P=P (single-CUBE SP)
+            _auto_dim_remap=False,
+        )
+
+    return run_bench(
+        topology=topo,
+        bench_fn=_bench_fn,
+        device=resolve_device(None),
+        engine_factory=_engine_factory,
+    )
+
+
+def _count(op_log, name: str) -> int:
+    return sum(1 for r in op_log if r.op_name == name)
+
+
+# ── Root-only write ────────────────────────────────────────────────────
+
+
+def test_sp_chain_reduce_root_only_writes_o():
+    """ADR-0060 §A.2: only the root rank (PE 0) writes O. Baseline today
+    has every rank write the full final O (bidirectional fan-out)."""
+    result = _run_decode_sp(h_q=1, h_kv=1, P=8, S_kv=64)
+    assert result.completion.ok, f"P=8 chain reduce failed: {result.completion}"
+    n_writes = _count(result.engine.op_log, "dma_write")
+    assert n_writes == 1, (
+        f"reduce-to-root must produce exactly 1 dma_write (PE 0); "
+        f"got {n_writes}"
+    )
+
+
+# ── Chain step count ───────────────────────────────────────────────────
+
+
+def test_sp_chain_reduce_p_minus_one_ipcq_pairs():
+    """Chain reduce-to-root has P-1 send→recv pairs along the W chain;
+    each pair logs one ``ipcq_copy`` (inbound DMA, per
+    ``milestone_gqa_llama70b._summarize_op_log``). Each chain step ships
+    the triplet (m, ℓ, O) → 3 handles per step → 7 steps × 3 = 21."""
+    result = _run_decode_sp(h_q=1, h_kv=1, P=8, S_kv=64)
+    assert result.completion.ok, f"P=8 chain reduce failed: {result.completion}"
+    n_copy = _count(result.engine.op_log, "ipcq_copy")
+    expected = (8 - 1) * 3
+    assert n_copy == expected, (
+        f"chain reduce: expected {expected} ipcq_copy (P-1=7 steps × "
+        f"3 handles m/ℓ/O); got {n_copy}"
+    )
+
+
+# ── Real GQA × SP combined ─────────────────────────────────────────────
+
+
+def test_sp_real_gqa_h_q_eight_h_kv_one_p_eight():
+    """The combined unlock: real GQA (h_q=G·h_kv with G=8) AND SP
+    (P=8) together — neither expressible by the baseline."""
+    result = _run_decode_sp(h_q=8, h_kv=1, P=8, S_kv=64)
+    assert result.completion.ok, (
+        f"real GQA + SP combined run failed: {result.completion}"
+    )
+    n_writes = _count(result.engine.op_log, "dma_write")
+    assert n_writes == 1, (
+        f"root-only write must hold under M-fold too; got {n_writes}"
+    )
+
+
+# ── Degenerate P=1 ────────────────────────────────────────────────────
+
+
+def test_sp_p_one_degenerate_no_ipcq_traffic():
+    """P=1: SP degenerates to a single rank. No IPCQ traffic; one dma_write."""
+    result = _run_decode_sp(h_q=8, h_kv=1, P=1, S_kv=16)
+    assert result.completion.ok, f"P=1 degenerate failed: {result.completion}"
+    n_send = _count(result.engine.op_log, "ipcq_send")
+    n_recv = _count(result.engine.op_log, "ipcq_recv")
+    assert n_send == 0, f"P=1 must have no ipcq_send; got {n_send}"
+    assert n_recv == 0, f"P=1 must have no ipcq_recv; got {n_recv}"
+    n_writes = _count(result.engine.op_log, "dma_write")
+    assert n_writes == 1, f"P=1: one dma_write; got {n_writes}"