kernbench2/tests/attention/test_gqa_prefill_long_pe_sp.py

"""Tests for intra-CUBE PE-SP in prefill_long (query-axis split).

ADR-0060 §5.5 last bullet + §B-item-3: split the head's query rows
``[T_q, d_head]`` across the ``P`` PEs of each CUBE so all P PEs work
in parallel. Output rows are disjoint across PEs ⇒ no intra-CUBE
reduce needed.

Same-lane SFR wiring: PE ``i`` in CUBE A has its own E/W ring link to
PE ``i`` in CUBE B (the snake's prev/next). All P PEs of a CUBE see
the same K, V (HBM-resident, ``pe="replicate"``) ⇒ P parallel rings
run in lockstep, each PE rotating its own K/V copies via its own IPCQ
channels.

Activation contract:
  - ``P == 1`` (default; omitted from launch args) → existing
    PE-0-only behavior. Byte-for-byte unchanged.
  - ``P > 1`` → all P PEs active; each handles ``T_q // P`` query
    rows. Requires ``T_q % P == 0`` (degenerate T_q < P is rejected
    — caller must use ``P=1`` for that workload, ADR-0060 §B-item-3).

Phase 1: tests only — production code lands in Phase 2.
T1, T2, T3, T5 fail today (TypeError: kernel signature has no P).
T4 passes today as the backward-compat anchor.
"""
from __future__ import annotations

from pathlib import Path

import pytest

from kernbench.benches._gqa_attention_prefill_long import gqa_attention_prefill_long_kernel  # noqa: F401
from kernbench.ccl.install import load_ccl_config, resolve_algorithm_config
from kernbench.ccl.sfr_config import configure_sfr_intercube_ring
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)


def _run_prefill_pe_sp(
    *, T_q: int, S_kv: int, C: int, P: int | None,
    snake: bool,
):
    """Drive a prefill_long launch with optional intra-CUBE PE-SP.

    ``P=None`` → omit P from the launch args (exercises the kernel's
    default behaviour).
    ``snake=True`` → install the 2×4 snake ring SFR (Increment 1);
    else install the 1D-row ring at ``ring_size=C``.
    """
    topo = resolve_topology(str(TOPOLOGY_DEFAULT))

    def _bench_fn(ctx):
        if snake:
            configure_sfr_intercube_ring(
                ctx.engine, ctx.spec, _ccl_cfg(),
                submesh_shape=(2, 4),
            )
        else:
            configure_sfr_intercube_ring(
                ctx.engine, ctx.spec, _ccl_cfg(),
                ring_size=C,
            )
        num_pes = P if (P is not None and P > 1) else 1
        # When PE-SP is active, Q and O are split row-wise across PEs;
        # K, V remain replicated within a CUBE.
        q_pe = "row_wise" if num_pes > 1 else "replicate"
        o_pe = "row_wise" if num_pes > 1 else "replicate"
        dp_q = DPPolicy(cube="replicate", pe=q_pe,
                        num_cubes=C, num_pes=num_pes)
        dp_kv = DPPolicy(cube="row_wise" if C > 1 else "replicate",
                         pe="replicate", num_cubes=C, num_pes=num_pes)
        dp_o = DPPolicy(cube="row_wise" if C > 1 else "replicate",
                        pe=o_pe, num_cubes=C, num_pes=num_pes)
        suffix = f"c{C}_p{P}_t{T_q}_s{S_kv}"
        q = ctx.zeros((T_q, D_HEAD), dtype=DTYPE, dp=dp_q,
                      name=f"q_pe_sp_{suffix}")
        k = ctx.zeros((S_kv, D_HEAD), dtype=DTYPE, dp=dp_kv,
                      name=f"k_pe_sp_{suffix}")
        v = ctx.zeros((S_kv, D_HEAD), dtype=DTYPE, dp=dp_kv,
                      name=f"v_pe_sp_{suffix}")
        o = ctx.empty((T_q * C, D_HEAD), dtype=DTYPE, dp=dp_o,
                      name=f"o_pe_sp_{suffix}")
        launch_args = [q, k, v, o, T_q, S_kv, D_HEAD, C]
        if P is not None:
            launch_args.append(P)
        ctx.launch(
            f"gqa_prefill_long_pe_sp_{suffix}",
            gqa_attention_prefill_long_kernel,
            *launch_args,
            _auto_dim_remap=False,
        )

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


# ── T1: C=8 P=8 completes end-to-end ─────────────────────────────────


