Add fused Modulated RMSNorm for DiT-style scale/shift conditioning

[yukiu00]

The feature, motivation

I would like to add a fused Modulated RMSNorm kernel for AdaLN/FiLM-style conditioning patterns used in diffusion Transformers.

The operation is:

y = rms_norm(x, weight, eps) * (1 + scale) + shift

with shift optional:

y = rms_norm(x, weight, eps) * (1 + scale)

This pattern shows up in DiT-style blocks where a conditioning vector produces per-layer modulation parameters. In current PyTorch-style implementations, RMSNorm and the scale/shift modulation are usually separate ops, which means extra kernel launches and an intermediate normalized tensor.

Liger already has a fast RMSNorm kernel, so this seems like a small and useful extension for diffusion/video Transformer workloads.

Proposed API:

liger_modulated_rms_norm(
    X,
    W,
    scale,
    shift=None,
    eps=1e-6,
    offset=0.0,
    casting_mode="llama",
    in_place=True,
)

and a module wrapper:

LigerModulatedRMSNorm(hidden_size, eps=1e-6, ...)

Initial scope:

• Fuse RMSNorm + modulation scale + optional shift
• Support scale/shift shaped per-row or per-batch, broadcast over tokens
• Support W=None / elementwise_affine=False
• Keep existing RMSNorm options where possible: offset, casting_mode, in_place
• Add correctness tests for forward and backward against a PyTorch reference
• Add a benchmark comparing:
  • PyTorch/HF RMSNorm + modulation
  • LigerRMSNorm + modulation
  • fused LigerModulatedRMSNorm

Out of scope for the first PR:

• Fusing the modulation MLP that produces scale and shift
• Fusing residual gates
• Diffusers monkey-patching
• LayerNorm + modulation

Alternatives

The current alternative is to use LigerRMSNorm and then apply scale/shift with regular PyTorch ops. That works, but still materializes the normalized output and launches extra elementwise kernels.

Additional context

This is a common pattern in recent diffusion Transformer architectures:

• FiLM introduced feature-wise affine modulation:
  https://ojs.aaai.org/index.php/AAAI/article/view/11671
• DiT uses adaptive layer norm / adaLN-Zero conditioning:
  https://www.wpeebles.com/DiT
  https://github.com/facebookresearch/DiT/blob/main/models.py
• FLUX uses per-block shift, scale, and gate modulation:
  https://github.com/black-forest-labs/flux/blob/main/src/flux/modules/layers.py
• Diffusers supports adaptive normalization with RMSNorm variants:
  https://huggingface.co/docs/diffusers/api/normalization
  https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/normalization.py

[yukiu00]

@Tcc0403 What do you think about this proposal?

[yukiu00]

Opened a draft PR for this proposal: #1225

Benchmarked on an RTX 3090 with bf16, hidden size 4096, row-wise scale and shift.

Speed vs LigerRMSNorm + torch modulation:

• BT=1024 full: 0.1403 ms vs 0.2642 ms, 1.9x faster
• BT=4096 full: 0.4787 ms vs 1.3793 ms, 2.9x faster
• Forward is about 2.5x-2.6x faster

Memory vs LigerRMSNorm + torch modulation:

• Forward: about 43% lower peak memory
• Backward/full: about 8% lower peak memory

Correctness test passes: test/transformers/test_modulated_rms_norm.py with 152 passed.

[Tcc0403]

Sure, +1 on it. I wonder how it performs against torch.compile. cc @Mecoli1219 @vaibhavjindal

[yukiu00]

@Tcc0403 I ran a quick experiment against torch.compile.

I kept the benchmark change on a separate branch/commit and did not mix it into the current draft PR #1225:

• branch: https://github.com/yukiu00/Liger-Kernel/tree/codex-modulated-rms-norm-torch-compile-bench
• commit: eae3e79

Setup: RTX 3090, PyTorch 2.10.0+cu128, bf16, hidden size 4096, row-wise scale and shift, p50 latency. The torch_compile baseline is a simple torch.compile(NaiveModulatedRMSNorm(...)) provider. This was a small local sweep because the machine only had ~2 GB free VRAM at the time.

BT	Liger fwd	torch.compile fwd	Liger bwd	torch.compile bwd	Liger full mem	torch.compile full mem
1024	0.043 ms	0.053 ms	0.077 ms	0.148 ms	89 MB	104 MB
2048	0.080 ms	0.103 ms	0.134 ms	0.291 ms	177 MB	208 MB
4096	0.156 ms	0.203 ms	0.255 ms	0.569 ms	353 MB	416 MB
8192	0.305 ms	0.399 ms	0.493 ms	0.607 ms	705 MB	840 MB

So in this quick comparison, the fused Liger kernel is still faster than torch.compile: about 1.24-1.31x on forward and 1.23-2.23x on backward, with about 14-16% lower full forward+backward peak memory.

[Tcc0403]

good to know, looking forward to your contribution!