Blame for code/tests/test_attention_backends.py

"""Backend probe/select/benchmark guards, both directions.

Run: uv run pytest tests/test_attention_backends.py -v -m "not gpu"

Every test here runs the probe on device="cpu" on purpose: the CPU suite must
not initialise a CUDA context on a shared box, and the selection LOGIC is
device-independent. The device-specific numbers are the gpu-marked benchmark.
"""

from __future__ import annotations

import math
import os
import sys

import pytest

from align.model import attention as attn
from align.model.attention import (
    TRITON_FA2_GATE_MODULE,
    BackendReport,
    benchmark_backends,
    probe_backends,
    select_attn_implementation,
)
from align.model.gates import load_sdpa_baseline


@pytest.fixture
def no_triton_kernel():
    """Make section-06's gate module unimportable for the duration of a test.

    The 'kernel absent' branch is a permanent code path -- it is what any fresh
    checkout, and any box where the Triton build fails, actually hits -- so it
    is tested by BLOCKING the import rather than by relying on section-06 not
    having landed yet. A test that silently stops testing anything the day a
    sibling lane ships is worse than no test.
    """

    class _Block:
        def find_spec(self, name, path=None, target=None):
            if name == TRITON_FA2_GATE_MODULE:
                raise ImportError(f"blocked by test fixture: {name}")
            return None

    blocker = _Block()
    saved = sys.modules.pop(TRITON_FA2_GATE_MODULE, None)
    sys.meta_path.insert(0, blocker)
    try:
        yield
    finally:
        sys.meta_path.remove(blocker)
        if saved is not None:
            sys.modules[TRITON_FA2_GATE_MODULE] = saved


@pytest.fixture
def gate_dir(tmp_workdir):
    """A scratch equivalence-gate directory, so these tests never read -- or
    write -- the real artifact under .cache/kernel_gates."""
    return os.path.join(tmp_workdir, "kernel_gates")


def _write_passing_gate(gate_dir):
    """Write a REAL gate artifact for the current kernel source and
    environment, using section-06's own writer. Stubbing a boolean here would
    test my mock; this tests the contract the two sections actually share."""
    from align.kernels import gate as gate_mod

    gate_mod.write_equivalence_gate(
        tolerances={"out": 1e-3},
        shapes_covered=["backend-selection test"],
        gate_dir=gate_dir,
    )


def _reports(device="cpu", gate_dir=None):
    return {r.name: r for r in probe_backends(device=device, gate_dir=gate_dir)}


def test_probe_reports_every_backend_with_a_reason(gate_dir):
    """A probe that omits a backend is indistinguishable from one that found it
    unavailable. Every declared backend gets a row, and every row carries a
    non-empty reason -- the reason string is what ends up in the launch banner
    when a run cannot use the accelerator someone expected."""
    reports = _reports(gate_dir=gate_dir)
    assert set(reports) == set(attn.ALL_BACKENDS)
    for name, r in reports.items():
        assert isinstance(r, BackendReport)
        assert r.reason.strip(), name


def test_select_triton_fa2_raises_before_equivalence_gate(gate_dir):
    """select_attn_implementation("triton_fa2") RAISES while the kernel's
    SDPA-equivalence gate (section-06) has not passed. It must never silently
    fall back -- a run must not claim an accelerator it never used. (Both
    directions: with a valid gate artifact it returns the backend.)"""
    # direction 1: no gate artifact -> refuse
    with pytest.raises(RuntimeError, match="equivalence gate"):
        select_attn_implementation("triton_fa2", device="cpu", gate_dir=gate_dir)

    # direction 2: a gate artifact for this kernel source and environment
    _write_passing_gate(gate_dir)
    assert select_attn_implementation("triton_fa2", device="cpu", gate_dir=gate_dir) == "triton_fa2"


def test_select_triton_fa2_raises_when_the_kernel_is_absent(no_triton_kernel):
    """The fresh-checkout path: section-06's module is not importable at all.
    That is a different branch from 'gate not passed' and it must also refuse
    rather than fall back."""
    with pytest.raises(RuntimeError, match="not yet implemented"):
        select_attn_implementation("triton_fa2", device="cpu")


def test_select_triton_fa2_refuses_a_gate_that_no_longer_matches_the_kernel(gate_dir):
    """THE STALE-GATE INCIDENT, at the selector: an artifact written before a
    kernel edit would bless code that was never run against SDPA. The gate
    pins the kernel source hash; the selector must surface that refusal rather
    than treat 'an artifact exists' as availability."""
    from align.kernels import gate as gate_mod

    _write_passing_gate(gate_dir)
    assert select_attn_implementation("triton_fa2", device="cpu", gate_dir=gate_dir) == "triton_fa2"

    original = gate_mod.kernel_source_hash
    gate_mod.kernel_source_hash = lambda: "0" * 64
    try:
        with pytest.raises(RuntimeError, match="kernel source changed"):
            select_attn_implementation("triton_fa2", device="cpu", gate_dir=gate_dir)
    finally:
        gate_mod.kernel_source_hash = original


def test_triton_fa2_is_never_auto_selected(gate_dir):
    """Even with the gate passed, auto-selection stays on sdpa. A custom
    kernel is opted INTO explicitly so the banner records a deliberate choice;
    a kernel that switches itself on turns every later benchmark into a
    comparison against an unknown baseline."""
    _write_passing_gate(gate_dir)
    assert select_attn_implementation(None, device="cpu", gate_dir=gate_dir) == "sdpa"


