torchnorm

Fused Triton kernels for normalisation layers. Drop-in replacements for PyTorch's norm modules with automatic dispatch to optimised CUDA kernels when triton is available.

Installation

pip install torchnorm            # PyTorch fallback only
pip install "torchnorm[triton]"  # with Triton fused kernels

From source:

git clone https://github.com/liodon-ai/torchnorm
cd torchnorm
pip install -e ".[triton]"

Modules

RMSNorm

Root Mean Square Layer Normalisation (Zhang & Sennrich, 2019). Used in LLaMA, Mistral, Gemma, Qwen, Falcon, and most modern LLMs.

from torchnorm import RMSNorm

norm = RMSNorm(dim=4096).cuda()
x = torch.randn(4, 512, 4096, device="cuda", dtype=torch.float16)
out = norm(x)  # fused Triton kernel on CUDA, PyTorch fallback elsewhere

The fused kernel reads each element once, computes the RMS, and applies the weight in a single pass — no intermediate allocation.

FusedAddRMSNorm

Fuses the residual add and RMSNorm into one kernel. Canonical pattern in transformer pre-norm blocks:

from torchnorm import FusedAddRMSNorm

norm = FusedAddRMSNorm(dim=4096).cuda()

# Instead of:
#   residual = x + residual
#   hidden   = rms_norm(residual)

# Write:
hidden, residual = norm(x, residual)

Returns both the normalised output and the updated residual so the next block can use it directly.

LayerNorm

Layer Normalisation (Ba et al., 2016). Identical interface to torch.nn.LayerNorm.

from torchnorm import LayerNorm

norm = LayerNorm(768, bias=False).cuda()   # BERT-style, no bias
out  = norm(x)

GroupNorm

Group Normalisation (Wu & He, 2018). Standard for diffusion UNets and CNNs.

from torchnorm import GroupNorm

norm = GroupNorm(num_groups=32, num_channels=512).cuda()
out  = norm(x)   # x: (N, C, H, W)

InstanceNorm

Instance Normalisation (Ulyanov et al., 2016). Normalises each sample–channel pair over spatial dims. Standard in style transfer.

from torchnorm import InstanceNorm

norm = InstanceNorm(num_features=512, affine=True).cuda()
out  = norm(x)   # x: (N, C, H, W) or (N, C)

ScaleNorm

ScaleNorm (Nguyen & Salazar, 2019). L2-normalises and scales by a single learnable scalar. Cheaper than LayerNorm — no mean subtraction, no per-element parameters.

from torchnorm import ScaleNorm

norm = ScaleNorm(dim=512).cuda()
out  = norm(x)

Dispatch logic

Every module auto-selects the fastest available path:

No flags to set. Installing triton is the only opt-in required.

Compatibility

Related projects

torchnorm vs Liger Kernel: Liger patches entire model classes at import time — ideal when you want the full optimization suite with use_liger_kernel=True. torchnorm provides individual nn.Module drop-ins without model patching — useful when building custom architectures or when you only need the norm layers.

Part of the Liodon AI stack

• torchembed — fused kernels for embedding layers (RoPE, patch, categorical)
• torchnorm — fused kernels for normalisation layers