A minimal re-implementation of Flash Attention with CUDA and PyTorch. The official implementation can be quite daunting for a CUDA beginner (like myself), so this repo tries to be small and educational.
- The entire forward pass is written in ~100 lines in
flash.cu. - The variable names follow the notations from the original paper.
- PyTorch (with CUDA)
Ninjafor loading in C++
Compare the wall-clock time between manual attention and minimal flash attention:
python bench.py
Sample output on a T4:
=== profiling manual attention ===
...
Self CPU time total: 52.389ms
Self CUDA time total: 52.545ms
=== profiling minimal flash attention ===
...
Self CPU time total: 11.452ms
Self CUDA time total: 3.908ms
Speed-up achieved!
Try out this online colab demo.
- No backward pass! To be honest, I found it a lot more complex than the forward pass, which was enough to show the use of shared memory to avoid large N^2 read/writes.
- In the inner loop, I assign each thread to a row of the output matrix. This differs from the original implementation.
- This thread-per-row simplification makes the matrix multiplications very slow. This is probably why for longer sequences and larger block sizes, this gets slower than the manual implementation.
- Q,K,Vs are in float32, unlike the original implementation which uses float16.
- The block size is fixed at compile time to 32.
- Add backward pass
- Speed up matmults
- Dynamically set block size