Inicio Explorar Axuda
Rexistro Iniciar sesión
david
/
flash-attention
réplica de https://github.com/Dao-AILab/flash-attention
1
0
Fork 0
Ficheiros Incidencias 0 Wiki
Árbore: 0dfb281743
Ramas Etiquetas
flash-attention
/
csrc
/
fused_dense_lib
sclarkson 1feb711f46 Fix compilation with clang on ARM64 (#1285) hai 1 semana
..
README.md 43ab0b5205 Mention that some CUDA extensions have only been tested on A100s %!s(int64=2) %!d(string=hai) anos
fused_dense.cpp 1feb711f46 Fix compilation with clang on ARM64 (#1285) hai 1 semana
fused_dense_cuda.cu 27f8f890df [FusedDense] Allocate lt_workspace on input device hai 1 ano
setup.py 50896ec574 Make nvcc threads configurable via environment variable (#885) hai 9 meses

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.

It has only been tested on A100s.

cd csrc/fused_dense_lib && pip install .