Repository: pytorch
Issue: [RFC] Improve performance for layer_norm op for cuda with revectorized
Labels: module: nn,module: cuda,triaged,topic: performance

### 🚀 The feature, motivation and pitch

I found there is big perf drop if the size of layer_norm's inside size is not multiple of 4, there is a micro test case:
```python
import torch

DEVICE=torch.device('cuda')

# Time cost for near 1024
for cnt in range(2040, 2050):
    x = torch.randn(4096, cnt, device=DEVICE, dtype=torch.float32)
    w_shape = (x.shape[-1], )

    #warm up
    need_warmup = True
    round = 5
    if need_warmup:
        for _ in range(round):
            output = torch.nn.functional.layer_norm(x, w_shape)

    torch.cuda.synchronize()
    start_time = torch.cuda.Event(enable_timing=True)
    end_time = torch.cuda.Event(enable_timing=True)

    # Start time
    start_time.record()

    # Apply layernorm
    for _ in range(round):
        output = torch.nn.functional.layer_norm(x, w_shape)

    # End time
    end_time.record()

    torch.cuda.synchronize()

    # Calculate elapsed time
    elapsed_time_ms = start_time.elapsed_time(end_time)
    # print(f"CUDA Time: {elapsed_time_ms:.6f} ms")
    gbps = lambda ms: round * 2 * x.numel() * x.element_size() * 1e-9 / (ms * 1e-3)
    print(f"n as {cnt} of softmax: {gbps(elapsed_time_ms):.6f} gb/s")
```

Its output is:
```
n as 2040 of softmax: 483.555543 gb/s
n as 2041 of softmax: 345.531858 gb/s
n as 2042 of softmax: 345.369984 gb/s
n as 2043 of softmax: 347.825623 gb/s
n as 2044 of softmax: 489.580822 gb/s
n as 2045 of softmax: 345.591973 gb/s
n as 2046 of softmax: 346.057951 gb/s
n as 2047 of softmax: 345.850055 gb/s
n as 2048 of softmax: 470.192376 gb/s
n as 2049 of softmax: 347.012446 gb/s
```
We could see for the perf for input with `N = 2040, 2044, 2048` is obvious greater than other input `(480 vs 340)`. 

Therefore, what I would like to do is **mitigating the perf gap between these input**.

### Alternatives

The root cause is that there is two kernels for `layer_norm`:
https://github.com/pytorch/pytorch/blob/5e4cf3e6ad6f1f06436f409b394ae02e5ed5583d/aten/src/ATen/native/cuda/layer_norm_kernel.cu#L824-L837
We could see if `N` is multiple of `num_vec_elems (4)`, it will call a kernel called `launch_vectorized_layer_norm_kernel`, which could load vectorized elements.

So what we could do is to also enable vectorized elements load for those case whose `N` is not multiple for `num_vec_elems (4)`.

I tried a draft implement to achieve it, it could improve performance same as the vectorized case, for this case it's about **~40%**, `(480 vs 340) take 2042 as example`.

Optimized data:
```
n as 2040 of softmax: 459.758796 gb/s
n as 2041 of softmax: 497.804886 gb/s
n as 2042 of softmax: 479.061575 gb/s
n as 2043 of softmax: 477.197096 gb/s
n as 2044 of softmax: 473.285078 gb/s
n as 2045 of softmax: 473.795206 gb/s
n as 2046 of softmax: 495.999985 gb/s
n as 2047 of softmax: 499.649134 gb/s
n as 2048 of softmax: 455.111104 gb/s
n as 2049 of softmax: 473.928423 gb/s
```

### Additional context

Here is a [doc](https://docs.google.com/document/d/1TFIbJAO3tek1-EltVvMC_0TgYuEeBZw8rxDrSb_Wx8Q/edit?tab=t.0) contains some stuffs about it.

And there is another optimization, we could see there is still perf gap for `layer_norm` in `pytorch` and `triton`, we may mitigate the gap by using **register to cache the data in pytorch kernel**, because the `launch_vectorized_layer_norm_kernel` load data from gmem twice yet, but I guess the `N` for `layer_norm` op may usually be a very big number (> 10k for some 2d layer_norm), so it may introduce much register pressure, WDYT?

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