README.md 9.92 KB
Newer Older
1
# FlashAttention
Tri Dao's avatar
Tri Dao committed
2
3
4
This repository provides the official implementation of FlashAttention and
FlashAttention-2 from the
following papers.
5

Tri Dao's avatar
Tri Dao committed
6
**FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness**  
Tri Dao's avatar
Tri Dao committed
7
Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, Christopher Ré  
8
9
Paper: https://arxiv.org/abs/2205.14135  
IEEE Spectrum [article](https://spectrum.ieee.org/mlperf-rankings-2022) about our submission to the MLPerf 2.0 benchmark using FlashAttention.
10
![FlashAttention](assets/flashattn_banner.jpg)
11

Tri Dao's avatar
Tri Dao committed
12
13
**FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning**  
Tri Dao
Tri Dao's avatar
Tri Dao committed
14

Tri Dao's avatar
Tri Dao committed
15
Paper: https://tridao.me/publications/flash2/flash2.pdf
Tri Dao's avatar
Tri Dao committed
16

Tri Dao's avatar
Tri Dao committed
17
![FlashAttention-2](assets/flashattention_logo.png)
Tri Dao's avatar
Tri Dao committed
18
19


Tri Dao's avatar
Tri Dao committed
20
## Usage
Tri Dao's avatar
Tri Dao committed
21

Tri Dao's avatar
Tri Dao committed
22
23
24
We've been very happy to see FlashAttention being widely adopted in such a short
time after its release. This [page](https://github.com/Dao-AILab/flash-attention/blob/main/usage.md)
contains a partial list of places where FlashAttention is being used.
Tri Dao's avatar
Tri Dao committed
25

Tri Dao's avatar
Tri Dao committed
26
27
FlashAttention and FlashAttention-2 are free to use and modify (see LICENSE).
Please cite and credit FlashAttention if you use it.
Tri Dao's avatar
Tri Dao committed
28

Tri Dao's avatar
Tri Dao committed
29
## Installation and features
Tri Dao's avatar
Tri Dao committed
30

31
Requirements:
32
- CUDA 11.6 and above.
33
- PyTorch 1.12 and above.
34
35
- Linux. Windows is not supported for now. If you have ideas on how to modify
  the code to support Windows, please reach out via Github issue.
36

Tri Dao's avatar
Tri Dao committed
37
38
39
40
We recommend the
[Pytorch](https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pytorch)
container from Nvidia, which has all the required tools to install FlashAttention.

41
To install:
Tri Dao's avatar
Tri Dao committed
42
43
44
45
46
1. Make sure that PyTorch is installed.
2. Make sure that `packaging` is installed (`pip install packaging`)
3. Make sure that `ninja` is installed and that it works correctly (e.g. `ninja
--version` then `echo $?` should return exit code 0). If not (sometimes `ninja
--version` then `echo $?` returns a nonzero exit code), uninstall then reinstall
Tri Dao's avatar
Tri Dao committed
47
`ninja` (`pip uninstall -y ninja && pip install ninja`). Without `ninja`,
Tri Dao's avatar
Tri Dao committed
48
49
50
compiling can take a very long time (2h) since it does not use multiple CPU
cores. With `ninja` compiling takes 3-5 minutes on a 64-core machine.
4. Then:
51
```sh
Tri Dao's avatar
Tri Dao committed
52
pip install flash-attn --no-build-isolation
53
54
```
Alternatively you can compile from source:
Ian Timmis's avatar
Ian Timmis committed
55
```sh
Tri Dao's avatar
Tri Dao committed
56
python setup.py install
Tri Dao's avatar
Tri Dao committed
57
58
```

59
60
61
62
If your machine has less than 96GB of RAM and lots of CPU cores, `ninja` might
run too many parallel compilation jobs that could exhaust the amount of RAM. To
limit the number of parallel compilation jobs, you can set the environment
variable `MAX_JOBS`:
Ian Timmis's avatar
Ian Timmis committed
63
```sh
64
65
66
MAX_JOBS=4 pip install flash-attn --no-build-isolation
```

Tri Dao's avatar
Tri Dao committed
67
Interface: `src/flash_attention_interface.py`
Tri Dao's avatar
Tri Dao committed
68

Tri Dao's avatar
Tri Dao committed
69
70
71
72
73
74
FlashAttention-2 currently supports:
1. Ampere, Ada, or Hopper GPUs (e.g., A100, RTX 3090, RTX 4090, H100). Support for Turing
   GPUs (T4, RTX 2080) is coming soon, please use FlashAttention 1.x for Turing
   GPUs for now.
2. Datatype fp16 and bf16 (bf16 requires Ampere, Ada, or Hopper GPUs).
3. All head dimensions up to 256. Head dim > 192 backward requires A100/A800 or H100/H800.
Tri Dao's avatar
Tri Dao committed
75

76
77
78

## How to use FlashAttention

Tri Dao's avatar
Tri Dao committed
79
80
The main functions implement scaled dot product attention (softmax(Q @ K^T *
softmax_scale) @ V):
Ian Timmis's avatar
Ian Timmis committed
81
```python
Tri Dao's avatar
Tri Dao committed
82
from flash_attn import flash_attn_qkvpacked_func, flash_attn_func
83
84
```

Ian Timmis's avatar
Ian Timmis committed
85
```python
Tri Dao's avatar
Tri Dao committed
86
87
88
89
90
91
92
93
94
95
96
97
98
flash_attn_qkvpacked_func(qkv, dropout_p=0.0, softmax_scale=None, causal=False):
"""dropout_p should be set to 0.0 during evaluation
If Q, K, V are already stacked into 1 tensor, this function will be faster than
calling flash_attn_func on Q, K, V since the backward pass avoids explicit concatenation
of the gradients of Q, K, V.
Arguments:
    qkv: (batch_size, seqlen, 3, nheads, headdim)
    dropout_p: float. Dropout probability.
    softmax_scale: float. The scaling of QK^T before applying softmax.
        Default to 1 / sqrt(headdim).
    causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).
Return:
    out: (batch_size, seqlen, nheads, headdim).
Ian Timmis's avatar
Ian Timmis committed
99
"""
Tri Dao's avatar
Tri Dao committed
100
```
101

Ian Timmis's avatar
Ian Timmis committed
102
```python
Tri Dao's avatar
Tri Dao committed
103
104
105
flash_attn_func(q, k, v, dropout_p=0.0, softmax_scale=None, causal=False):
"""dropout_p should be set to 0.0 during evaluation
Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads
Tri Dao's avatar
Tri Dao committed
106
than Q. Note that the number of heads in Q must be divisible by the number of heads in KV.
Tri Dao's avatar
Tri Dao committed
107
108
109
110
111
112
113
114
115
116
117
118
119
For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head
0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V.

Arguments:
    q: (batch_size, seqlen, nheads, headdim)
    k: (batch_size, seqlen, nheads_k, headdim)
    v: (batch_size, seqlen, nheads_k, headdim)
    dropout_p: float. Dropout probability.
    softmax_scale: float. The scaling of QK^T before applying softmax.
        Default to 1 / sqrt(headdim).
    causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).
Return:
    out: (batch_size, seqlen, nheads, headdim).
Ian Timmis's avatar
Ian Timmis committed
120
"""
121
122
```

Tri Dao's avatar
Tri Dao committed
123
124
To see how these functions are used in a multi-head attention layer (which
includes QKV projection, output projection), see the MHA [implementation](https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/modules/mha.py).
125

Tri Dao's avatar
Tri Dao committed
126
## Upgrading from FlashAttention (1.x) to FlashAttention-2
127

Tri Dao's avatar
Tri Dao committed
128
129
130
131
These functions have been renamed:
- `flash_attn_unpadded_func` -> `flash_attn_varlen_func`
- `flash_attn_unpadded_qkvpacked_func` -> `flash_attn_varlen_qkvpacked_func`
- `flash_attn_unpadded_kvpacked_func` -> `flash_attn_varlen_kvpacked_func`
132

Tri Dao's avatar
Tri Dao committed
133
134
If the inputs have the same sequence lengths in the same batch, it is simpler
and faster to use these functions:
Ian Timmis's avatar
Ian Timmis committed
135
```python
Tri Dao's avatar
Tri Dao committed
136
flash_attn_qkvpacked_func(qkv, dropout_p=0.0, softmax_scale=None, causal=False)
137
```
Ian Timmis's avatar
Ian Timmis committed
138
```python
Tri Dao's avatar
Tri Dao committed
139
flash_attn_func(q, k, v, dropout_p=0.0, softmax_scale=None, causal=False)
140
```
141
142
143
144
145
## Changes in v2.1 (compared to v2.0)

If seqlen_q != seqlen_k and causal=True, the causal mask is aligned to the
bottom right corner of the attention matrix, instead of the top-left corner.

Tri Dao's avatar
Tri Dao committed
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 =
masked out) is:  
v2.0:  
    1 0 0 0 0  
    1 1 0 0 0  
v2.1:  
    1 1 1 1 0  
    1 1 1 1 1  

If seqlen_q = 5 and seqlen_k = 2, the causal mask is:  
v2.0:  
    1 0  
    1 1  
    1 1  
    1 1  
    1 1  
v2.1:  
    0 0  
    0 0  
    0 0  
    1 0  
    1 1  
168
If the row of the mask is all zero, the output will be zero.
169

Tri Dao's avatar
Tri Dao committed
170
## Performance
Dan Fu's avatar
Dan Fu committed
171

Dan Fu's avatar
Dan Fu committed
172
173
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).

