kernbench2/tests/attention/test_gqa_tile_sweep.py

"""Phase 1 spec test for P3b: S_kv tile sweep in the GQA kernels
(ADR-0063 §A.2 + ADR-0060 §B "long/short context split").

The local one-shot partial in three kernels — ``_gqa_attention_decode_long.py``,
``_gqa_attention_decode_short.py``, ``_gqa_attention_prefill_short.py`` — loads its rank's
entire ``(d_head, S_local)`` / ``(S_local, d_head)`` KV slice in one
shot, so per-rank scratch grows linearly with ``S_local``. The
``_attention_*`` baselines hit a TCM ceiling once ``S_local`` exceeds
the 1 MiB pool divided by the per-rank intermediate footprint.

P3b replaces the one-shot partial with a tile sweep:
  - Module-level ``TILE_S_KV = 1024`` per kernel file.
  - Tile 0 establishes the persistent ``(m_local, l_local, O_local)``.
  - Tiles 1..n_tiles-1 wrap their intermediates (K_T tile, V tile,
    scores, centered, exp_scores, partials) in ``tl.scratch_scope()``
    and persist the merged ``(m, ℓ, O)`` to the persistent handles via
    ``tl.copy_to`` (ADR-0063 §D3 / §D3.1).

When ``S_local ≤ TILE_S_KV`` the loop body never runs and the op_log
is structurally identical to today's one-shot path — every existing
validation-scale test continues to pass.

``_gqa_attention_prefill_long.py`` is **out of scope** for P3b. Its inner step
is IPCQ partial-recv (Ring KV rotation), not an HBM load; tiling it
intersects the ring-step structure and is a separate phase.

Phase 1 (this commit): tests only — production code lands in Phase 2.
The four multi-tile tests fail today because the kernels never call
``copy_to`` outside the chain-reduce merges (so isolating to a
chain-free config gives ``copy_to == 0`` today). The 128K test fails
today with a ``TLContext`` scratch overflow.
"""
from __future__ import annotations

from pathlib import Path

from kernbench.benches._gqa_attention_decode_long import gqa_attention_decode_long_kernel  # noqa: F401
from kernbench.benches._gqa_attention_decode_short import gqa_attention_decode_short_kernel  # noqa: F401
from kernbench.benches._gqa_attention_prefill_short import gqa_attention_prefill_short_kernel  # noqa: F401
from kernbench.ccl.install import load_ccl_config, resolve_algorithm_config
from kernbench.ccl.sfr_config import configure_sfr_intercube_multisip
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 _ccl_cfg():
    return resolve_algorithm_config(
        load_ccl_config(), name="lrab_hierarchical_allreduce",
    )


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


def _count(op_log, name: str) -> int:
    return sum(1 for r in op_log if r.op_name == name)


# ── decode_long runner (chain-free configs isolate tile-sweep behavior) ──


def _run_decode_long(*, C: int, P: int, S_kv: int, h_q: int = 1, h_kv: int = 1):
    """Run the long-context decode kernel.

    For tile-sweep isolation, callers pass C=1, P=1 — this leaves both
    chain-reduce levels inactive so any ``copy_to`` in op_log comes
    solely from the tile-merge body.
    """
    topo = resolve_topology(str(TOPOLOGY_DEFAULT))

    def _bench_fn(ctx):
        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" if P > 1 else "replicate",
                         num_cubes=C, num_pes=P)
        q = ctx.zeros((1, h_q * D_HEAD),
                      dtype=DTYPE, dp=dp_full, name=f"q_tl_c{C}_p{P}_s{S_kv}")
        k = ctx.zeros((S_kv, h_kv * D_HEAD),
                      dtype=DTYPE, dp=dp_kv, name=f"k_tl_c{C}_p{P}_s{S_kv}")
        v = ctx.zeros((S_kv, h_kv * D_HEAD),
                      dtype=DTYPE, dp=dp_kv, name=f"v_tl_c{C}_p{P}_s{S_kv}")
        o = ctx.empty((1, h_q * D_HEAD),
                      dtype=DTYPE, dp=dp_full, name=f"o_tl_c{C}_p{P}_s{S_kv}")
        ctx.launch(
            f"gqa_decode_long_tile_c{C}_p{P}_s{S_kv}",
            gqa_attention_decode_long_kernel,
            q, k, v, o,
            1, S_kv, h_q, h_kv, D_HEAD, C, P,
            _auto_dim_remap=False,
        )

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


