code/tests/test_attention_backends.py · viai957/align at main

align / code /tests /test_attention_backends.py

Add files using upload-large-folder tool

bcc8d35 verified 16 days ago

11.2 kB

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

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