Repository: pytorch
Issue: nn.LayerNorm is slower than naive implementation when dimension is low
Labels: module: performance,module: nn,triaged,module: norms and normalization

### 🐛 Describe the bug

```python
import torch
import torch.nn as nn
import time
import matplotlib.pyplot as plt
from tqdm import tqdm

class ElementwiseLayerNorm(nn.Module):
    def __init__(self, dim, eps=1e-5, elementwise_affine=True):
        super(ElementwiseLayerNorm, self).__init__()
        self.eps = eps
        self.elementwise_affine = elementwise_affine

        if self.elementwise_affine:
            self.weight = nn.Parameter(torch.ones(dim))
            self.bias = nn.Parameter(torch.zeros(dim))
        else:
            self.register_parameter('weight', None)
            self.register_parameter('bias', None)

    def forward(self, x):
        mean = x.mean(dim=-1, keepdim=True)
        var = x.var(dim=-1, keepdim=True, unbiased=False)
        x_normalized = (x - mean) / torch.sqrt(var + self.eps)

        if self.elementwise_affine:
            return x_normalized * self.weight + self.bias
        else:
            return x_normalized

embedding_sizes = [8, 16, 32, 64, 128,256]
seq_lens = [8, 16, 32, 64, 128,256]

results = {}
for el in embedding_sizes:
  results[el] = {}
  for el2 in seq_lens:
    results[el][el2] =  {"builtin": [], "custom": [], "builtin_compiled": [], "custom_compiled": []}
  

# Benchmark parameters
n_trials = 100
n_runs = 100
warmup_trials = 20
batch_size = 1024

with torch.no_grad():
  for i in tqdm(range(n_runs)):
    for seq_len in seq_lens:
      for embedding_size in embedding_sizes:
          # Input tensor
          x = torch.randn(batch_size, seq_len, embedding_size).cuda()

          # Built-in LayerNorm
          layernorm_builtin = nn.LayerNorm(embedding_size).cuda()
          layernorm_builtin_compiled = torch.compile(layernorm_builtin)

          layer_norm_custom = ElementwiseLayerNorm(embedding_size).cuda()
          layernorm_custom_compiled = torch.compile(layer_norm_custom)

          # Benchmark built-in LayerNorm
          torch.cuda.synchronize()
          start_time = time.time()
          for _ in range(n_trials):
              _ = layernorm_builtin(x)
          torch.cuda.synchronize()
          if i>warmup_trials: 
            results[embedding_size][seq_len]["builtin"].append((time.time() - start_time) / n_trials)

          torch.cuda.synchronize()
          start_time = time.time()
          for _ in range(n_trials):
              _ = layer_norm_custom(x)
          torch.cuda.synchronize()
          if i>warmup_trials: 
            results[embedding_size][seq_len]["custom"].append((time.time() - start_time) / n_trials)

          torch.cuda.synchronize()
          start_time = time.time()
          for _ in range(n_trials):
              _ = layernorm_builtin_compiled(x)
          torch.cuda.synchronize()
          if i>warmup_trials: 
            results[embedding_size][seq_len]["builtin_compiled"].append((time.time() - start_time) / n_trials)

          torch.cuda.synchronize()
          start_time = time.time()
          for _ in range(n_trials):
              _ = layernorm_custom_compiled(x)
          torch.cuda.synchronize()
          if i>warmup_trials: 
            results[embedding_size][seq_len]["custom_compiled"].append((time.time() - start_time) / n_trials)

```
```
from statistics import mean
for dim in results:
  for seq in results[dim]:
    for impl in results[dim][seq]:
      print(f"dim={dim}",f"seq_len={seq}",impl, mean(results[dim][seq][impl]))
    print(f"dim={dim}",f"seq_len={seq}","builtin/custom",mean(results[dim][seq]["builtin"])/mean(results[dim][seq]["custom"]))
    print(f"dim={dim}",f"seq_len={seq}","builtin_compiled/custom_compiled",mean(results[dim][seq]["builtin_compiled"])/mean(results[dim][seq]["custom_compiled"]))
    print() 
```
output
```
dim=8 seq_len=8 builtin 9.678224974040744e-05
dim=8 seq_len=8 custom 0.000110510874398147
dim=8 seq_len=8 builtin_compiled 8.088670199430441e-05
dim=8 seq_len=8 custom_compiled 0.00012810453583922567
dim=8 seq_len=8 builtin/custom 0.8757712783243551
dim=8 seq_len=8 builtin_compiled/custom_compiled 0.631411694085674

dim=8 seq_len=16 builtin 0.00016347541084772423
dim=8 seq_len=16 custom 0.00010370945628685287
dim=8 seq_len=16 builtin_compiled 0.00014441260808630834
dim=8 seq_len=16 custom_compiled 0.00012443533426598658
dim=8 seq_len=16 builtin/custom 1.5762825946706638
dim=8 seq_len=16 builtin_compiled/custom_compiled 1.1605434174960254

dim=8 seq_len=32 builtin 0.00032524501221089424
dim=8 seq_len=32 custom 0.00010442410843281806
dim=8 seq_len=32 builtin_compiled 0.0002761233003833626
dim=8 seq_len=32 custom_compiled 0.00012560237812090524
dim=8 seq_len=32 builtin/custom 3.1146544326987744
dim=8 seq_len=32 builtin_compiled/custom_compiled 2.1983922957061006

dim=8 seq_len=64 builtin 0.0006188744834706753
dim=8 seq_len=64 custom 0.00012607903420170651
dim=8 seq_len=64 builtin_compiled 0.0005796107762976538
dim=8 seq_len=64 custom_compiled 0.0001332149928129172
dim=8 seq_len=64 builtin/custom 4.90862328847296
dim=8 seq_len=64 builtin_compiled/custom_compiled 4.350942518246728

dim=8 seq_len=128 builtin 0.0012050605423842805
dim=8 seq_len=128 custom 0.0002844446218466457
dim=8 seq_len=128 builtin_compiled 0.0011724761467945726
dim=8 seq_len=128 custom_compiled 0.0002772958369194707
dim=8 seq_len=128 builtin/custom 4.236538327077149
dim=8 seq_len=128 builtin_compiled/custom_compiled 4.228250087775644

dim=8 seq_len=256 builtin 0.002282433721083629
dim=8 seq_len=256 custom 0.0005463570280920101
dim=8 seq_len=256 builtin_compiled 0.002378028012529204
dim=8 seq_len=256 custom_compiled 0.0005396363101428069
dim=8 seq_len=256 builtin/custom 4.1775498506065745
dim=8 seq_len=256 builtin_compiled/custom_compiled 4.406723505873602

dim=16 seq_len=8 builtin 6.951238535627534e-05
dim=16 seq_len=8 custom 0.00010947852195063724
dim=16 seq_len=8 builtin_compiled 6.628769862500927e-05
dim=16 seq_len=8 custom_compiled 0.00012823382510414606
dim=16 seq_len=8 builtin/custom 0.6349408460923301
dim=16 seq_len=8 builtin_compiled/custom_compiled 0.5169283421996749

dim=16 seq_len=16 builtin 0.00013065977941585492
dim=16 seq_len=16 custom 0.00010479531710660911
dim=16 seq_len=16 builtin_compiled 0.00012786113763157327
dim=16 seq_len=16 custom_compiled 0.00012438583977614776
dim=16 seq_len=16 builtin/custom 1.2468093329297691
dim=16 seq_len=16 builtin_compiled/custom_compiled 1.0279396582575626

dim=16 seq_len=32 builtin 0.0002697024767911887
dim=16 seq_len=32 custom 0.00011774214008186437
dim=16 seq_len=32 builtin_compiled 0.00028936654706544513
dim=16 seq_len=32 custom_compiled 0.00012548968761782103
dim=16 seq_len=32 builtin/custom 2.2906197951189653
dim=16 seq_len=32 builtin_compiled/custom_compiled 2.305899014958992

dim=16 seq_len=64 builtin 0.0005626950384695319
dim=16 seq_len=64 custom 0.0002793135522287103
dim=16 seq_len=64 builtin_compiled 0.0006131733520121514
dim=16 seq_len=64 custom_compiled 0.00028033552290518074
dim=16 seq_len=64 builtin/custom 2.0145640409484336
dim=16 seq_len=64 builtin_compiled/custom_compiled 2.1872838149717757

dim=16 seq_len=128 builtin 0.0011427697652502905
dim=16 seq_len=128 custom 0.0005438440057295787
dim=16 seq_len=128 builtin_compiled 0.0012262067915518073
dim=16 seq_len=128 custom_compiled 0.0005439215973962712
dim=16 seq_len=128 builtin/custom 2.101282266993528
dim=16 seq_len=128 builtin_compiled/custom_compiled 2.254381509066022

dim=16 seq_len=256 builtin 0.0022997066642664655
dim=16 seq_len=256 custom 0.0011024737961684602
dim=16 seq_len=256 builtin_compiled 0.0023997911018661306
dim=16 seq_len=256 custom_compiled 0.0011086066765121267
dim=16 seq_len=256 builtin/custom 2.0859513144519815
dim=16 seq_len=256 builtin_compiled/custom_compiled 2.164691186432594

dim=32 seq_len=8 builtin 6.437950496432147e-05
dim=32 seq_len=8 custom 0.00010460446152505995
dim=32 seq_len=8 builtin_compiled 6.390468983710567e-05
dim=32 seq_len=8 custom_compiled 0.00012611235244364678
dim=32 seq_len=8 builtin/custom 0.6154565878521172
dim=32 seq_len=8 builtin_compiled/custom_compiled 0.5067282355680538

dim=32 seq_len=16 builtin 0.00013311337821091278
dim=32 seq_len=16 custom 0.00011503971075709862
dim=32 seq_len=16 builtin_compiled 0.00014025141921224473
dim=32 seq_len=16 custom_compiled 0.00012984505182580102
dim=32 seq_len=16 builtin/custom 1.1571080745498041
dim=32 seq_len=16 builtin_compiled/custom_compiled 1.0801445048549467

dim=32 seq_len=32 builtin 0.00029121899906593034
dim=32 seq_len=32 custom 0.0002813394763801671
dim=32 seq_len=32 builtin_compiled 0.0003161558320250692
dim=32 seq_len=32 custom_compiled 0.0002892188180851031
dim=32 seq_len=32 builtin/custom 1.0351160200227758
dim=32 seq_len=32 builtin_compiled/custom_compiled 1.0931371413461757

dim=32 seq_len=64 builtin 0.0005759487574613547
dim=32 seq_len=64 custom 0.000560284656814382
dim=32 seq_len=64 builtin_compiled 0.000689772955978973
dim=32 seq_len=64 custom_compiled 0.0005597002898590475
dim=32 seq_len=64 builtin/custom 1.0279573971131641
dim=32 seq_len=64 builtin_compiled/custom_compiled 1.2323969961721521

dim=32 seq_len=128 builtin 0.0011702884601641306
dim=32 seq_len=128 custom 0.0011110351659074614
dim=32 seq_len=128 builtin_compiled 0.0012642960005168674
dim=32 seq_len=128 custom_compiled 0.0011257228066649618
dim=32 seq_len=128 builtin/custom 1.0533316100829921
dim=32 seq_len=128 builtin_compiled/custom_compiled 1.1230970830753968

dim=32 seq_len=256 builtin 0.0023688038693198674
dim=32 seq_len=256 custom 0.0022182996665375144
dim=32 seq_len=256 builtin_compiled 0.0024688906307461897
dim=32 seq_len=256 custom_compiled 0.002304515567006944
dim=32 seq_len=256 builtin/custom 1.0678466507716116
dim=32 seq_len=256 builtin_compiled/custom_compiled 1.0713273826797067

dim=64 seq_len=8 builtin 6.682106211215635e-05
dim=64 seq_len=8 custom 0.00010525736627699453
dim=64 seq_len=8 builtin_compiled 6.782447235493721e-05
dim=64 seq_len=8 custom_compiled 0.00013141333302365074
dim=64 seq_len=8 builtin/custom 0.6348350189221961
dim=64 seq_len=8 builtin_compiled/custom_compiled 0.5161156086249687

dim=64 seq_len=16 builtin 0.00015420886534678785
dim=64 seq_len=16 custom 0.00023138170000873036
dim=64 seq_len=16 builtin_compiled 0.0001589152480982527
dim=64 seq_len=16 custom_compiled 0.00023725310458412653
dim=64 seq_len=16 builtin/custom 0.6664695839859822
dim=64 seq_len=16 builtin_compiled/custom_compiled 0.6698131448134692

dim=64 seq_len=32 builtin 0.00029486447949952715
dim=64 seq_len=32 custom 0.0004519757138022894
dim=64 seq_len=32 builtin_compiled 0.0003595852851867676
dim=64 seq_len=32 custom_compiled 0.00045755332029318507
dim=64 seq_len=32 builtin/custom 0.652390096403524
dim=64 seq_len=32 builtin_compiled/custom_compiled 0.7858871725733674

dim=64 seq_len=64 builtin 0.0005870274048817308
dim=64 seq_len=64 custom 0.0008933664273612107
dim=64 seq_len=64 builtin_compiled 0.0006842062443117552
dim=64 seq_len=64 custom_compiled 0.0008906397034850301
dim=64 seq_len=64 builtin/custom 0.6570958868643278
dim=64 seq_len=64 builtin_compiled/custom_compiled 0.7682188898995735

dim=64 seq_len=128 builtin 0.0012294705306427388
dim=64 seq_len=128 custom 0.0017652433733396891
dim=64 seq_len=128 builtin_compiled 0.0013268066056167022
dim=64 seq_len=128 custom_compiled 0.0017648941655702228
dim=64 seq_len=128 builtin/custom 0.6964878323359379
dim=64 seq_len=128 builtin_compiled/custom_compiled 0.7517768665680993

dim=64 seq_len=256 builtin 0.00250308896921858
dim=64 seq_len=256 custom 0.0035207106795492053
dim=64 seq_len=256 builtin_compiled 0.0026342577270314664
dim=64 seq_len=256 custom_compiled 0.0036150625687611255
dim=64 seq_len=256 builtin/custom 0.7109612794252885
dim=64 seq_len=256 builtin_compiled/custom_compiled 0.7286893869541575

dim=128 seq_len=8 builtin 7.666261890266515e-05
dim=128 seq_len=8 custom 0.00021290081965772412
dim=128 seq_len=8 builtin_compiled 7.994509950468812e-05
dim=128 seq_len=8 custom_compiled 0.00021764854841594455
dim=128 seq_len=8 builtin/custom 0.3600860674276122
dim=128 seq_len=8 builtin_compiled/custom_compiled 0.3673128081327993

dim=128 seq_len=16 builtin 0.0001518993438044681
dim=128 seq_len=16 custom 0.00040543532069725326
dim=128 seq_len=16 builtin_compiled 0.00017205464689037467
dim=128 seq_len=16 custom_compiled 0.00041323366044442863
dim=128 seq_len=16 builtin/custom 0.3746574016867524
dim=128 seq_len=16 builtin_compiled/custom_compiled 0.41636164562521755

dim=128 seq_len=32 builtin 0.0002870179429838929
dim=128 seq_len=32 custom 0.0007859975778603856
dim=128 seq_len=32 builtin_compiled 0.0003618722324129901
dim=128 seq_len=32 custom_compiled 0.0007874307753164557
dim=128 seq_len=32 builtin/custom 0.36516390262321524
dim=128 seq_len=32 builtin_compiled/custom_compiled 0.45956069251618914

dim=128 seq_len=64 builtin 0.0006045847603037387
dim=128 seq_len=64 custom 0.0015403484996361068
dim=128 seq_len=64 builtin_compiled 0.0007178449328941635
dim=128 seq_len=64 custom_compiled 0.0015334002881110468
dim=128 seq_len=64 builtin/custom 0.3924986848408437
dim=128 seq_len=64 builtin_compiled/custom_compiled 0.4681392969988657

dim=128 seq_len=128 builtin 0.0013390024100677876
dim=128 seq_len=128 custom 0.003025425234927407
dim=128 seq_len=128 builtin_compiled 0.001430505529234681
dim=128 seq_len=128 custom_compiled 0.003132360253152968
dim=128 seq_len=128 builtin/custom 0.4425832093318664
dim=128 seq_len=128 builtin_compiled/custom_compiled 0.45668614515037476

dim=128 seq_len=256 builtin 0.0027464558504804782
dim=128 seq_len=256 custom 0.006132948941822294
dim=128 seq_len=256 builtin_compiled 0.003066002900087381
dim=128 seq_len=256 custom_compiled 0.00601917876472956
dim=128 seq_len=256 builtin/custom 0.44781978075043605
dim=128 seq_len=256 builtin_compiled/custom_compiled 0.5093722947810033

dim=256 seq_len=8 builtin 9.337358836886249e-05
dim=256 seq_len=8 custom 0.00039830766146696067
dim=256 seq_len=8 builtin_compiled 9.896996655041659e-05
dim=256 seq_len=8 custom_compiled 0.0004038087023964411
dim=256 seq_len=8 builtin/custom 0.2344257904178119
dim=256 seq_len=8 builtin_compiled/custom_compiled 0.24509121760643077

dim=256 seq_len=16 builtin 0.00017236368565619746
dim=256 seq_len=16 custom 0.0007606318630749666
dim=256 seq_len=16 builtin_compiled 0.00021217062503476686
dim=256 seq_len=16 custom_compiled 0.0007664575154268289
dim=256 seq_len=16 builtin/custom 0.2266059233429846
dim=256 seq_len=16 builtin_compiled/custom_compiled 0.2768198116194505

dim=256 seq_len=32 builtin 0.0003308805634703817
dim=256 seq_len=32 custom 0.00148523535909532
dim=256 seq_len=32 builtin_compiled 0.0004414746127551115
dim=256 seq_len=32 custom_compiled 0.001480541108529779
dim=256 seq_len=32 builtin/custom 0.2227798856552447
dim=256 seq_len=32 builtin_compiled/custom_compiled 0.29818463682748314

dim=256 seq_len=64 builtin 0.0008268203916428964
dim=256 seq_len=64 custom 0.0029062589512595647
dim=256 seq_len=64 builtin_compiled 0.0009188679803775836
dim=256 seq_len=64 custom_compiled 0.002903341281263134
dim=256 seq_len=64 builtin/custom 0.2844964628098796
dim=256 seq_len=64 builtin_compiled/custom_compiled 0.31648638288152564

dim=256 seq_len=128 builtin 0.0018138287037233764
dim=256 seq_len=128 custom 0.00574048962774156
dim=256 seq_len=128 builtin_compiled 0.002191464750072624
dim=256 seq_len=128 custom_compiled 0.005744126416459868
dim=256 seq_len=128 builtin/custom 0.31597107935843066
dim=256 seq_len=128 builtin_compiled/custom_compiled 0.38151401817915315

dim=256 seq_len=256 builtin 0.0036927097658567792
dim=256 seq_len=256 custom 0.01159543124935295
dim=256 seq_len=256 builtin_compiled 0.003931530155713045
dim=256 seq_len=256 custom_compiled 0.011449624496170237
dim=256 seq_len=256 builtin/custom 0.31846247771619884
dim=256 seq_len=256 builtin_compiled/custom_compiled 0.34337634016103274
```