# ── decode_short / prefill_short runners ─────────────────────────────


def _run_decode_short(*, kv_per_cube: int, C: int, P: int, S_kv: int,
                      h_q: int = 8, h_kv: int = 8):
    """For tile-sweep isolation: kv_per_cube=P → group_size=1 (no chain)."""
    topo = resolve_topology(str(TOPOLOGY_DEFAULT))

    def _bench_fn(ctx):
        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((1, h_q * D_HEAD),
                      dtype=DTYPE, dp=dp_full, name=f"q_dsh_s{S_kv}")
        k = ctx.zeros((h_kv * S_kv, D_HEAD),
                      dtype=DTYPE, dp=dp_kv, name=f"k_dsh_s{S_kv}")
        v = ctx.zeros((h_kv * S_kv, D_HEAD),
                      dtype=DTYPE, dp=dp_kv, name=f"v_dsh_s{S_kv}")
        o = ctx.empty((1, h_q * D_HEAD),
                      dtype=DTYPE, dp=dp_full, name=f"o_dsh_s{S_kv}")
        ctx.launch(
            f"gqa_decode_short_tile_s{S_kv}",
            gqa_attention_decode_short_kernel,
            q, k, v, o,
            1, S_kv, h_q, h_kv, D_HEAD, C, P, kv_per_cube,
            _auto_dim_remap=False,
        )

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


def _run_prefill_short(*, kv_per_cube: int, C: int, P: int,
                       T_q: int, S_kv: int, h_kv: int = 8):
    topo = resolve_topology(str(TOPOLOGY_DEFAULT))

    def _bench_fn(ctx):
        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, h_kv * D_HEAD),
                      dtype=DTYPE, dp=dp_full, name=f"q_psh_s{S_kv}")
        k = ctx.zeros((h_kv * S_kv, D_HEAD),
                      dtype=DTYPE, dp=dp_kv, name=f"k_psh_s{S_kv}")
        v = ctx.zeros((h_kv * S_kv, D_HEAD),
                      dtype=DTYPE, dp=dp_kv, name=f"v_psh_s{S_kv}")
        o = ctx.empty((T_q, h_kv * D_HEAD),
                      dtype=DTYPE, dp=dp_full, name=f"o_psh_s{S_kv}")
        ctx.launch(
            f"gqa_prefill_short_tile_s{S_kv}",
            gqa_attention_prefill_short_kernel,
            q, k, v, o,
            T_q, S_kv, h_kv, D_HEAD, C, P, kv_per_cube,
            _auto_dim_remap=False,
        )

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


# ── T1: single-tile path is op_log-stable (regression guard) ─────────


def test_decode_long_single_tile_path_unchanged():
    """ADR-0063 §A.2: when S_local <= TILE_S_KV the tile-sweep loop
    body never runs and op_log is structurally identical to today's
    one-shot path.

    Passes today (one-shot path) AND after Phase 2 (n_tiles=1 skips
    the merge loop). The witness: with chain reduce inactive (C=1,
    P=1), there must be zero ``copy_to`` entries — neither today nor
    after Phase 2 should the kernel emit tile-merge writebacks here.
    """
    result = _run_decode_long(C=1, P=1, S_kv=64)
    assert result.completion.ok, (
        f"single-tile decode_long must complete; got {result.completion}"
    )
    n_copy = _count(result.engine.op_log, "copy")
    assert n_copy == 0, (
        f"S_local <= TILE_S_KV must not emit tile-merge copy_to; got {n_copy}"
    )


# ── T2: decode_long multi-tile emits tile-merge copy_to ──────────────


