squall/flash-attention
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
2e33fc8e36
flash-attention/csrc/fused_dense_lib
History
Tri Dao 2e33fc8e36 Add GPT and ViT models
2022-11-13 22:30:23 -08:00
..
fused_dense_cuda.cu Add fused_dense and dropout_add_layernorm CUDA extensions 2022-11-13 21:59:20 -08:00
fused_dense.cpp Add fused_dense and dropout_add_layernorm CUDA extensions 2022-11-13 21:59:20 -08:00
README.md Add GPT and ViT models 2022-11-13 22:30:23 -08:00
setup.py Add fused_dense and dropout_add_layernorm CUDA extensions 2022-11-13 21:59:20 -08:00

README.md

This CUDA extension implements fused matmul + bias (forward and backward), and fused matmul + bias + gelu (forward and backward), adapted from Apex's FusedDense. We make it work for bfloat16.

For best performance, you should use CUDA >= 11.8. CuBLAS versions before this doesn't have the best matmul + bias + gelu performance for bfloat16.

cd csrc/fused_dense_lib && pip install .