metal-flash-attention

Why we included this project

The FlashAttention algorithm, ported to Metal for Apple silicon, is worth studying for the engineering notes as much as the code. The implementation stays compact, a set of Swift sources JIT-compiled at runtime rather than a prebuilt binary, and the README explains the hard trade-offs in detail. Register pressure and the lack of native FP32 atomics on Apple hardware forced a redesigned backward pass, one that splits into separate dQ and dK/dV kernels and trades more GEMMs for full parallelization across both dimensions of the attention matrix. Performance is reported as gigainstructions per second against a roofline model instead of raw GFLOPS, with a consistent 4400 GI/s on an M1 Max at 83% ALU utilization, which gives you a concrete way to judge whether your own Metal shaders are leaving ALU throughput on the table. The register-spilling and block-tiling strategies transfer directly to anyone porting CUDA attention kernels to Apple GPUs.