def test_select_unavailable_backend_raises_with_reason(gate_dir):
    """Explicitly requesting any unavailable backend raises RuntimeError
    carrying the probe's reason string; auto-selection (requested=None) picks
    the best available and logs which."""
    reports = _reports(gate_dir=gate_dir)
    unavailable = [n for n, r in reports.items() if not r.available]
    assert unavailable, "probe found nothing unavailable -- it is not probing"
    for name in unavailable:
        with pytest.raises(RuntimeError) as exc:
            select_attn_implementation(name, device="cpu", gate_dir=gate_dir)
        # the probe's own reason must survive into the error the user sees
        assert reports[name].reason[:40] in str(exc.value), name

    assert select_attn_implementation(None, device="cpu", gate_dir=gate_dir) == "sdpa"


def test_flashmla_refused_even_if_it_were_available():
    """FlashMLA is refused for structural reasons, not availability ones: its
    kernel requires head_dim_v=512 MLA tensors and Mistral is GQA head_dim
    128, so no GPU -- Hopper included -- can serve this model with it. The
    refusal must not be reachable-around by a future SM90 box."""
    with pytest.raises(RuntimeError, match="head_dim"):
        select_attn_implementation("flashmla", device="cpu")


def test_select_rejects_an_unknown_backend_name():
    """A typo'd backend name must not fall through to the default. The
    predecessor's whole failure class is a plausible substitute quietly
    standing in for what was asked for."""
    with pytest.raises(RuntimeError, match="unknown backend"):
        select_attn_implementation("sdpa_flash", device="cpu")


def test_benchmark_records_nan_not_zero_for_unusable_backend():
    """A backend that throws at these shapes records float('nan'), never 0.0
    and never a swallowed pass (the EFFICIENT-on-GQA incident: NaN is a real
    finding, recorded not hidden). Undefined stats are NaN, never 0.

    Run on CPU at tiny shapes, where EFFICIENT_ATTENTION and CUDNN_ATTENTION
    genuinely have no kernel -- the same shape of unavailability the pinned
    A10G baseline recorded for EFFICIENT on GQA inputs.
    """
    results = benchmark_backends(
        batch=1,
        seq=16,
        n_heads=4,
        n_kv_heads=2,
        head_dim=16,
        iters=2,
        device="cpu",
    )
    assert set(results) >= {"FLASH_ATTENTION", "EFFICIENT_ATTENTION", "MATH"}
    unusable = [k for k, v in results.items() if math.isnan(v)]
    assert unusable, "no backend was unusable at CPU tiny shapes -- expected at least one"
    for name, ms in results.items():
        assert ms != 0.0, f"{name} recorded 0.0; an unusable backend is NaN, never 0"
        assert math.isnan(ms) or ms > 0.0, f"{name}={ms}"


def test_benchmark_does_not_swallow_the_reason(monkeypatch):
    """The exception a backend raises is a finding. It is logged and attached
    to the result, not discarded -- 'NaN' with no reason cannot be triaged
    six weeks later when a torch upgrade changes which kernels exist."""
    results, reasons = attn.benchmark_backends_with_reasons(
        batch=1,
        seq=16,
        n_heads=4,
        n_kv_heads=2,
        head_dim=16,
        iters=2,
        device="cpu",
    )
    for name, ms in results.items():
        if math.isnan(ms):
            assert reasons[name] and reasons[name] != "ok", name
        else:
            assert reasons[name] == "ok", name


def test_pinned_sdpa_baseline_is_the_bar_and_records_nan():
    """The measured A10G baseline is pinned in-repo so section-06's 'beat the
    bar' comparison and a torch/driver regression of the SDPA path itself are
    both detectable. EFFICIENT is NaN in the pinned artifact for the same
    reason it is NaN at runtime -- that is the recorded finding, not a hole."""
    base = load_sdpa_baseline()
    assert set(base) == {"512", "1024", "2048"}
    for seq, row in base.items():
        assert row["FLASH_ATTENTION"] > 0.0
        assert math.isnan(row["EFFICIENT_ATTENTION"]), (
            "EFFICIENT must be NaN, not 0 and not absent, in the pinned baseline"
        )
        # flash must beat math by a wide margin or the pin is of the wrong run
        assert row["MATH"] > 10 * row["FLASH_ATTENTION"], seq


def test_banner_states_which_backend_actually_ran():
    """Every training/inference job's banner reports the backend that was
    resolved, so 'we ran the custom kernel' is a recorded fact rather than an
    intention. attention_backend is already a mandatory banner key; this test
    pins that the selector's output is what fills it."""
    from align.registry import MANDATORY_BANNER_KEYS

    assert "attention_backend" in MANDATORY_BANNER_KEYS
    resolved = select_attn_implementation(None, device="cpu")
    assert attn.banner_value(resolved, device="cpu").startswith(resolved)


@pytest.mark.gpu
def test_benchmark_reproduces_the_pinned_a10g_baseline():
    """On a real A10G, re-running the benchmark at the pinned shapes must land
    within 2x of the recorded numbers and must still record EFFICIENT as NaN.
    A drastic regression of the SDPA path (torch upgrade, driver change) is
    exactly what this catches."""
    import torch

    if not torch.cuda.is_available():
        pytest.skip("no CUDA device visible")
    base = load_sdpa_baseline()
    for seq in (512, 1024, 2048):
        got = benchmark_backends(batch=4, seq=seq, device="cuda")
        ref = base[str(seq)]
        assert math.isnan(got["EFFICIENT_ATTENTION"])
        assert got["FLASH_ATTENTION"] < 2.0 * ref["FLASH_ATTENTION"], (seq, got, ref)