squall/flash-attention
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Link to Triton implementation

Browse Source
This commit is contained in:
Tri Dao 2022-07-11 15:57:45 -07:00
parent bc2c210254
commit 4577151ff8
1 changed files with 10 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
README.md
View File

@ -8,6 +8,16 @@ 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. IEEE Spectrum [article](https://spectrum.ieee.org/mlperf-rankings-2022) about our submission to the MLPerf 2.0 benchmark using FlashAttention.
![FlashAttention](assets/flashattn_banner.jpg) ![FlashAttention](assets/flashattn_banner.jpg)
#### Triton implementation of FlashAttention
Phil Tillet (OpenAI) has an 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.
## Alpha release (0.1). ## Alpha release (0.1).
To compile (requiring CUDA 11, NVCC, and an Turing or Ampere GPU): To compile (requiring CUDA 11, NVCC, and an Turing or Ampere GPU):