[Feature]: FlashAttention on macOS

[qdrddr]

I'd like to benefit from flash attention on Mac platform

https://github.com/philipturner/metal-flash-attention
https://github.com/bghira/universal-metal-flash-attention

[marcogva-hub]

I wanted an MLX Flash Attention version, and with the help of Claude, ended up creating one...

If anyone wants to try, feel free to have a look at https://github.com/marcogva-hub/mlx-flashattention-steel/

[WiseFarAI]

I wanted an MLX Flash Attention version, and with the help of Claude, ended up creating one...

If anyone wants to try, feel free to have a look at https://github.com/marcogva-hub/mlx-flashattention-steel/

That seems really interesting. Wondering how these numbers will translate if something similar is implemented natively. Looks like there is still a lot of software optimization headroom left.

[marcogva-hub]

That seems really interesting. Wondering how these numbers will translate if something similar is implemented natively. Looks like there is still a lot of software optimization headroom left.

There is, particularly with the M5 Metal Accelerators (Metal 4 TensorOps), but I don't have the hardware to test it right now, and I won't be replacing my M1 Max for a few months.

For the M1->M4 generations, one thing that was great was simdgroup_async_copy, but Apple has blocked __asm in JIT compilation... I'm testing a solution, let's hope it works: I should have my answer in a few hours.

[WiseFarAI]

That seems really interesting. Wondering how these numbers will translate if something similar is implemented natively. Looks like there is still a lot of software optimization headroom left.

There is, particularly with the M5 Metal Accelerators (Metal 4 TensorOps), but I don't have the hardware to test it right now, and I won't be replacing my M1 Max for a few months.

For the M1->M4 generations, one thing that was great was simdgroup_async_copy, but Apple has blocked __asm in JIT compilation... I'm testing a solution, let's hope it works: I should have my answer in a few hours.

Also still rocking an M1 Max. Please keep us updated on your progress. Very interested!

[marcogva-hub]

Also still rocking an M1 Max. Please keep us updated on your progress. Very interested!

I ended up precompiling my async_copy metallib kernel through GitHub Actions with macOS 15/Xcode 16, but Apple not only removed entirely async_copy from their librairies on macOS 26 and Xcode 26, when a compiled metallib is executed, macOS 26 now automatically and silently converts async opcodes to synchronous loads, defeating the purpose... 🤬

I'll keep looking for optimization ways, but this one was the most promising.