def test_prefill_c8_p8_pe_sp_completes():
    """ADR-0060 §5.5 + §B-item-3 design target: 64 ranks active
    (8 CUBEs × 8 PEs), query-axis split across PEs, snake-mapped
    Ring KV. Guards no scratch overflow, no deadlock across the
    P parallel same-lane rings.
    """
    result = _run_prefill_pe_sp(
        T_q=8, S_kv=8192, C=8, P=8, snake=True,
    )
    assert result.completion.ok, (
        f"prefill at C=8, P=8 (PE-SP) must complete; "
        f"got {result.completion}"
    )


# ── T2: 64 dma_writes (one per PE, disjoint query rows) ──────────────


def test_prefill_c8_p8_pe_sp_64_dma_writes():
    """With query-axis split, each PE owns ``T_q/P = 1`` query row
    and stores its own ``(1, d_head)`` slice. Across all 64 ranks
    (8 CUBEs × 8 PEs), ``dma_write_count == 64``.

    This is the structural signal that PE-SP is actually wired:
    PE-0-only would give 8 dma_writes (one per CUBE).
    """
    result = _run_prefill_pe_sp(
        T_q=8, S_kv=8192, C=8, P=8, snake=True,
    )
    assert result.completion.ok
    n_writes = _count(result.engine.op_log, "dma_write")
    assert n_writes == 64, (
        f"PE-SP at C=8, P=8: expected 64 dma_writes "
        f"(8 cubes × 8 PEs, each writes its T_q/P=1 query row); "
        f"got {n_writes}"
    )


# ── T3: P parallel rings — ipcq_copy scales by P ─────────────────────


def test_prefill_c8_p8_pe_sp_ring_ipcq_count():
    """Per-PE same-lane rings: each PE runs its own ring traffic via
    its own IPCQ channels. Total inter-CUBE ipcq_copy:
    ``(C-1) · n_tiles · 2 · C · P`` (= existing C=8 formula × P).

    Configuration: T_q=8, S_kv=8192, C=8 → S_local=1024, n_tiles=1
    (TILE_S_KV=1024). Expected = ``7·1·2·8·8`` = **896**.
    """
    C = 8
    P = 8
    n_tiles = 1
    result = _run_prefill_pe_sp(
        T_q=8, S_kv=8192, C=C, P=P, snake=True,
    )
    assert result.completion.ok
    n_copy = _count(result.engine.op_log, "ipcq_copy")
    expected = (C - 1) * n_tiles * 2 * C * P
    assert n_copy == expected, (
        f"PE-SP parallel rings at C=8, P=8: expected {expected} "
        f"ipcq_copy ((C-1)·n_tiles·2·C·P = "
        f"{C - 1}·{n_tiles}·2·{C}·{P}); got {n_copy}"
    )


# ── T4: backward-compat — P omitted → existing PE-0-only behaviour ───


def test_prefill_p_default_1_backward_compat():
    """When ``P`` is omitted from the launch args, the kernel must
    behave identically to today: only PE 0 of each CUBE participates;
    one head per CUBE; one dma_write per CUBE.

    At C=4 this gives 4 dma_writes (matches the existing
    ``test_prefill_long_tile_ring_dma_write_count``). Phase 2 must
    not regress this path.

    Passes today AND after Phase 2.
    """
    C = 4
    result = _run_prefill_pe_sp(
        T_q=4, S_kv=8192, C=C, P=None, snake=False,
    )
    assert result.completion.ok, (
        f"prefill at C=4 with default P (PE-0-only) must complete; "
        f"got {result.completion}"
    )
    n_writes = _count(result.engine.op_log, "dma_write")
    assert n_writes == C, (
        f"default P=1 (PE-0-only): expected {C} dma_writes "
        f"(one per CUBE); got {n_writes}"
    )


# ── T5: validation — T_q must be divisible by P when P > 1 ───────────


def test_prefill_pe_sp_rejects_non_divisible_t_q():
    """ADR-0060 §B-item-3 fallback (KV-block split + intra-CUBE
    reduce for T_q < P) is deferred. Callers must request ``P=1`` for
    workloads where T_q < P or T_q % P != 0.

    C=4, P=8, T_q=4 violates ``T_q % P == 0`` (and also T_q < P).
    The kernel must raise ValueError with a clear error message
    rather than silently producing wrong results.
    """
    with pytest.raises((ValueError, AssertionError), match=r"T_q"):
        _run_prefill_pe_sp(
            T_q=4, S_kv=8192, C=4, P=8, snake=False,
        )