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,148 @@
+"""Phase 1 spec test for P7: headline milestone-gqa bench with real GQA.
+
+P7 wires the new ``_gqa_decode`` and ``_gqa_prefill`` kernels into a
+new 4-panel milestone bench (independent from the existing
+``milestone-gqa-llama70b`` which still covers the baseline kernels).
+Real GQA (``h_q > h_kv`` with G=8) runs end-to-end through a
+milestone-style sweep + sweep.json output.
+
+Restriction in P7 first cut:
+  - C ≤ 4 (single-row inter-CUBE ring SFR; ADR-0060 §B leaves
+    multi-row rings for follow-on)
+  - Single SIP (default ``topology.yaml`` 4×4 cube mesh; 4-SIP
+    headline left to follow-on)
+  - No figure renderers (defer to a separate cycle)
+
+Panels (4 total):
+  single_user_prefill_gqa: C=1, T_q=4, S_kv=16   (no ring)
+  multi_user_prefill_gqa : C=4, T_q=4, S_kv=16   (Ring KV, 3 steps)
+  single_user_decode_gqa : C=1, P=8, h_q=8, h_kv=1, S_kv=64    (M-fold + intra-cube chain)
+  multi_user_decode_gqa  : C=4, P=8, h_q=8, h_kv=1, S_kv=128   (M-fold + 2-level chain)
+
+Phase 1 (this commit): tests only — bench module lands in Phase 2.
+"""
+from __future__ import annotations
+
+import json
+
+import kernbench.benches.milestone_gqa_headline as bench_mod  # noqa: F401 (Phase 2)
+from kernbench.benches.registry import resolve
+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
+
+BENCH_NAME = "milestone-gqa-headline"
+
+PANELS = (
+    "single_user_prefill_gqa",
+    "multi_user_prefill_gqa",
+    "single_user_decode_gqa",
+    "multi_user_decode_gqa",
+)
+
+
+def _run_validation():
+    topo = resolve_topology("topology.yaml")
+    return run_bench(
+        topology=topo,
+        bench_fn=resolve(BENCH_NAME).run,
+        device=resolve_device(None),
+        engine_factory=lambda t, d: GraphEngine(
+            getattr(t, "topology_obj", t), enable_data=True,
+        ),
+    )
+
+
+def _sweep_json(monkeypatch) -> dict:
+    monkeypatch.setenv("GQA_HEADLINE_RUN", "1")
+    out = bench_mod._OUTPUT_DIR / "sweep.json"
+    if not out.exists():
+        result = _run_validation()
+        assert result.completion.ok, result.completion
+    assert out.exists(), f"missing {out}"
+    return json.loads(out.read_text())
+
+
+# ── Registration ───────────────────────────────────────────────────────
+
+
+def test_bench_registered():
+    spec = resolve(BENCH_NAME)
+    assert spec.name == BENCH_NAME
+    assert callable(spec.run)
+    assert spec.description.strip(), "description must be non-empty"
+
+
+# ── Validation run completes ──────────────────────────────────────────
+
+
+def test_validation_run_completes_ok(monkeypatch):
+    monkeypatch.setenv("GQA_HEADLINE_RUN", "1")
+    result = _run_validation()
+    assert result.completion.ok, (
+        f"headline validation run failed: {result.completion}"
+    )
+
+
+# ── sweep.json shape ──────────────────────────────────────────────────
+
+
+def test_sweep_json_has_four_panels(monkeypatch):
+    data = _sweep_json(monkeypatch)
+    assert set(data["panels"]) == set(PANELS), (
+        f"panels mismatch: expected {set(PANELS)}, got {set(data['panels'])}"
+    )
+    assert len(data["rows"]) == 4
+    assert {r["panel"] for r in data["rows"]} == set(PANELS)
+
+
+def test_decode_panels_use_real_gqa(monkeypatch):
+    """ADR-0060 §A.1 headline: decode panels must use h_q = G·h_kv with G>1."""
+    data = _sweep_json(monkeypatch)
+    cfg = data["config"]
+    assert cfg["h_q_decode"] > cfg["h_kv_decode"], (
+        f"decode must use real GQA (h_q > h_kv); got "
+        f"h_q={cfg['h_q_decode']}, h_kv={cfg['h_kv_decode']}"
+    )
+
+
+# ── Per-panel architectural assertions ────────────────────────────────
+
+
+def _row(rows, panel: str) -> dict:
+    for r in rows:
+        if r["panel"] == panel:
+            return r
+    raise AssertionError(f"missing row for panel {panel!r}")
+
+
+def test_prefill_ring_panel_has_ipcq_traffic(monkeypatch):
+    """multi_user_prefill_gqa uses Ring KV → IPCQ traffic > 0."""
+    data = _sweep_json(monkeypatch)
+    row = _row(data["rows"], "multi_user_prefill_gqa")
+    n_copy = row["op_log_summary"].get("ipcq_copy_count", 0)
+    assert n_copy > 0, (
+        f"multi_user_prefill_gqa must have Ring KV traffic; got "
+        f"ipcq_copy_count={n_copy}"
+    )
+
+
+def test_decode_reduce_panel_writes_once(monkeypatch):
+    """multi_user_decode_gqa: chain reduce-to-root → exactly 1 dma_write."""
+    data = _sweep_json(monkeypatch)
+    row = _row(data["rows"], "multi_user_decode_gqa")
+    n_writes = row["op_log_summary"]["dma_write_count"]
+    assert n_writes == 1, (
+        f"multi_user_decode_gqa root-only write: expected 1; got {n_writes}"
+    )
+
+
+def test_prefill_panel_distributes_output(monkeypatch):
+    """multi_user_prefill_gqa: per-CUBE distributed output → dma_write_count == C."""
+    data = _sweep_json(monkeypatch)
+    row = _row(data["rows"], "multi_user_prefill_gqa")
+    n_writes = row["op_log_summary"]["dma_write_count"]
+    assert n_writes == 4, (
+        f"multi_user_prefill_gqa per-CUBE distributed: expected 4; got {n_writes}"
+    )