kernbench2/tests/attention/test_attention_8kv_groups_diag.py

"""Diagnostic harness for the Llama-70B "1 Q-head per cube" target.

Per the GQA Llama-70B sharding study at
``llm_paper_review/notes/GQA_MHA_sharding/scripts/_gen_llama70b_1M_4cases.py``,
the 1 Q-head/cube baseline uses 64 cubes (4 SIPs × 16 cubes/SIP) organized
into 8 KV-groups of 8 cubes each. Each KV-group occupies a ``2×4``
sub-mesh within a SIP's ``4×4`` cube grid and runs the C2 2D row-then-col
AllReduce-mlo (ADR-0059 extension). This harness probes the gap between
validation and headline in three incrementally-larger steps:

  step_1_single_kv_group_at_full_breadth
     ONE multi_user_decode launch on a 2×4 sub-mesh (8 cubes) of the
     4-SIP topology, via the 2D mesh-mlo kernel. Verifies the per-KV-group
     2D AllReduce works at full breadth. Smallest dim possible
     (S_q=1, S_kv=16, h=1, d_head=64) to keep wall time bounded.

  step_2_four_kv_groups_one_per_sip
     Four sequential multi_user_decode launches, each targeting a
     different SIP. Verifies that per-SIP isolation works (each SIP holds
     its own 2×4 KV-group; the SFR install only writes intra-SIP
     E/W + N/S edges so the 4 groups don't see each other).

  step_3_eight_kv_groups_two_per_sip
     The actual study target: 8 KV-groups, two per SIP (cubes 0..7 vs
     cubes 8..15 within each SIP). Expected to FAIL with current infra —
     DPPolicy doesn't take a cube offset and target_device is SIP-level,
     so back-to-back launches both land on cubes 0..7 of their target SIP.
     Captures what's needed to lift the 4-group cap to 8.

Each step prints what it observed; the test asserts only the documented
expected outcomes so we can land it, watch CI, and iterate. Steps that
are *expected* to fail (step 3) are marked xfail with a precise reason.
"""
from __future__ import annotations

import traceback
from pathlib import Path

import pytest

from kernbench.benches._attention_mesh_mlo_2d import attention_mesh_mlo_2d_kernel
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 DeviceSelector, resolve_device
from kernbench.sim_engine.engine import GraphEngine
from kernbench.topology.builder import resolve_topology

TOPOLOGY_4SIP = (
    Path(__file__).resolve().parents[2] / "topologies" / "llama70b_4sip.yaml"
)
TOPOLOGY_DEFAULT = Path(__file__).resolve().parents[2] / "topology.yaml"

S_Q_DECODE = 1
S_KV_PER_RANK = 16
H_Q = 1
H_KV = 1
D_HEAD = 64
# 2×4 sub-mesh per KV-group (study: 8 cubes per KV-group at Q/cube=1).
MESH_ROWS = 2
MESH_COLS = 4
N_CUBES_PER_KV_GROUP = MESH_ROWS * MESH_COLS
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 _make_one_kv_group_bench(mesh_rows: int, mesh_cols: int):
    """Return a bench_fn that runs ONE multi_user_decode kernel on a
    ``mesh_rows × mesh_cols`` sub-mesh."""
    n_cubes = mesh_rows * mesh_cols

    def _bench_fn(ctx):
        configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
        dp_full = DPPolicy(cube="replicate", pe="replicate",
                           num_cubes=n_cubes, num_pes=8)
        dp_kv = DPPolicy(cube="row_wise", pe="replicate",
                         num_cubes=n_cubes, num_pes=8)
        q = ctx.zeros((S_Q_DECODE, H_Q * D_HEAD),
                      dtype=DTYPE, dp=dp_full, name="q")
        k = ctx.zeros((S_KV_PER_RANK * n_cubes, H_KV * D_HEAD),
                      dtype=DTYPE, dp=dp_kv, name="k")
        v = ctx.zeros((S_KV_PER_RANK * n_cubes, H_KV * D_HEAD),
                      dtype=DTYPE, dp=dp_kv, name="v")
        o = ctx.empty((S_Q_DECODE, H_Q * D_HEAD),
                      dtype=DTYPE, dp=dp_full, name="o")
        ctx.launch(
            f"single_kv_group_{mesh_rows}x{mesh_cols}",
            attention_mesh_mlo_2d_kernel,
            q, k, v, o,
            S_Q_DECODE, S_KV_PER_RANK, H_Q, H_KV, D_HEAD,
            mesh_rows, mesh_cols,
            1,                               # rank_axis=1 → cube-level ring
            0,                               # cube_start=0 — single sub-mesh launch
            _auto_dim_remap=False,
        )

    return _bench_fn


