This repository provides the official implementation of FlashAttention from the following paper.
FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness
Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, Christopher Ré
Paper: https://arxiv.org/abs/2205.14135
IEEE Spectrum article about our submission to the MLPerf 2.0 benchmark using FlashAttention.
Phil Tillet (OpenAI) has an experimental implementation of FlashAttention in Triton: https://github.com/openai/triton/blob/master/python/tutorials/06-fused-attention.py
As Triton is a higher-level language than CUDA, it might be easier to understand and experiment with. The notations in the Triton implementation are also closer to what's used in our paper.
This branch contains the rewrite of FlashAttention forward pass to use Cutlass. This simplifies the code and supports more head dimensions. We support head dimensions that are multiples of 8 up to 128 (previously we supported head dimensions 16, 32, 64, 128).
To compile (requiring CUDA 11, NVCC, and an Turing or Ampere GPU):
python setup.py install
Interface: src/flash_attention.py
To run the benchmark against PyTorch standard attention:
PYTHONPATH=$PWD python benchmarks/benchmark_flash_attention.py
FlashAttention currently supports:
Our tentative roadmap:
We present expected speedup (combined forward + backward pass) and memory savings from using FlashAttention against PyTorch standard attention, depending on sequence length, on different GPUs (speedup depends on memory bandwidth - we see more speedup on slower GPU memory).
We currently have benchmarks for these GPUs:
We display FlashAttention speedup using these parameters (similar to BERT-base):
Our graphs show sequence lengths between 128 and 4096 (when standard attention runs out of memory on an A100), but FlashAttention can scale up to sequence length 64K.
We generally see 2-4X speedup at sequence lengths between 128 and 4K, and we see more speedup when using dropout and masking, since we fuse the kernels. At sequence lengths that are popular with language models like 512 and 1K, we see speedups up to 4X when using dropout and masking.
We show memory savings in this graph (note that memory footprint is the same no matter if you use dropout or masking). Memory savings are proportional to sequence length -- since standard attention has memory quadratic in sequence length, whereas FlashAttention has memory linear in sequence length. We see 10X memory savings at sequence length 2K, and 20X at 4K. As a result, FlashAttention can scale to much longer sequence lengths.
We show speedup with head dimension 128. Here we show batch size 16 with 12 heads. Speedup is less than with the smaller head sizes, since we have to make the block size smaller in the tiling. But speedup is still significant, especially with a causal mask.
For the RTX 3090, we use batch size 12 with 12 attention heads. Memory savings are the same as on an A100, so we'll only show speedup here.
We see slightly higher speedups (between 2.5-4.5x) on the GTX 3090, since memory bandwidth on the GDDR6X is lower than A100 HBM (~900 GB/s vs. ~1.5 TB/s).
We again use batch size 12 with 12 attention heads.
T4 SRAM is smaller than the newer GPUs (64 KB), so we see less speedup (we need to make the block sizes smaller, so we end up doing more R/W). This matches the IO complexity analysis from section 3.2 of our paper.
T4 GPUs are commonly used for inference, so we also measure speedup on the forward pass only (note that these are not directly comparable to the graphs above):
We see speedups between 2.5x-4.5x on the forward pass.
We test that FlashAttention produces the same output and gradient as a reference implementation, up to some numerical tolerance. In particular, we check that the maximum numerical error of FlashAttention is at most twice the numerical error of a baseline implementation in Pytorch (for different head dimensions, input dtype, sequence length, causal / non-causal).
To run the tests:
pytest -q -s tests/test_flash_attn.py
This alpha release of FlashAttention contains code written for a research project to validate ideas on speeding up attention. We have tested it on several models (BERT, GPT2, ViT). However, there might still be bugs in the implementation that we hope to iron out in the next few months.
If you encounter any of these bugs, please open a respective GitHub Issue!
Our implementation uses Apex's FMHA code as a starting point.
We thank Young-Jun Ko for the in-depth explanation of his FMHA implementation and for his thoughtful answers to our questions about CUDA.
If you use this codebase, or otherwise found our work valuable, please cite:
@article{dao2022flashattention,
title={FlashAttention: 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},
journal={arXiv preprint arXiv:2205.14135},
year={2022}
}
此处可能存在不合适展示的内容,页面不予展示。您可通过相关编辑功能自查并修改。
如您确认内容无涉及 不当用语 / 纯广告导流 / 暴力 / 低俗色情 / 侵权 / 盗版 / 虚假 / 无价值内容或违法国家有关法律法规的内容,可点击提交进行申诉,我们将尽快为您处理。