"""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)