File size: 4,954 Bytes
bcc8d35 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 | """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
#: Kept in step with tests/test_flash_attention.py. Recorded in the artifact so
#: a reader knows what "verified" actually covered.
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) # a stale report must never vouch for this run
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
# Exit 0 is necessary, not sufficient: an all-skipped run also exits 0.
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())
|