flash-attention/README.md

FlashAttention - Alpha release (0.1).

To compile (requiring NVCC and an A100 GPU):

cd csrc/flash_attn
python setup.py install

Interface: flash_attention.py

To run the benchmark against PyTorch standard attention:

PYTHONPATH=$PWD python benchmarks/benchmark_flash_attention.py

FlashAttention currently supports:

1. A100 GPUs.
2. fp16.
3. Head dimensions 16, 32, 64.

Our roadmap to broaden the support:

1. Refactor to use Cutlass.
2. Support SM86 GPUs (e.g. RTX 3080, 3090), support SM75 GPUs (e.g. T4).
3. Support bf16.
4. Support head dimension 128.
5. Make package pip-installable.
6. Support SM70 GPUs (V100).
7. Fused rotary embedding.
8. Attention linear bias (e.g. ALiBi).

Speedup and Memory Savings

We present expected speedup (combined forward + backward pass) and memory savings from using FlashAttention against PyTorch standard attention, depending on sequence length. We display FlashAttention speedup using these parameters (similar to BERT-base):

• Batch size 8
• Head dimension 64
• 12 attention heads

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.

Speedup

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.

Memory

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.

Acknowledgments

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.