Example how to use with Python app

[qdrddr]

Would you be so kind as to provide an example of using this implementation of MetalFlashAttention compiled as a library as a replacement for the original flash_attn with Python code running a model such as https://huggingface.co/microsoft/Phi-4-multimodal-instruct?

And thanks for your article!
Thank you!

[philipturner]

I am not likely to do that. To port it to Python, you first have to understand how the code works internally, then port it yourself. In addition, understand where performance can unexpectedly degrade, and have comprehensive unit tests to diagnose and correct performance drops. This is the case with pretty much any endeavor, whether software or real-world precision analog electronics. That is, endeavors where performance is at a premium, and high performance is extremely difficult and time-consuming to achieve.

[qdrddr]

Maybe some PoC with a bare minimum at least to give the community a starting point? @philipturner
Could you please consider providing a very simple Python wrapper maybe to use existing library billd with swift or something along those lines?

[philipturner]

I am not likely to provide assistance on this, unfortunately. It was a project that I was paid to do, but I am no longer at that occupation, so it no longer maintained. I designed the code, so that other people could build on it, in my absence.

[bghira]

see https://github.com/bghira/universal-metal-flash-attention for examples of pytorch custom op with MTLBuffer sharing

[qdrddr]

https://github.com/marcogva-hub/mlx-flashattention-steel/