def _run_one_kv_group(topology_path: Path, mesh_rows: int, mesh_cols: int,
                      target_device=None):
    topo = resolve_topology(str(topology_path))
    captured: dict = {"engine": None}

    def factory(t, d):
        eng = _engine_factory(t, d)
        captured["engine"] = eng
        return eng

    exc = None
    result = None
    try:
        result = run_bench(
            topology=topo,
            bench_fn=_make_one_kv_group_bench(mesh_rows, mesh_cols),
            device=target_device or resolve_device(None),
            engine_factory=factory,
        )
    except BaseException as e:  # noqa: BLE001
        exc = e
    return exc, result, captured["engine"]


# ── Step 1 — single KV-group at the study's full breadth ──────────


def test_step_1_single_kv_group_at_full_breadth():
    """One multi_user_decode launch on a 2×4 sub-mesh, 4-SIP topology.

    Uses the C2 2D row-then-col AllReduce-mlo kernel: stage 1 reduces
    across cols (E/W) within each row, stage 2 reduces across rows (N/S).
    Expected to PASS — N/S edges are wired by
    ``configure_sfr_intercube_multisip`` and the 2D fan-out avoids the
    row-boundary IpcqInvalidDirection that the 1D kernel hit at cube 4.
    """
    if not TOPOLOGY_4SIP.exists():
        pytest.skip(f"4-SIP topology missing: {TOPOLOGY_4SIP}")
    exc, result, engine = _run_one_kv_group(
        TOPOLOGY_4SIP, mesh_rows=MESH_ROWS, mesh_cols=MESH_COLS,
    )
    if exc is not None:
        oplog_len = len(getattr(engine, "op_log", []) or []) if engine else 0
        print(f"\nstep_1 FAIL — op_log records before crash: {oplog_len}")
        traceback.print_exception(type(exc), exc, exc.__traceback__)
        raise AssertionError(f"step_1 failed: {exc}") from exc
    assert result is not None and result.completion.ok, (
        f"step_1: completion not ok — {result.completion if result else None}"
    )


# ── Step 2 — 4 KV-groups, one per SIP, sequential launches ────────


def _make_multi_sip_bench_fn(sip_groups: list[tuple[int, str, int]]):
    """One bench_fn that does one 2×4 multi_user_decode launch per item.

    Each ``(sip, tag, cube_start)`` tuple becomes one launch:
      - ``ctx.ahbm.set_device(sip)`` switches allocations to that SIP
        (mirrors ``milestone_1h_ccl.py:283-292``).
      - ``cube_start`` selects which 8-cube sub-mesh within the SIP:
        ``0`` → cubes 0..7 (rows 0..1), ``8`` → cubes 8..15 (rows 2..3).
      - ``tag`` disambiguates tensor names so launches in the same
        run_bench don't collide on the allocator namespace.
    """
    n_cubes = MESH_ROWS * MESH_COLS

    def _bench_fn(ctx):
        configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
        for sip, tag, cube_start in sip_groups:
            ctx.ahbm.set_device(sip)
            dp_full = DPPolicy(cube="replicate", pe="replicate",
                               num_cubes=n_cubes, num_pes=8,
                               cube_start=cube_start)
            dp_kv = DPPolicy(cube="row_wise", pe="replicate",
                             num_cubes=n_cubes, num_pes=8,
                             cube_start=cube_start)
            q = ctx.zeros((S_Q_DECODE, H_Q * D_HEAD),
                          dtype=DTYPE, dp=dp_full, name=f"q_{tag}")
            k = ctx.zeros((S_KV_PER_RANK * n_cubes, H_KV * D_HEAD),
                          dtype=DTYPE, dp=dp_kv, name=f"k_{tag}")
            v = ctx.zeros((S_KV_PER_RANK * n_cubes, H_KV * D_HEAD),
                          dtype=DTYPE, dp=dp_kv, name=f"v_{tag}")
            o = ctx.empty((S_Q_DECODE, H_Q * D_HEAD),
                          dtype=DTYPE, dp=dp_full, name=f"o_{tag}")
            ctx.launch(
                f"kv_group_{tag}", attention_mesh_mlo_2d_kernel,
                q, k, v, o,
                S_Q_DECODE, S_KV_PER_RANK, H_Q, H_KV, D_HEAD,
                MESH_ROWS, MESH_COLS,
                1,                               # rank_axis=1 → cube-level ring
                cube_start,                      # converts physical id → launch-local rank
                _auto_dim_remap=False,
            )

    return _bench_fn


