kernbench2/tests/attention/test_milestone_gqa_headline.py

"""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_kv_group_prefill_gqa_c8_p8",
    "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_expected_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"]) == len(PANELS)
    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}"
    )