https://github.com/gindachen/flexflashattention3

FlexAttention w/ FlashAttention3 Support
https://github.com/gindachen/flexflashattention3

Last synced: 17 days ago
JSON representation

FlexAttention w/ FlashAttention3 Support

• Host: GitHub
• URL: https://github.com/gindachen/flexflashattention3
• Owner: GindaChen
• Created: 2024-09-24T23:08:58.000Z (4 months ago)
• Default Branch: main
• Last Pushed: 2024-10-05T01:59:56.000Z (3 months ago)
• Last Synced: 2024-10-19T05:15:55.858Z (3 months ago)
• Language: Python
• Size: 26.4 KB
• Stars: 24
• Watchers: 2
• Forks: 1
• Open Issues: 0
• Metadata Files:
  • Readme: README.md

Awesome Lists containing this project

README

        
⚠️WIP: We have not put the flexattention/FA3 logic here yet! Coming soon...

Support FlexAttention in FlashAttention3

## Build environment

These instructions currently only tested in CUDA 12.4. We eventually need CUDA 12.3 for the build, and this is still under dev.

Install mamba

```bash

# https://mamba.readthedocs.io/en/latest/installation/micromamba-installation.html

"${SHELL}" <(curl -L micro.mamba.pm/install.sh)

```

Build environment to benchmark FlashAttention3 with nightly torch

```bash

mkdir -p envs

mamba create -p ~/envs/flashattn_3 python=3.10 -y

ln -s ~/envs/flashattn_3 ./envs/flashattn_3

mamba activate ./envs/flashattn_3

mamba install pytorch torchvision torchaudio pytorch-cuda=12.4 -c pytorch-nightly -c nvidia -y

mamba install ipython nvitop ninja cmake -y

pip install packaging

# Install FlashAttention3

MAX_JOBS=80 python setup.py install

cd hopper

MAX_JOBS=80 python setup.py install

```

Build environment to benchmark FlexAttention in nightly build torch

```bash

mkdir -p envs

mamba create -p ~/envs/flexattn-nightly-built python=3.10 -y

mamba activate ~/envs/flexattn-nightly-built

mamba install pytorch torchvision torchaudio pytorch-cuda=12.4 -c pytorch-nightly -c nvidia -y

mamba install ipython nvitop -y

```

Build environment for Flex attention gym

```bash

mamba activate ~/envs/flexattn-nightly-built

git clone https://github.com/pytorch-labs/attention-gym.git

cd attention-gym

pip install .

```