kernbench2/src/kernbench/benches/milestone_gqa_headline.py

"""milestone-gqa-headline bench: real GQA + 2-level SP + Ring KV.

Wires the new ``_gqa_decode`` and ``_gqa_prefill`` kernels through 4
panels with real GQA (h_q = G·h_kv, G > 1 on the decode side), writing
per-panel ``op_log_summary`` into ``sweep.json``. Independent from the
existing ``milestone-gqa-llama70b`` validation-scale bench (which stays
on the legacy baseline kernels).

Restrictions (P7 first cut):
  - C ≤ 4 (single-row inter-CUBE ring SFR; multi-row deferred)
  - Single SIP (default ``topology.yaml`` 4×4 cube mesh; 4-SIP
    headline deferred)
  - No figure renderers (defer to a separate cycle)

Panels:
  single_user_prefill_gqa : prefill C=1, T_q=4, S_kv=16
  multi_user_prefill_gqa  : prefill C=4 Ring KV, T_q=4, S_kv=16
  single_user_decode_gqa  : decode C=1, P=8, h_q=8, h_kv=1, S_kv=64
                            (M-fold + intra-cube row-then-col chain)
  multi_user_decode_gqa   : decode C=4, P=8, h_q=8, h_kv=1, S_kv=128
                            (M-fold + 2-level chain reduce-to-root)

Gated by ``GQA_HEADLINE_RUN=1``.
"""
from __future__ import annotations

import json
import os
from pathlib import Path

from kernbench.benches._gqa_attention_decode_long import gqa_attention_decode_long_kernel
from kernbench.benches._gqa_attention_prefill_long import gqa_attention_prefill_long_kernel
from kernbench.benches.registry import bench
from kernbench.ccl.install import load_ccl_config, resolve_algorithm_config
from kernbench.ccl.sfr_config import (
    configure_sfr_intercube_multisip,
    configure_sfr_intercube_ring,
)
from kernbench.policy.placement.dp import DPPolicy

_OUTPUT_DIR = Path(__file__).resolve().parent / "1H_milestone_output" / "gqa_headline"
_SWEEP_JSON = _OUTPUT_DIR / "sweep.json"

# ── Panel configs ────────────────────────────────────────────────────

_DTYPE = "f16"
_D_HEAD = 64
_T_Q_PREFILL = 4
_T_Q_DECODE = 1
_S_KV_PREFILL = 16
_H_Q_DECODE = 8        # real GQA: G = H_Q_DECODE / H_KV_DECODE = 8
_H_KV_DECODE = 1

_PANELS = (
    "single_user_prefill_gqa",
    "multi_user_prefill_gqa",
    "single_user_decode_gqa",
    "multi_user_decode_gqa",
)

# Each entry: (kind, panel-specific params)
_PANEL_DISPATCH: dict[str, tuple[str, dict]] = {
    "single_user_prefill_gqa": ("prefill", {"C": 1, "S_kv": _S_KV_PREFILL}),
    "multi_user_prefill_gqa":  ("prefill", {"C": 4, "S_kv": _S_KV_PREFILL}),
    "single_user_decode_gqa":  ("decode",  {"C": 1, "P": 8, "S_kv": 64}),
    "multi_user_decode_gqa":   ("decode",  {"C": 4, "P": 8, "S_kv": 128}),
}


def _ccl_cfg():
    return resolve_algorithm_config(
        load_ccl_config(), name="lrab_hierarchical_allreduce",
    )


# ── Per-kind launch helpers ──────────────────────────────────────────


def _run_prefill_panel(ctx, *, panel: str, C: int, S_kv: int) -> None:
    if C > 1:
        configure_sfr_intercube_ring(
            ctx.engine, ctx.spec, _ccl_cfg(), ring_size=C,
        )
    dp_q = DPPolicy(cube="replicate", pe="replicate",
                    num_cubes=C, num_pes=1)
    dp_kv = DPPolicy(cube="row_wise" if C > 1 else "replicate",
                     pe="replicate", num_cubes=C, num_pes=1)
    dp_o = DPPolicy(cube="row_wise" if C > 1 else "replicate",
                    pe="replicate", num_cubes=C, num_pes=1)
    q = ctx.zeros((_T_Q_PREFILL, _D_HEAD),
                  dtype=_DTYPE, dp=dp_q, name=f"{panel}_q")
    k = ctx.zeros((S_kv, _D_HEAD),
                  dtype=_DTYPE, dp=dp_kv, name=f"{panel}_k")
    v = ctx.zeros((S_kv, _D_HEAD),
                  dtype=_DTYPE, dp=dp_kv, name=f"{panel}_v")
    o = ctx.empty((_T_Q_PREFILL * C, _D_HEAD),
                  dtype=_DTYPE, dp=dp_o, name=f"{panel}_o")
    ctx.launch(
        panel, gqa_attention_prefill_long_kernel,
        q, k, v, o,
        _T_Q_PREFILL, S_kv, _D_HEAD, C,
        _auto_dim_remap=False,
    )


