kernbench2/tests/attention/test_gqa_prefill_ring.py

"""Phase 1 spec test for P6b GQA prefill Ring KV (head-parallel, C>1).

P6b adds the Ring KV rotation (ADR-0060 §5.5) to the prefill kernel.
Each CUBE owns one Q head + its KV slice; over C ring steps the KV
blocks rotate around the C CUBEs so every CUBE sees every block. The
running ``(m, ℓ, O)`` is folded inside the loop. No reduce — each CUBE
writes its own head's output.

Requires a new SFR install ``configure_sfr_intercube_ring(ring_size=C)``
that wires:
  - ``intra_*``  : 2×4 PE grid within a CUBE (same as multisip)
  - ``E/W``      : 1D ring of CUBEs 0..ring_size-1 WITH WRAP
                    (symmetric to ``configure_sfr_intracube_pe_ring``,
                     applied at CUBE level)
  - ``global_*`` : SIP-level (same as multisip)

The kernel ring body:
  for step in range(1, C):
      tl.send(dir="W", src=Kc)
      tl.send(dir="W", src=Vc)
      Kc = tl.recv(dir="E", ...)
      Vc = tl.recv(dir="E", ...)
      ... local partial + online-softmax merge into (m, ℓ, O) ...

Per CUBE per step: 2 sends (K, V) → 2 ``ipcq_copy``. Across all CUBEs:
``(C-1) * 2 * C`` total ipcq_copy.

Restriction in P6b first cut: ``C ∈ {1, 2, 4}`` (single row of the 4×4
cube mesh). C=8 ring would span rows — follow-on.

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

from pathlib import Path

from kernbench.benches._gqa_prefill_long import gqa_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  # noqa: F401 (Phase 2)
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 _run_prefill_ring(*, T_q: int, S_kv: int, C: int):
    """Head-parallel prefill with Ring KV across C CUBEs."""
    topo = resolve_topology(str(TOPOLOGY_DEFAULT))

    def _bench_fn(ctx):
        # P6b: new SFR install with cube-level ring wrap.
        configure_sfr_intercube_ring(
            ctx.engine, ctx.spec, _ccl_cfg(), ring_size=C,
        )
        # Q replicated on every CUBE (zeros for testing; per-CUBE head
        # indexing is implicit). KV sequence-sharded by CUBE. O
        # distributed — each CUBE writes its slice of (T_q*C, d_head).
        dp_q = DPPolicy(cube="replicate", pe="replicate",
                        num_cubes=C, num_pes=1)
        dp_kv = DPPolicy(cube="row_wise", pe="replicate",
                         num_cubes=C, num_pes=1)
        dp_o = DPPolicy(cube="row_wise", pe="replicate",
                        num_cubes=C, num_pes=1)
        q = ctx.zeros((T_q, D_HEAD), dtype=DTYPE, dp=dp_q,
                      name=f"q_t{T_q}_c{C}_ring")
        k = ctx.zeros((S_kv, D_HEAD), dtype=DTYPE, dp=dp_kv,
                      name=f"k_t{T_q}_c{C}_ring")
        v = ctx.zeros((S_kv, D_HEAD), dtype=DTYPE, dp=dp_kv,
                      name=f"v_t{T_q}_c{C}_ring")
        # O: (T_q * C, d_head), each CUBE writes (T_q, d_head) at its
        # slice. dma_write_count = C (per-CUBE distributed output).
        o = ctx.empty((T_q * C, D_HEAD), dtype=DTYPE, dp=dp_o,
                      name=f"o_t{T_q}_c{C}_ring")
        ctx.launch(
            f"gqa_prefill_ring_t{T_q}_s{S_kv}_c{C}",
            gqa_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)


# ── C=2 Ring KV ────────────────────────────────────────────────────────


def test_prefill_ring_c_two_completes():
    """C=2: 1 ring step rotates KV between the 2 CUBEs."""
    result = _run_prefill_ring(T_q=4, S_kv=16, C=2)
    assert result.completion.ok, (
        f"prefill ring C=2 failed: {result.completion}"
    )


def test_prefill_ring_c_two_distributed_output():
    """C=2: per-CUBE distributed output, no reduce → dma_write_count == 2."""
    result = _run_prefill_ring(T_q=4, S_kv=16, C=2)
    assert result.completion.ok
    n_writes = _count(result.engine.op_log, "dma_write")
    assert n_writes == 2, (
        f"C=2 prefill: expected 2 dma_write (one per CUBE); got {n_writes}"
    )


def test_prefill_ring_c_two_ipcq_count():
    """C=2: 1 ring step × 2 handles (K, V) × 2 CUBEs = 4 ipcq_copy."""
    result = _run_prefill_ring(T_q=4, S_kv=16, C=2)
    assert result.completion.ok
    n_copy = _count(result.engine.op_log, "ipcq_copy")
    expected = (2 - 1) * 2 * 2
    assert n_copy == expected, (
        f"C=2 ring: expected {expected} ipcq_copy "
        f"((C-1)·2·C = 1·2·2); got {n_copy}"
    )


# ── C=4 Ring KV (combined assertions) ─────────────────────────────────


def test_prefill_ring_c_four_combined():
    """C=4: 3 ring steps rotate KV around 4 CUBEs.
    Expected: completes; 4 dma_writes; (4-1)·2·4 = 24 ipcq_copy."""
    result = _run_prefill_ring(T_q=4, S_kv=32, C=4)
    assert result.completion.ok, (
        f"prefill ring C=4 failed: {result.completion}"
    )
    n_writes = _count(result.engine.op_log, "dma_write")
    assert n_writes == 4, (
        f"C=4 prefill: expected 4 dma_write (one per CUBE); got {n_writes}"
    )
    n_copy = _count(result.engine.op_log, "ipcq_copy")
    expected = (4 - 1) * 2 * 4
    assert n_copy == expected, (
        f"C=4 ring: expected {expected} ipcq_copy "
        f"((C-1)·2·C = 3·2·4); got {n_copy}"
    )