mjun0812/flash-attention-prebuild-wheels

Provide with pre-build flash-attention 2 and 3 package wheels on Linux and Windows using GitHub Actions

What it solves

Building the flash-attention library from source is a resource-intensive process that takes a very long time. This project provides pre-built binary wheels for various combinations of Python, CUDA, and PyTorch versions, allowing users to install the library quickly without needing to compile it themselves.

How it works

The project uses a combination of GitHub-hosted runners, self-hosted runners, and AWS CodeBuild to automate the compilation of flash-attention (including Flash Attention 3) across multiple platforms. It generates a wide matrix of wheels for Linux (x86_64 and ARM64) and Windows, ensuring compatibility with various glibc versions and hardware requirements (such as Hopper SM90 GPUs for Flash Attention 3).

Who it’s for

Developers and researchers working with Large Language Models who need to use Flash Attention for faster, memory-efficient attention mechanisms but want to avoid the lengthy and complex build process.

Highlights

  • Extensive Coverage: Provides over 800 pre-built wheels across Linux and Windows.
  • Broad Compatibility: Supports various combinations of Python, CUDA, and PyTorch versions, including those not officially distributed.
  • Cross-Platform: Available for Linux x86_64, Linux ARM64, and Windows.
  • Simplified Installation: Allows for direct installation via pip install using a provided URL.

Related

  • Project
    wildminder/AI-windows-whl

    A curated collection of pre-compiled Python wheels for AI/ML libraries that are difficult to install on Windows, eliminating the need for source compilation.

  • Project
    meta-pytorch/attention-gym

    A collection of kernels, guides, and examples for FlexAttention and other novel attention variants, providing tools for efficient sparse and linear attention.

  • Project
    fla-org/flash-linear-attention

    A library of hardware-efficient, training-ready implementations for linear attention, state space models, and hybrid LLM architectures across multiple hardware platforms.

  • Project
    flagos-ai/FlagAttention

    A collection of memory-efficient attention operators implemented in Triton, allowing for easier customization of attention mechanisms compared to CUDA-based FlashAttention.

  • Project
    facebookresearch/xformers

    A toolbox of optimized, customizable building blocks for Transformer models that accelerates research by providing memory-efficient and high-performance CUDA kernels.