def _run_decode_panel(ctx, *, panel: str, C: int, P: int, S_kv: int) -> None:
    configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
    dp_full = DPPolicy(cube="replicate", pe="replicate",
                       num_cubes=C, num_pes=P)
    dp_kv = DPPolicy(cube="row_wise" if C > 1 else "replicate",
                     pe="row_wise", num_cubes=C, num_pes=P)
    q = ctx.zeros((_T_Q_DECODE, _H_Q_DECODE * _D_HEAD),
                  dtype=_DTYPE, dp=dp_full, name=f"{panel}_q")
    k = ctx.zeros((S_kv, _H_KV_DECODE * _D_HEAD),
                  dtype=_DTYPE, dp=dp_kv, name=f"{panel}_k")
    v = ctx.zeros((S_kv, _H_KV_DECODE * _D_HEAD),
                  dtype=_DTYPE, dp=dp_kv, name=f"{panel}_v")
    o = ctx.empty((_T_Q_DECODE, _H_Q_DECODE * _D_HEAD),
                  dtype=_DTYPE, dp=dp_full, name=f"{panel}_o")
    ctx.launch(
        panel, gqa_attention_decode_long_kernel,
        q, k, v, o,
        _T_Q_DECODE, S_kv, _H_Q_DECODE, _H_KV_DECODE, _D_HEAD, C, P,
        _auto_dim_remap=False,
    )


def _make_bench_fn(panel: str):
    kind, params = _PANEL_DISPATCH[panel]

    def _bench_fn(ctx):
        if kind == "prefill":
            _run_prefill_panel(ctx, panel=panel, **params)
        else:
            _run_decode_panel(ctx, panel=panel, **params)

    return _bench_fn


# ── Op-log summary ──────────────────────────────────────────────────


def _summarize_op_log(op_log) -> dict[str, int]:
    """Per-panel op_log counts (gemm, ipcq_copy, dma_read, dma_write)."""
    gemm_count = 0
    ipcq_copy_count = 0
    dma_read_count = 0
    dma_write_count = 0
    for r in op_log:
        if r.op_kind == "gemm":
            gemm_count += 1
        elif r.op_name == "dma_read":
            dma_read_count += 1
        elif r.op_name == "dma_write":
            dma_write_count += 1
        elif r.op_name == "ipcq_copy":
            ipcq_copy_count += 1
    return {
        "gemm_count": gemm_count,
        "ipcq_copy_count": ipcq_copy_count,
        "dma_read_count": dma_read_count,
        "dma_write_count": dma_write_count,
    }


def _run_panel(panel: str, topology: str) -> dict:
    """Run one panel in a fresh engine; return its row dict."""
    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

    topo = resolve_topology(topology)
    result = run_bench(
        topology=topo, bench_fn=_make_bench_fn(panel),
        device=resolve_device(None),
        engine_factory=lambda t, d: GraphEngine(
            getattr(t, "topology_obj", t), enable_data=True,
        ),
    )
    if not result.completion.ok:
        raise RuntimeError(
            f"milestone-gqa-headline panel {panel!r} failed: "
            f"{result.completion}"
        )
    kind, params = _PANEL_DISPATCH[panel]
    return {
        "panel": panel,
        "kind": kind,
        **params,
        "op_log_summary": _summarize_op_log(result.engine.op_log),
    }


# ── Bench entry ──────────────────────────────────────────────────────


@bench(
    name="milestone-gqa-headline",
    description="Headline GQA milestone — real GQA h_q=8/h_kv=1 + 2-level SP (decode) + Ring KV (prefill).",
)
def run(torch) -> None:
    """Drive 4 headline panels through the new GQA kernels; write sweep.json.

    Gated by GQA_HEADLINE_RUN=1.
    """
    _OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    if not os.environ.get("GQA_HEADLINE_RUN"):
        raise RuntimeError(
            "milestone-gqa-headline needs GQA_HEADLINE_RUN=1."
        )

    topology = os.environ.get("GQA_HEADLINE_TOPOLOGY", "topology.yaml")
    rows = [_run_panel(panel, topology) for panel in _PANELS]
    sweep = {
        "version": 1,
        "panels": list(_PANELS),
        "config": {
            "T_q_prefill": _T_Q_PREFILL,
            "T_q_decode": _T_Q_DECODE,
            "S_kv_prefill": _S_KV_PREFILL,
            "h_q_decode": _H_Q_DECODE,
            "h_kv_decode": _H_KV_DECODE,
            "d_head": _D_HEAD,
        },
        "rows": rows,
    }
    _SWEEP_JSON.write_text(json.dumps(sweep, indent=2))
    print(f"  milestone-gqa-headline: {len(rows)} rows -> {_SWEEP_JSON}")

    # Sentinel tensor (ADR-0045 D4 / ADR-0054 D2 carve-out).
    torch.zeros(
        (1, 1), dtype="f16",
        dp=DPPolicy(cube="row_wise", pe="replicate", num_cubes=1, num_pes=1),
        name="milestone_gqa_headline_sentinel",
    )