Custom GPU kernel for Root Mean Square layer normalization layer with fused operations
- Implemented with Triton compiler for high performance and parallel computations on GPUs.
- Includes both forward and backward passes of RMS layer normalization.
Requirements
torch==2.1.0+cu121
torchaudio==2.1.0+cu121
torchvision==0.16.0+cu121
triton==2.1.0
You can install the package using pip3 install -e .
:
pip3 install -e .
Includes benchmarking against LayerNorm, Pytorchs jit compiled RMSNorm, and Triton RMSNorm with the following configurations:
batch_size = 250
last_dim_sizes = [1024, 2048, 4096, 8192]
num_passes = 1000
do_backward