flash-attention pre-build wheels
This repository provides wheels for the pre-built flash-attention.
Since building flash-attention takes a very long time and is resource-intensive, I also build and provide combinations of CUDA and PyTorch that are not officially distributed.
The building Github Actions Workflow can be found here.
The built packages are available on the release page.
This repository uses a self-hosted runner and AWS CodeBuild for building the wheels. If you find this project helpful, please consider sponsoring to help maintain the infrastructure!
Special thanks to @KiralyCraft for providing the computing resources used to build wheels. Thank you!!
Install
- Select the versions for Python, CUDA, PyTorch, and flash_attn.
flash_attn-[flash_attn Version]+cu[CUDA Version]torch[PyTorch Version]-cp[Python Version]-cp[Python Version]-linux_x86_64.whl
# Example: Python 3.11, CUDA 12.4, PyTorch 2.5, and flash_attn 2.6.3
flash_attn-2.6.3+cu124torch2.5-cp312-cp312-linux_x86_64.whl
-
Find the corresponding version of a wheel from the Packages page and releases page.
-
Direct Install or Download and Local Install
# Direct Install
pip install https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.0.0/flash_attn-2.6.3+cu124torch2.5-cp312-cp312-linux_x86_64.whl
# Download and Local Install
wget https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.0.0/flash_attn-2.6.3+cu124torch2.5-cp312-cp312-linux_x86_64.whl
pip install ./flash_attn-2.6.3+cu124torch2.5-cp312-cp312-linux_x86_64.whl
Packages
Note
Since v0.5.0, wheels are built with a local version label indicating the CUDA and PyTorch versions.
Example:pip list->flash_attn==2.8.3 -> flash_attn==2.8.3+cu130torch2.9
See ./docs/packages.md for the full list of available packages.
History
History of this repository is available here.
Citation
If you use this repository in your research and find it helpful, please cite the following paper!
@misc{flash-attention-prebuild-wheels,
author = {Morioka, Junya},
year = {2025},
title = {mjun0812/flash-attention-prebuild-wheels},
url = {https://github.com/mjun0812/flash-attention-prebuild-wheels},
howpublished = {https://github.com/mjun0812/flash-attention-prebuild-wheels},
}
Acknowledgments
- @okaris : Sponsored me!
- @xhiroga : Sponsored me!
- cjustus613 : Buy me a coffee!
- @KiralyCraft : Provided with computing resource!
Star History and Download Statistics
|
|
|
Original Repository
@inproceedings{dao2022flashattention,
title={Flash{A}ttention: Fast and Memory-Efficient Exact Attention with {IO}-Awareness},
author={Dao, Tri and Fu, Daniel Y. and Ermon, Stefano and Rudra, Atri and R{\'e}, Christopher},
booktitle={Advances in Neural Information Processing Systems (NeurIPS)},
year={2022}
}
@inproceedings{dao2023flashattention2,
title={Flash{A}ttention-2: Faster Attention with Better Parallelism and Work Partitioning},
author={Dao, Tri},
booktitle={International Conference on Learning Representations (ICLR)},
year={2024}
}
Self build
If you cannot find the version you are looking for, you can fork this repository and create a wheel on GitHub Actions.
- Fork this repository
- Edit Python script
create_matrix.pyto set the version you want to build. - Add tag
v*.*.*to trigger the build workflow.git tag v*.*.* && git push --tags
Please note that depending on the combination of versions, it may not be possible to build.
Self-Hosted Runner Build
In some version combinations, you cannot build wheels on GitHub-hosted runners due to job time limitations. To build the wheels for these versions, you can use self-hosted runners.
git clone https://github.com/mjun0812/flash-attention-prebuild-wheels.git
cd self-hosted-runner
cp env.template env
Edit env file to set the environment variables.
# Registry Token for GitHub Personal Access Token
PERSONAL_ACCESS_TOKEN=[Github Personal Access Token]
# or Registry Token for GitHub Actions Runner
REGISTRY_TOKEN=[Runner Registry Token]
# Optional
RUNNER_LABELS=Linux,self-hosted
Edit compose.yml file if you use repository folked from this repository.
services:
runner:
privileged: true
restart: always
env_file:
- .env
environment:
REPOSITORY_URL: https://github.com/[OWNER]/[REPOSITORY]
RUNNER_NAME: self-hosted-runner
RUNNER_GROUP: default
TARGET_ARCH: x64
build:
context: .
dockerfile: Dockerfile
args:
GH_RUNNER_VERSION: 2.329.0
TARGET_ARCH: x64
Then, build and run the docker container.
# Build and run
docker compose build
docker compose up -d
Getting One-Time Registry Token for GitHub Actions Runner
gh api \
-X POST \
/repos/[OWNER]/[REPOSITORY]/actions/runners/registration-token
