kernbench2/tests/attention/test_gqa_decode.py

"""Phase 1 spec test for P1a GQA decode kernel (real GQA via M-fold).

P1a is the first phase of the DDD-0060 plan, split out of the original
P1 (the composite-hybrid swap is P1b, deferred until the tl.composite
output-handle question is decided). P1a is the *correctness* unlock:
the kernel processes ONE KV head at a time and folds the G query heads
into the matmul M (row) dimension so a single Q·Kᵀ serves all G heads
sharing one K (ADR-0060 §5.2). This lifts the baseline's
``h_q == h_kv == 1`` cap pinned at
``tests/attention/test_milestone_gqa_llama70b.py:137-142``.

P1a stays inside the existing ``tl`` API: the two attention GEMMs use the
blocking ``tl.dot`` so the chain ``Q·Kᵀ → softmax → P·V → store`` fits
without any composite-output chaining. The composite swap (P1b) will
revisit this once the API for feeding a composite's output into a
downstream MATH op is settled.

Phase 1 (this commit): tests only — production code lands in Phase 2.
Tests fail at import in Phase 1 with ModuleNotFoundError; Phase 2 makes
them pass.
"""
from __future__ import annotations

from pathlib import Path

from kernbench.benches._gqa_decode_long import gqa_decode_long_kernel  # noqa: F401  (Phase 2 deliverable)
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"

# Decode shapes — P1a is one-shot, no tiling, single rank. P3 will tile.
T_Q = 1
D_HEAD = 64
S_KV = 16
DTYPE = "f16"


def _engine_factory(t, d):
    return GraphEngine(getattr(t, "topology_obj", t), enable_data=True)


def _run_decode_p1(*, h_q: int, h_kv: int):
    """One-shot GQA decode on a single PE in a single CUBE (P=1, no SP).

    Tensor layout (natural K — same as the P2a SP tests; the kernel
    reshapes K to (d_head, S_local) via byte-conserving load):
      Q : (T_q, h_q · d_head)   — natural Q layout; kernel reshapes to
          (G·T_q, d_head) — byte-conserving and math-correct because
          T_q=1 collapses axis ordering.
      K : (S_kv, h_kv · d_head) — natural K layout; kernel loads as
          (d_head, S_local) — reshape-as-transpose caveat (ADR-0060 §3
          / §B item 2), correct for zero inputs used here.
      V : (S_kv, h_kv · d_head) — natural V layout.
      O : (T_q, h_q · d_head)   — same shape as Q.
    """
    topo = resolve_topology(str(TOPOLOGY_DEFAULT))

    def _bench_fn(ctx):
        dp = DPPolicy(cube="replicate", pe="replicate",
                      num_cubes=1, num_pes=1)
        q = ctx.zeros((T_Q, h_q * D_HEAD),
                      dtype=DTYPE, dp=dp, name=f"q_h{h_q}_kv{h_kv}")
        k = ctx.zeros((S_KV, h_kv * D_HEAD),
                      dtype=DTYPE, dp=dp, name=f"k_h{h_q}_kv{h_kv}")
        v = ctx.zeros((S_KV, h_kv * D_HEAD),
                      dtype=DTYPE, dp=dp, name=f"v_h{h_q}_kv{h_kv}")
        o = ctx.empty((T_Q, h_q * D_HEAD),
                      dtype=DTYPE, dp=dp, name=f"o_h{h_q}_kv{h_kv}")
        ctx.launch(
            f"gqa_decode_p1_h{h_q}_kv{h_kv}",
            gqa_decode_long_kernel,
            q, k, v, o,
            T_Q, S_KV, h_q, h_kv, D_HEAD,
            1, 1,                               # C=1, P=1 (no SP, degenerate)
            _auto_dim_remap=False,
        )

    return run_bench(
        topology=topo,
        bench_fn=_bench_fn,
        device=resolve_device(None),
        engine_factory=_engine_factory,
    )


def _dma_read_count(op_log) -> int:
    return sum(1 for r in op_log if r.op_name == "dma_read")


# ── Headline unlock: real GQA (h_q = G·h_kv) runs in data mode ──────────


def test_real_gqa_h_q_eight_h_kv_one_completes_in_data_mode():
    """ADR-0060 §A.1 headline unlock — the baseline raises
    ``ValueError: Shape mismatch …`` in MemoryStore at h_q=8, h_kv=1
    because ``_view(K, (h_q·d, S_kv))`` only byte-conserves when h_q==h_kv.
    M-fold processes one KV head at a time using only byte-conserving
    reshapes, so this completes.
    """
    result = _run_decode_p1(h_q=8, h_kv=1)
    assert result.completion.ok, (
        f"real GQA (h_q=8, h_kv=1) decode failed: {result.completion}"
    )


# ── M-fold property: K/V loads do not scale with G ─────────────────────


def test_kv_dma_read_count_independent_of_g():
    """ADR-0060 TL;DR / §5.2: M-fold loads K and V once per KV head and
    folds the G query heads into the GEMM M dim. The dma_read_count must
    therefore be identical between (G=1, h_kv=1) and (G=8, h_kv=1) — both
    issue exactly 3 reads (Q + K + V). This pins the GQA-reuse property
    that the rest of the plan (composite streaming in P4, etc.) builds on.
    """
    g1 = _run_decode_p1(h_q=1, h_kv=1)
    g8 = _run_decode_p1(h_q=8, h_kv=1)
    n_g1 = _dma_read_count(g1.engine.op_log)
    n_g8 = _dma_read_count(g8.engine.op_log)
    assert n_g1 == 3, f"G=1 dma_read_count must be 3 (Q+K+V); got {n_g1}"
    assert n_g8 == 3, f"G=8 dma_read_count must be 3 (Q+K+V); got {n_g8}"
    assert n_g1 == n_g8, (
        f"K/V dma_read_count must be independent of G; "
        f"got G=1 -> {n_g1}, G=8 -> {n_g8}"
    )


# ── Backward-compat: degenerate G=1 still works ────────────────────────


def test_degenerate_g_equals_one_still_works():
    """G=1 (h_q == h_kv == 1) is the baseline-compatible config. M-fold
    degenerates to (T_q, d) = (1, 64) — the same shape the baseline
    already exercises — so this proves no regression on that path.
    """
    result = _run_decode_p1(h_q=1, h_kv=1)
    assert result.completion.ok, (
        f"degenerate G=1 decode failed: {result.completion}"
    )