def _run_multi_sip(sip_groups: list[tuple[int, str, int]]):
    """Run a single run_bench call covering all (sip, tag) groups."""
    topo = resolve_topology(str(TOPOLOGY_4SIP))
    captured: dict = {"engine": None}

    def factory(t, d):
        eng = _engine_factory(t, d)
        captured["engine"] = eng
        return eng

    exc = None
    result = None
    try:
        result = run_bench(
            topology=topo,
            bench_fn=_make_multi_sip_bench_fn(sip_groups),
            device=resolve_device(None),  # "all" SIPs in scope
            engine_factory=factory,
        )
    except BaseException as e:  # noqa: BLE001
        exc = e
    return exc, result, captured["engine"]


def test_step_2_four_kv_groups_one_per_sip():
    """Four multi_user_decode launches, one per SIP, in ONE run_bench call.

    Uses the CCL milestone pattern (``milestone_1h_ccl.py:283-292``):
    ``target_device="all"`` scopes the runtime to every SIP; then
    ``ctx.ahbm.set_device(sip)`` before each ``ctx.zeros``/``launch``
    switches which SIP the next allocation+launch lands on. This is the
    canonical sequential per-SIP pattern in the codebase — four separate
    ``run_bench`` calls with ``DeviceSelector("sip:N")`` is a misuse.

    Expected to PASS — the SFR install draws intra-SIP edges only, so the
    4 KV-groups can't see each other.
    """
    if not TOPOLOGY_4SIP.exists():
        pytest.skip(f"4-SIP topology missing: {TOPOLOGY_4SIP}")
    sip_groups = [(sip, f"sip{sip}", 0) for sip in range(4)]
    exc, result, engine = _run_multi_sip(sip_groups)
    if exc is not None:
        oplog_len = len(getattr(engine, "op_log", []) or []) if engine else 0
        print(f"\nstep_2 FAIL — op_log records before crash: {oplog_len}")
        traceback.print_exception(type(exc), exc, exc.__traceback__)
        raise AssertionError(f"step_2 failed: {exc}") from exc
    assert result is not None and result.completion.ok, (
        f"step_2: completion not ok — {result.completion if result else None}"
    )


# ── Step 3 — 8 KV-groups, two per SIP (study target) ──────────────


def test_step_3_eight_kv_groups_two_per_sip():
    """Two launches per SIP × 4 SIPs = 8 KV-groups total in one run_bench.

    The headline target: 64 cubes serving 8 KV-groups, two disjoint 2×4
    sub-meshes per SIP. ``cube_start=0`` puts the first KV-group on
    cubes 0..7 (rows 0..1); ``cube_start=8`` puts the second on cubes
    8..15 (rows 2..3). This is the use case ``DPPolicy.cube_start`` was
    added to enable.

    Expected to PASS — the SFR install draws intra-SIP edges only, so
    the 8 KV-groups can't see each other; ``cube_start`` ensures the
    two halves of each SIP land on disjoint cubes.
    """
    if not TOPOLOGY_4SIP.exists():
        pytest.skip(f"4-SIP topology missing: {TOPOLOGY_4SIP}")
    sip_groups = [
        (sip, f"sip{sip}_half{half}", half * N_CUBES_PER_KV_GROUP)
        for sip in range(4) for half in (0, 1)
    ]
    exc, result, engine = _run_multi_sip(sip_groups)
    if exc is not None:
        raise AssertionError(f"step_3 failed: {exc}") from exc
    assert result is not None and result.completion.ok, (
        f"step_3: completion not ok — {result.completion if result else None}"
    )