Flash Attention Benchmarking on B200 #12
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| name: Flash Attention Benchmark | |
| # To remotely trigger a FA Benchmarking run, use the following: | |
| # curl -XPOST -H "Accept: application/vnd.github.v3+json" -H "Content-Type: application/json" https://api.github.com/repos/pytorch/pytorch-integration-testing/dispatches --data '{"event_type": "benchmark_flash_attention"}' | |
| on: | |
| schedule: | |
| - cron: "0 6 * * *" # Run every day at 6AM | |
| push: | |
| paths: | |
| - .github/workflows/flash_attention.yml | |
| repository_dispatch: | |
| types: benchmark_flash_attention | |
| workflow_dispatch: | |
| jobs: | |
| benchmark-flash-attn: | |
| name: Flash Attention CuTe DSL Benchmark | |
| runs-on: B200 | |
| container: | |
| # https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pytorch/ | |
| image: nvcr.io/nvidia/pytorch:25.06-py3 | |
| options: --gpus all --ipc=host --ulimit memlock=-1 --ulimit stack=67108864 | |
| steps: | |
| - uses: actions/checkout@v4 | |
| with: | |
| repository: 'Dao-AILab/flash-attention' | |
| path: 'fa4' | |
| - name: Install CuTe DSL | |
| run: | | |
| set -x | |
| echo "Installing nvidia-cutlass-dsl" | |
| pip install nvidia-cutlass-dsl==4.1.0.dev0 | |
| - name: Buid and Run FlashAttention CuTe DSL | |
| run: | | |
| set -x | |
| pushd fa4 | |
| python setup.py install | |
| echo '<h1>B200 1000W</h1>' >> $GITHUB_STEP_SUMMARY | |
| nvidia-smi | |
| export PYTHONPATH=$(pwd) | |
| python benchmarks/benchmark_attn.py >> $GITHUB_STEP_SUMMARY | |
| popd |