Dan Fu's avatar
T4  
Dan Fu committed
174
175
We currently have benchmarks for these GPUs:
* [A100](#a100)
Tri Dao's avatar
Tri Dao committed
176
177
178
* [H100](#h100)
<!-- * [RTX 3090](#rtx-3090) -->
<!-- * [T4](#t4) -->
Dan Fu's avatar
T4  
Dan Fu committed
179

Dan Fu's avatar
Dan Fu committed
180
181
### A100

Tri Dao's avatar
Tri Dao committed
182
183
184
185
We display FlashAttention speedup using these parameters:
* Head dimension 64 or 128, hidden dimension 2048 (i.e. either 32 or 16 heads).
* Sequence length 512, 1k, 2k, 4k, 8k, 16k.
* Batch size set to 16k / seqlen.
Dan Fu's avatar
Dan Fu committed
186

Dan Fu's avatar
Dan Fu committed
187
#### Speedup
Dan Fu's avatar
Dan Fu committed
188

Tri Dao's avatar
Tri Dao committed
189
![FlashAttention speedup on A100 80GB SXM5 with FP16/BF16](assets/flash2_a100_fwd_bwd_benchmark.png)
Dan Fu's avatar
Dan Fu committed
190

Dan Fu's avatar
Dan Fu committed
191
#### Memory
Dan Fu's avatar
Dan Fu committed
192

193
![FlashAttention memory](assets/flashattn_memory.jpg)
Dan Fu's avatar
Dan Fu committed
194
195
196
197
198

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.
Tri Dao's avatar
Tri Dao committed
199

Tri Dao's avatar
Tri Dao committed
200
### H100
Dan Fu's avatar
Dan Fu committed
201

Tri Dao's avatar
Tri Dao committed
202
![FlashAttention speedup on H100 SXM5 with FP16/BF16](assets/flash2_h100_fwd_bwd_benchmark.png)
Dan Fu's avatar
Dan Fu committed
203

Tri Dao's avatar
Tri Dao committed
204
## Full model code and training script
Dan Fu's avatar
Dan Fu committed
205

Tri Dao's avatar
Tri Dao committed
206
207
208
209
210
211
212
We have released the full GPT model
[implementation](https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/models/gpt.py).
We also provide optimized implementations of other layers (e.g., MLP, LayerNorm,
cross-entropy loss, rotary embedding). Overall this speeds up training by 3-5x
compared to the baseline implementation from Huggingface, reaching up to 225
TFLOPs/sec per A100, equivalent to 72% model FLOPs utilization (we don't need
any activation checkpointing).
Dan Fu's avatar
T4  
Dan Fu committed
213

Tri Dao's avatar
Tri Dao committed
214
215
216
We also include a training
[script](https://github.com/Dao-AILab/flash-attention/tree/main/training) to
train GPT2 on Openwebtext and GPT3 on The Pile.
Dan Fu's avatar
T4  
Dan Fu committed
217

Tri Dao's avatar
Tri Dao committed
218
## Triton implementation of FlashAttention
Dan Fu's avatar
T4  
Dan Fu committed
219

Tri Dao's avatar
Tri Dao committed
220
221
Phil Tillet (OpenAI) has an experimental implementation of FlashAttention in Triton:
https://github.com/openai/triton/blob/master/python/tutorials/06-fused-attention.py
Dan Fu's avatar
T4  
Dan Fu committed
222

Tri Dao's avatar
Tri Dao committed
223
224
225
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.
Dan Fu's avatar
T4  
Dan Fu committed
226

Tri Dao's avatar
Tri Dao committed
227
228
229
We also have an experimental implementation in Triton that support attention
bias (e.g. ALiBi):
https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/flash_attn_triton.py
Dan Fu's avatar
T4  
Dan Fu committed
230
231


Tri Dao's avatar
Tri Dao committed
232
233
234
235
236
237
238
239
## Tests
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:
Ian Timmis's avatar
Ian Timmis committed
240
```sh
Tri Dao's avatar
Tri Dao committed
241
242
pytest -q -s tests/test_flash_attn.py
```
Tri Dao's avatar
Tri Dao committed
243
244
## When you encounter issues

Tri Dao's avatar
Tri Dao committed
245
This new release of FlashAttention-2 has been tested on several GPT-style
Tri Dao's avatar
Tri Dao committed
246
models, mostly on A100 GPUs.
Tri Dao's avatar
Tri Dao committed
247

Tri Dao's avatar
Tri Dao committed
248
If you encounter bugs, please open a GitHub Issue!
Dan Fu's avatar
Dan Fu committed
249
250
251
252

## Citation
If you use this codebase, or otherwise found our work valuable, please cite:
```
Tri Dao's avatar
Tri Dao committed
253
@inproceedings{dao2022flashattention,
254
  title={Flash{A}ttention: Fast and Memory-Efficient Exact Attention with {IO}-Awareness},
Dan Fu's avatar
Dan Fu committed
255
  author={Dao, Tri and Fu, Daniel Y. and Ermon, Stefano and Rudra, Atri and R{\'e}, Christopher},
Tri Dao's avatar
Tri Dao committed
256
  booktitle={Advances in Neural Information Processing Systems},
Dan Fu's avatar
Dan Fu committed
257
258
  year={2022}
}
Tri Dao's avatar
Tri Dao committed
259
@article{dao2023flashattention2,
260
  title={Flash{A}ttention-2: Faster Attention with Better Parallelism and Work Partitioning},
Tri Dao's avatar
Tri Dao committed
261
262
263
  author={Dao, Tri},
  year={2023}
}
Dan Fu's avatar
Dan Fu committed
264
```