def test_decode_long_multi_tile_emits_copy_to_merges():
    """ADR-0063 §A.2 + §D3.1: when S_local > TILE_S_KV the kernel must
    wrap each subsequent tile's intermediates in ``scratch_scope`` and
    persist the merged ``(m, ℓ, O)`` via ``tl.copy_to``.

    Chain-free config (C=1, P=1) isolates the tile-merge body — any
    ``copy_to`` in op_log comes from the tile sweep, not chain reduce.

    S_kv=2048, C=1, P=1 → S_local=2048 → 2 tiles → 1 merge × 3 handles
    (m, ℓ, O) ⇒ 3 ``copy`` entries.

    Currently 0 because the kernel performs a one-shot partial.
    """
    result = _run_decode_long(C=1, P=1, S_kv=2048)
    assert result.completion.ok, (
        f"multi-tile decode_long must complete; got {result.completion}"
    )
    n_copy = _count(result.engine.op_log, "copy")
    assert n_copy >= 3, (
        f"decode_long multi-tile must emit >= 3 copy_to entries "
        f"(1 merge × 3 handles); got {n_copy}"
    )


# ── T3: decode_short multi-tile emits tile-merge copy_to ─────────────


def test_decode_short_multi_tile_emits_copy_to_merges():
    """Same property as T2 for the short decode kernel.

    kv_per_cube=8, P=8, C=1 → group_size=1 (no chain reduce); S_kv=2048
    → S_local=2048 → 2 tiles per PE → 3 ``copy`` entries per PE × 8 PEs
    = 24 total.

    Currently 0 because the kernel performs a one-shot partial.
    """
    result = _run_decode_short(kv_per_cube=8, C=1, P=8, S_kv=2048)
    assert result.completion.ok, (
        f"multi-tile decode_short must complete; got {result.completion}"
    )
    n_copy = _count(result.engine.op_log, "copy")
    assert n_copy >= 3 * 8, (
        f"decode_short multi-tile must emit >= 24 copy_to entries "
        f"(1 merge × 3 handles × 8 PEs); got {n_copy}"
    )


# ── T4: prefill_short multi-tile emits tile-merge copy_to ────────────


def test_prefill_short_multi_tile_emits_copy_to_merges():
    """Same property as T3 for the short prefill kernel.

    kv_per_cube=8, P=8, C=1, T_q=4, S_kv=2048 → group_size=1 (no chain),
    S_local=2048 → 2 tiles per PE → 3 ``copy`` entries per PE × 8 PEs
    = 24 total.

    Currently 0 because the kernel performs a one-shot partial.
    """
    result = _run_prefill_short(
        kv_per_cube=8, C=1, P=8, T_q=4, S_kv=2048,
    )
    assert result.completion.ok, (
        f"multi-tile prefill_short must complete; got {result.completion}"
    )
    n_copy = _count(result.engine.op_log, "copy")
    assert n_copy >= 3 * 8, (
        f"prefill_short multi-tile must emit >= 24 copy_to entries "
        f"(1 merge × 3 handles × 8 PEs); got {n_copy}"
    )


# ── T5: decode_long at S_kv=256K completes (TCM ceiling lifted) ──────


def test_decode_long_context_256k_completes():
    """ADR-0063 §A.2 (test req 3): headline ceiling-lift. C=1 P=8
    decode at S_kv=256K → S_local=32K. Today the per-rank score stack
    (3 ×M·S_local·2 bytes, M=G·T_q=8) is ~1.5 MB → exceeds the 1 MiB
    scratch pool → TLContext scratch overflow. After Phase 2 the tile
    sweep bounds per-tile score stack at 3·M·TILE_S_KV·2 ≈ 48 KB and
    the run completes regardless of S_kv.
    """
    result = _run_decode_long(C=1, P=8, S_kv=262144, h_q=8, h_kv=1)
    assert result.completion.ok, (
        f"decode_long at S_kv=256K must complete after tile sweep; "
        f"got {result.completion}"
    )