if builtin/custom is > 1, it means that nn.LayerNorm is slower than custom implementation


### Versions

Collecting environment information...
PyTorch version: 2.5.1+cu121
Is debug build: False
CUDA used to build PyTorch: 12.1
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.3 LTS (x86_64)
GCC version: (Ubuntu 11.4.0-1ubuntu1~22.04) 11.4.0
Clang version: 14.0.0-1ubuntu1.1
CMake version: version 3.31.2
Libc version: glibc-2.35

Python version: 3.10.12 (main, Nov  6 2024, 20:22:13) [GCC 11.4.0] (64-bit runtime)
Python platform: Linux-6.1.85+-x86_64-with-glibc2.35
Is CUDA available: True
CUDA runtime version: 12.2.140
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: Tesla T4
Nvidia driver version: 535.104.05
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.9.6
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.9.6
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.9.6
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.9.6
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.9.6
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.9.6
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.9.6
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:                         x86_64
CPU op-mode(s):                       32-bit, 64-bit
Address sizes:                        46 bits physical, 48 bits virtual
Byte Order:                           Little Endian
CPU(s):                               2
On-line CPU(s) list:                  0,1
Vendor ID:                            GenuineIntel
Model name:                           Intel(R) Xeon(R) CPU @ 2.00GHz
CPU family:                           6
Model:                                85
Thread(s) per core:                   2
Core(s) per socket:                   1
Socket(s):                            1
Stepping:                             3
BogoMIPS:                             4000.35
Flags:                                fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ss ht syscall nx pdpe1gb rdtscp lm constant_tsc rep_good nopl xtopology nonstop_tsc cpuid tsc_known_freq pni pclmulqdq ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt aes xsave avx f16c rdrand hypervisor lahf_lm abm 3dnowprefetch invpcid_single ssbd ibrs ibpb stibp fsgsbase tsc_adjust bmi1 hle avx2 smep bmi2 erms invpcid rtm mpx avx512f avx512dq rdseed adx smap clflushopt clwb avx512cd avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves arat md_clear arch_capabilities
Hypervisor vendor:                    KVM
Virtualization type:                  full
L1d cache:                            32 KiB (1 instance)
L1i cache:                            32 KiB (1 instance)
L2 cache:                             1 MiB (1 instance)
L3 cache:                             38.5 MiB (1 instance)
NUMA node(s):                         1
NUMA node0 CPU(s):                    0,1
Vulnerability Gather data sampling:   Not affected
Vulnerability Itlb multihit:          Not affected
Vulnerability L1tf:                   Mitigation; PTE Inversion
Vulnerability Mds:                    Vulnerable; SMT Host state unknown
Vulnerability Meltdown:               Vulnerable
Vulnerability Mmio stale data:        Vulnerable
Vulnerability Reg file data sampling: Not affected
Vulnerability Retbleed:               Vulnerable
Vulnerability Spec rstack overflow:   Not affected
Vulnerability Spec store bypass:      Vulnerable
Vulnerability Spectre v1:             Vulnerable: __user pointer sanitization and usercopy barriers only; no swapgs barriers
Vulnerability Spectre v2:             Vulnerable; IBPB: disabled; STIBP: disabled; PBRSB-eIBRS: Not affected; BHI: Vulnerable (Syscall hardening enabled)
Vulnerability Srbds:                  Not affected
Vulnerability Tsx async abort:        Vulnerable

Versions of relevant libraries:
[pip3] numpy==1.26.4
[pip3] nvidia-cublas-cu12==12.6.4.1
[pip3] nvidia-cuda-cupti-cu12==12.6.80
[pip3] nvidia-cuda-runtime-cu12==12.6.77
[pip3] nvidia-cudnn-cu12==9.6.0.74
[pip3] nvidia-cufft-cu12==11.3.0.4
[pip3] nvidia-curand-cu12==10.3.7.77
[pip3] nvidia-cusolver-cu12==11.7.1.2
[pip3] nvidia-cusparse-cu12==12.5.4.2
[pip3] nvidia-nccl-cu12==2.23.4
[pip3] nvidia-nvjitlink-cu12==12.6.85
[pip3] nvtx==0.2.10
[pip3] optree==0.13.1
[pip3] pynvjitlink-cu12==0.4.0
[pip3] torch==2.5.1+cu121
[pip3] torchaudio==2.5.1+cu121
[pip3] torchsummary==1.5.1
[pip3] torchvision==0.20.1+cu121
[pip3] triton==3.1.0
[conda] Could not collect

cc @msaroufim @albanD @mruberry @jbschlosser @walterddr @mikaylagawarecki