README.md 438 Bytes
Newer Older
Tri Dao's avatar
Tri Dao committed
1
This CUDA extension implements fused dropout + residual + LayerNorm, based on
2
3
Apex's [FastLayerNorm](https://github.com/NVIDIA/apex/tree/master/apex/contrib/layer_norm).
We add dropout and residual, and make it work for both pre-norm and post-norm architecture.
4

Tri Dao's avatar
Tri Dao committed
5
6
This only supports a limited set of dimensions, see `csrc/layer_norm/ln_fwd_cuda_kernel.cu`.

7
8
It has only been tested on A100s.

9
10
11
```sh
cd csrc/layer_norm && pip install .
```