kernbench2/tests/attention/test_gqa_prefill.py

"""Phase 1 spec test for P6a GQA prefill kernel (head-parallel, C=1 baseline).

P6a introduces ``_gqa_attention_prefill_long.py`` with the head-parallel structure (one
Q head per CUBE, per-CUBE distributed output, no reduce). C=1 is the
degenerate case — no Ring KV, no IPCQ traffic. Validates kernel
structure and T_q > 1 attention.

P6b (deferred) adds the Ring KV rotation for C > 1, which needs either
a new SFR install function (intra_* + wrapped E/W at CUBE level) or a
topology-specific config — separate design call.

The prefill kernel differs from decode (P1a/P2a/P2b) in three ways
(ADR-0060 §5.5 / TL;DR):
  1. Q has T_q > 1 rows (not just decode's single timestep).
  2. Head-parallel placement: each CUBE owns ONE Q head — no M-fold.
  3. Each CUBE writes its own head's output — NO reduce.

Phase 1 (this commit): tests only — production code lands in Phase 2.
"""
from __future__ import annotations

from pathlib import Path

from kernbench.benches._gqa_attention_prefill_long import gqa_attention_prefill_long_kernel  # noqa: F401
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"

D_HEAD = 64
DTYPE = "f16"


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


def _run_prefill(*, T_q: int, S_kv: int, C: int = 1):
    """C=1 head-parallel prefill: single CUBE owns the one head + full KV."""
    topo = resolve_topology(str(TOPOLOGY_DEFAULT))

    def _bench_fn(ctx):
        dp = DPPolicy(cube="replicate", pe="replicate",
                      num_cubes=C, num_pes=1)
        # Q: (T_q, d_head) — one head per CUBE (head-parallel; for C=1
        # only one head total). 2D layout matches what the kernel loads.
        # P6b will use a 3D (h_q, T_q, d_head) Q with cube_row_wise
        # sharding so each CUBE owns its head.
        q = ctx.zeros((T_q, D_HEAD), dtype=DTYPE, dp=dp,
                      name=f"q_t{T_q}_c{C}")
        # K, V: full local for C=1 (no ring). Kernel loads K as
        # (d_head, S_kv) via byte-conserving reshape.
        k = ctx.zeros((S_kv, D_HEAD), dtype=DTYPE, dp=dp,
                      name=f"k_t{T_q}_c{C}")
        v = ctx.zeros((S_kv, D_HEAD), dtype=DTYPE, dp=dp,
                      name=f"v_t{T_q}_c{C}")
        # O: (T_q, d_head) — per-CUBE distributed output.
        o = ctx.empty((T_q, D_HEAD), dtype=DTYPE, dp=dp,
                      name=f"o_t{T_q}_c{C}")
        ctx.launch(
            f"gqa_prefill_p6a_t{T_q}_s{S_kv}_c{C}",
            gqa_attention_prefill_long_kernel,
            q, k, v, o,
            T_q, S_kv, D_HEAD, C,
            _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)


def test_prefill_c_one_t_q_one_completes():
    """C=1, T_q=1: smallest workload (decode-like)."""
    result = _run_prefill(T_q=1, S_kv=16, C=1)
    assert result.completion.ok, (
        f"prefill C=1 T_q=1 failed: {result.completion}"
    )


def test_prefill_c_one_t_q_four_completes():
    """C=1, T_q=4: real prefill (Q has multiple rows) — distinguishes
    prefill from decode (T_q=1)."""
    result = _run_prefill(T_q=4, S_kv=16, C=1)
    assert result.completion.ok, (
        f"prefill C=1 T_q=4 failed: {result.completion}"
    )


def test_prefill_c_one_no_ipcq_traffic():
    """C=1: no ring step, no IPCQ traffic."""
    result = _run_prefill(T_q=4, S_kv=16, C=1)
    assert result.completion.ok
    n_copy = _count(result.engine.op_log, "ipcq_copy")
    assert n_copy == 0, (
        f"C=1 must have no IPCQ traffic (no ring); got {n_copy}"
    )


def test_prefill_c_one_one_dma_write():
    """C=1, one head: exactly one dma_write (per-CUBE distributed output;
    no reduce). For C > 1 in P6b this becomes dma_write_count == C."""
    result = _run_prefill(T_q=4, S_kv=16, C=1)
    assert result.completion.ok
    n_writes = _count(result.engine.op_log, "dma_write")
    assert n_writes == 1, (
        f"C=1 prefill: expected 1 dma_write (one head per cube); "
        f"got {n_writes}"
    )