Flash Attention Benchmarking on B200 #10
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 */2 * * *" # Run every 2 hours | |
| 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 | |
| pwd | |
| git log -1 | |
| python setup.py install | |
| export PYTHONPATH=$(pwd) | |
| echo '<h1>B200 1000W</h1>' >> $GITHUB_STEP_SUMMARY | |
| nvidia-smi | |
| python benchmarks/benchmark_attn.py >> $GITHUB_STEP_SUMMARY | |
| popd |