| """Run the triton_fa2 equivalence suite and, ONLY on success, write the gate. |
| |
| Run (pin ONE free device): |
| CUDA_VISIBLE_DEVICES=<free gpu> uv run python scripts/gate_flash_attention.py |
| |
| This is the mechanism behind the never-silently-claimed rule. The gate artifact |
| is the ONLY thing that makes `triton_fa2` selectable |
| (`align.kernels.gate.resolve_triton_fa2_backend` raises without it), and this |
| script is the only writer. It records the kernel source hash, the torch/triton |
| versions, the exact GPU, and the tolerances and shapes the suite covered, so a |
| later edit to the kernel or a move to a different card invalidates it |
| automatically. |
| |
| Deliberately a separate entry point rather than a pytest hook: writing the gate |
| must require an explicit human action on a machine with a real GPU, and must |
| never be a side effect of someone running the test suite with `-k` narrowing |
| the selection to a subset. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import os |
| import subprocess |
| import sys |
| import xml.etree.ElementTree as ET |
|
|
| REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) |
| sys.path.insert(0, REPO) |
|
|
| from align.kernels import gate as gate_mod |
|
|
| |
| |
| TOLERANCES = { |
| "policy": "max |err| vs fp32 reference <= max(3.0 x torch-SDPA-bf16 error, 2e-2)", |
| "slack_over_sdpa": 3.0, |
| "absolute_floor": 2e-2, |
| } |
|
|
| SHAPES_COVERED = [ |
| "gqa (n_q_heads, n_kv_heads) in [(8,2), (32,8)] plus MHA-degenerate (4,4)", |
| "head_dim in [64, 128]", |
| "seq in [192, 256, 512, 1000] (block-divisible AND non-divisible)", |
| "causal and non-causal", |
| "key_padding_mask, right-padded, non-block-aligned", |
| "every autotune config individually pinned (fwd + both bwd kernels)", |
| "dtype bfloat16 only (fp16/fp32 rejected)", |
| ] |
|
|
|
|
| def _count_outcomes(junit_path: str) -> tuple[int, int, int]: |
| """(passed, skipped, failed+errored) from a pytest junit report. |
| |
| pytest exits 0 when EVERY test is skipped (no CUDA, more than one device |
| visible, heavy-job lock held), so the exit status alone cannot prove the |
| kernel ran. Raises if the report is missing or unparseable. |
| """ |
| root = ET.parse(junit_path).getroot() |
| passed = skipped = bad = 0 |
| for case in root.iter("testcase"): |
| if case.find("skipped") is not None: |
| skipped += 1 |
| elif case.find("failure") is not None or case.find("error") is not None: |
| bad += 1 |
| else: |
| passed += 1 |
| return passed, skipped, bad |
|
|
|
|
| def main() -> int: |
| cvd = os.environ.get("CUDA_VISIBLE_DEVICES", "") |
| devices = [d.strip() for d in cvd.split(",")] |
| if len(devices) != 1 or not devices[0]: |
| print( |
| "refusing to run: set CUDA_VISIBLE_DEVICES to exactly ONE free device " |
| f"(got {cvd!r}; shared box, and the suite skips unless exactly one is visible).", |
| file=sys.stderr, |
| ) |
| return 2 |
|
|
| junit = os.path.join(REPO, ".cache", "gate_flash_attention_junit.xml") |
| os.makedirs(os.path.dirname(junit), exist_ok=True) |
| if os.path.exists(junit): |
| os.remove(junit) |
| cmd = [ |
| sys.executable, |
| "-m", |
| "pytest", |
| os.path.join(REPO, "tests", "test_flash_attention.py"), |
| "-m", |
| "gpu", |
| "-q", |
| "-rs", |
| f"--junitxml={junit}", |
| ] |
| print("running the equivalence suite:", " ".join(cmd), flush=True) |
| result = subprocess.run(cmd, cwd=REPO) |
|
|
| if result.returncode != 0: |
| print( |
| "\nequivalence suite FAILED -- gate NOT written. triton_fa2 stays " |
| "unselectable; SDPA remains the backend.", |
| file=sys.stderr, |
| ) |
| return result.returncode |
|
|
| |
| try: |
| passed, skipped, bad = _count_outcomes(junit) |
| except (OSError, ET.ParseError) as e: |
| print(f"\ncannot read the junit report ({e!r}) -- gate NOT written.", file=sys.stderr) |
| return 1 |
| if passed == 0 or skipped != 0 or bad != 0: |
| print( |
| f"\nequivalence suite did not EXECUTE in full (passed={passed}, " |
| f"skipped={skipped}, failed={bad}) -- gate NOT written. Skips mean " |
| "no CUDA device, a device count != 1, or the heavy-job lock held " |
| "by a live job; none of those verifies the kernel.", |
| file=sys.stderr, |
| ) |
| return 1 |
|
|
| path = gate_mod.write_equivalence_gate(tolerances=TOLERANCES, shapes_covered=SHAPES_COVERED) |
| ok, reason = gate_mod.gate_status() |
| print(f"\nequivalence suite passed ({passed} tests); gate written to {path}") |
| print(f"gate_status -> passed={ok} ({reason})") |
| print(f"kernel sha256: {gate_mod.kernel_source_hash()}") |
| return 0 if ok else 1 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
