flash-attn 3 - Windows ROCm pre-built wheels for AMD GPUs (aiter Triton backend)

Triton-backend only build of flash-attn 3 for Windows ROCm users who don't want to build from source.

Install

pip install https://github.com/0xDELUXA/flash-attention/releases/download/v3.0.0_win-rocm/flash_attn_3-3.0.0-py3-none-any.whl
pip install https://github.com/0xDELUXA/flash-attention/releases/download/v2.8.4_win-rocm/amd_aiter-0.0.0-py3-none-win_amd64.whl

Example

import torch
from flash_attn_interface import flash_attn_func

q = torch.randn(1, 1, 4, 16, device="cuda")
out = flash_attn_func(q, q, q)
print(out.shape)

Example output

UserWarning: flash_attn_3._C (which has ROCm/HIP kernels) not found, falling back to Triton implementation
[aiter] Windows: CK and HIP ops are not available. Triton ops only.
torch.Size([1, 1, 4, 16])