How can I use this Flash Attention Triton implementation with PyTorch+ROCM?

[farshadghodsian]

Great work on coding an implementation of Triton Flash Attention that works with Radeon GPUs (gfx1100) like the W7900. My question is would I be able to install this version of flash attention as a pip package for use in PyTorch on Linux?

I did not have any luck trying to adapt your code and update the Triton Flash Attention files in the ROCm Flash Attention repo and build from source to use your implementation: https://github.com/farshadghodsian/flash-attention/tree/flash_attention_for_rocm/flash_attn.

Any assistance or direction on this would be greatly appreciated. My goal is to get model training and inference using memory efficient attention working on my Radeon Pro W7900 using PyTorch and not just for Stable Diffusion. Thank you