Metadata-Version: 2.4
Name: torchnorm
Version: 0.1.2
Summary: Fused Triton kernels for RMSNorm, LayerNorm, FusedAddRMSNorm, GroupNorm, InstanceNorm, ScaleNorm
Author-email: Liodon AI <py.oss.ml@gmail.com>
License: MIT
Project-URL: Homepage, https://github.com/liodon-ai/torchnorm
Project-URL: Repository, https://github.com/liodon-ai/torchnorm
Project-URL: Issues, https://github.com/liodon-ai/torchnorm/issues
Project-URL: Part of, https://github.com/liodon-ai/torchembed
Keywords: pytorch,triton,normalization,rmsnorm,layernorm,deep-learning,transformers,llm,fused-kernels,gpu,cuda,training,llama,mistral
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Science/Research
Classifier: Intended Audience :: Developers
Classifier: License :: OSI Approved :: MIT License
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.9
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Topic :: Software Development :: Libraries :: Python Modules
Requires-Python: >=3.9
Description-Content-Type: text/markdown
Requires-Dist: torch>=2.0
Provides-Extra: triton
Requires-Dist: triton>=2.1; extra == "triton"
Provides-Extra: dev
Requires-Dist: pytest>=7.0; extra == "dev"
Requires-Dist: pytest-cov; extra == "dev"
Requires-Dist: ruff; extra == "dev"
Requires-Dist: triton>=2.1; extra == "dev"

# 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

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

From source:

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

## Modules

### RMSNorm

Root Mean Square Layer Normalisation ([Zhang & Sennrich, 2019](https://arxiv.org/abs/1910.07467)). Used in LLaMA, Mistral, Gemma, Qwen, Falcon, and most modern LLMs.

```python
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:

```python
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](https://arxiv.org/abs/1607.06450)). Identical interface to `torch.nn.LayerNorm`.

```python
from torchnorm import LayerNorm

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

### GroupNorm

Group Normalisation ([Wu & He, 2018](https://arxiv.org/abs/1803.08494)). Standard for diffusion UNets and CNNs.

```python
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](https://arxiv.org/abs/1607.08022)). Normalises each sample–channel pair over spatial dims. Standard in style transfer.

```python
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](https://arxiv.org/abs/1910.05895)). L2-normalises and scales by a single learnable scalar. Cheaper than LayerNorm — no mean subtraction, no per-element parameters.

```python
from torchnorm import ScaleNorm

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

## Dispatch logic

Every module auto-selects the fastest available path:

| Condition | Path |
|---|---|
| CUDA tensor + triton installed | Fused Triton kernel |
| CPU tensor or triton not installed | PyTorch (`F.rms_norm`, `F.layer_norm`, …) |

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

## Compatibility

| | |
|---|---|
| Python | ≥ 3.9 |
| PyTorch | ≥ 2.0 |
| Triton | ≥ 2.1 (optional) |
| CUDA | any version supported by your PyTorch build |

## Related projects

| Project | Focus | Install |
|---|---|---|
| [Liger Kernel](https://github.com/linkedin/Liger-Kernel) | Full training suite: RMSNorm + RoPE + SwiGLU + CrossEntropy + model patching | `pip install liger-kernel` |
| [flash-attn](https://github.com/Dao-AILab/flash-attention) | Flash Attention + fused ops including LayerNorm | build from source |
| [torchembed](https://github.com/liodon-ai/torchembed) | Companion library: RoPE (rotate-half & adjacent-pairs), ALiBi, patch, categorical embeddings | `pip install torchembed` |

**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**](https://github.com/liodon-ai/torchembed) — fused kernels for embedding layers (RoPE, patch, categorical)
- [**torchnorm**](https://github.com/liodon-ai/torchnorm) — fused kernels for normalisation layers
