flash-attention/csrc/fused_dense_lib
2024-03-13 20:46:57 -07:00
..
fused_dense_cuda.cu [FusedDense] Allocate lt_workspace on input device 2023-05-30 14:17:26 -07:00
fused_dense.cpp [FusedDense] Allocate lt_workspace on input device 2023-05-30 14:17:26 -07:00
README.md Mention that some CUDA extensions have only been tested on A100s 2022-11-15 07:10:25 -08:00
setup.py Make nvcc threads configurable via environment variable (#885) 2024-03-13 20:46:57 -07:00

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.

It has only been tested on A100s.

cd csrc/fused_dense_lib && pip install .