Repository: pytorch
Issue: rms_norm forward performance analysis
Labels: triaged,oncall: pt2,module: inductor

### 🐛 Describe the bug

In our previous tests, the performance of rms_norm operator compiled by inductor is worse than liger implementation. Here is the root cause analysis. 


## Reproduce tests

```
% python run.py --op rms_norm  --mode fwd  --precision fp32 --metrics latency,speedup --cudagraph
100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 6/6 [00:04<00:00,  1.24it/s]
       (M, H)    llama_rms-latency    liger_rms-speedup    liger_rms-latency    inductor_rms-speedup    inductor_rms-latency
-------------  -------------------  -------------------  -------------------  ----------------------  ----------------------
 (2048, 1024)             0.041888              3.50938             0.011936                 3.20833                0.013056
 (2048, 2048)             0.06192               3.01402             0.020544                 2.94521                0.021024
 (2048, 4096)             0.135552              3.80251             0.035648                 3.76533                0.036
 (2048, 8192)             0.25376               3.83092             0.06624                  3.81985                0.066432
(2048, 16384)             0.482496              3.71929             0.129728                 3.37996                0.142752
(2048, 32768)             0.926784              3.44909             0.268704                 2.5852                 0.358496
```

## Obtain NCU Profile Reports

```
python run.py --op rms_norm --mode fwd --precision fp32 --metrics ncu_rep 
```

## Report Analysis

### For input id 1-5

For input id 5, which is the largest input in our tests, the execution times are 0.35ms and 0.26ms for inductor and liger separately.  

The inductor compiled kernel is the following. 

```

# kernel path: /tmp/torchinductor_yhao/lc/clcsnr7zzaoxdu2qmftawo3xj5jh3rzkjdfrno47offog26pko6z.py
# Topologically Sorted Source Nodes: [pow_1, variance, add, rsqrt, hidden_states_1, mul_1], Original ATen: [aten.pow, aten.mean, aten.add, aten.rsqrt, aten.mul]
# Source node to ATen node mapping:
#   add => add
#   hidden_states_1 => mul
#   mul_1 => mul_1
#   pow_1 => pow_1
#   rsqrt => rsqrt
#   variance => mean
# Graph fragment:
#   %pow_1 : [num_users=1] = call_function[target=torch.ops.aten.pow.Tensor_Scalar](args = (%primals_1, 2), kwargs = {})
#   %mean : [num_users=1] = call_function[target=torch.ops.aten.mean.dim](args = (%pow_1, [-1], True), kwargs = {})
#   %add : [num_users=1] = call_function[target=torch.ops.aten.add.Tensor](args = (%mean, 1e-06), kwargs = {})
#   %rsqrt : [num_users=2] = call_function[target=torch.ops.aten.rsqrt.default](args = (%add,), kwargs = {})
#   %mul : [num_users=1] = call_function[target=torch.ops.aten.mul.Tensor](args = (%primals_1, %rsqrt), kwargs = {})
#   %mul_1 : [num_users=1] = call_function[target=torch.ops.aten.mul.Tensor](args = (%primals_2, %mul), kwargs = {})
triton_red_fused_add_mean_mul_pow_rsqrt_0 = async_compile.triton('triton_red_fused_add_mean_mul_pow_rsqrt_0', '''
import triton
import triton.language as tl
from triton.compiler.compiler import AttrsDescriptor

from torch._inductor.runtime import triton_helpers, triton_heuristics
from torch._inductor.runtime.triton_helpers import libdevice, math as tl_math
from torch._inductor.runtime.hints import AutotuneHint, ReductionHint, TileHint, DeviceProperties
triton_helpers.set_driver_to_gpu()

@triton_heuristics.reduction(
   size_hints=[2048, 32768],
   reduction_hint=ReductionHint.INNER,
   filename=__file__,
   triton_meta={'signature': {'in_out_ptr0': '*fp32', 'in_ptr0': '*fp32', 'in_ptr1': '*fp32', 'out_ptr0': '*fp32', 'xnumel': 'i32', 'rnumel': 'i32'}, 'device': DeviceProperties(type='cuda', index=0, cc=90, major=9, regs_per_multiprocessor=65536, max_threads_per_multi_processor=2048, multi_processor_count=132, warp_size=32), 'constants': {}, 'configs': [AttrsDescriptor(divisible_by_16=(0, 1, 2, 3, 4, 5), equal_to_1=())]},
   inductor_meta={'autotune_hints': set(), 'kernel_name': 'triton_red_fused_add_mean_mul_pow_rsqrt_0', 'mutated_arg_names': ['in_out_ptr0'], 'optimize_mem': False, 'no_x_dim': False, 'num_load': 3, 'num_reduction': 1, 'backend_hash': '1D985D9C0CD26A9BB5BAD8A051C971598DEF7D09C84A1EB5695128BF97C731E4', 'are_deterministic_algorithms_enabled': False, 'assert_indirect_indexing': True, 'autotune_local_cache': True, 'autotune_pointwise': True, 'autotune_remote_cache': None, 'force_disable_caches': False, 'dynamic_scale_rblock': True, 'max_autotune': False, 'max_autotune_pointwise': False, 'min_split_scan_rblock': 256, 'spill_threshold': 16, 'store_cubin': False}
)
@triton.jit
def triton_red_fused_add_mean_mul_pow_rsqrt_0(in_out_ptr0, in_ptr0, in_ptr1, out_ptr0, xnumel, rnumel, XBLOCK : tl.constexpr, RBLOCK : tl.constexpr):
   xnumel = 2048
   rnumel = 32768
   xoffset = tl.program_id(0) * XBLOCK
   xindex = xoffset + tl.arange(0, XBLOCK)[:, None]
   xmask = xindex < xnumel
   rbase = tl.arange(0, RBLOCK)[None, :]
   x0 = xindex
   _tmp3 = tl.full([XBLOCK, RBLOCK], 0, tl.float32)
   for roffset in range(0, rnumel, RBLOCK):
       rindex = roffset + rbase
       rmask = rindex < rnumel
       r1 = rindex
       tmp0 = tl.load(in_ptr0 + (r1 + (32768*x0)), rmask & xmask, eviction_policy='evict_last', other=0.0)
       tmp1 = tmp0 * tmp0
       tmp2 = tl.broadcast_to(tmp1, [XBLOCK, RBLOCK])
       tmp4 = _tmp3 + tmp2
       _tmp3 = tl.where(rmask & xmask, tmp4, _tmp3)
   tmp3 = tl.sum(_tmp3, 1)[:, None]
   tmp5 = 32768.0
   tmp6 = tmp3 / tmp5
   tmp7 = 1e-06
   tmp8 = tmp6 + tmp7
   tmp9 = libdevice.rsqrt(tmp8)
   tl.debug_barrier()
   tl.store(in_out_ptr0 + (x0), tmp9, xmask)
   for roffset in range(0, rnumel, RBLOCK):
       rindex = roffset + rbase
       rmask = rindex < rnumel
       r1 = rindex
       tmp10 = tl.load(in_ptr1 + (r1), rmask, eviction_policy='evict_last', other=0.0)
       tmp11 = tl.load(in_ptr0 + (r1 + (32768*x0)), rmask & xmask, eviction_policy='evict_first', other=0.0)
       tmp12 = tmp11 * tmp9
       tmp13 = tmp10 * tmp12
       tl.store(out_ptr0 + (r1 + (32768*x0)), tmp13, rmask & xmask)
''', device_str='cuda')

```

    
Line tmp11 is a redundant load instruction because all data has been loaded in line tmp0 before. However, we can’t reuse the tmp0 directly because it is in a loop. It is the root cause of the bad performance compared with liger’s implementation. 
The following are the memory charts for inductor and liger kernels.

![image](https://github.com/user-attachments/assets/e5fdea2b-e532-4b83-8aae-e56e592fe0f4)

Inductor version

![image](https://github.com/user-attachments/assets/8f40f8a9-fb38-4fd1-b775-88e77bea579d)
Liger version

As you can see, the total memory load from Device Memory to L2 Cache is 506MB for inductor version while 268MB for liger version. 

We observed the same pattern in input id 1-4 too. This can be a serious problem even for other ops we didn’t test. 

The equivalent implementation in liger kernel is [here](https://github.com/linkedin/Liger-Kernel/blob/7e0f459149d298c84f162363cc6f1347494b80f2/src/liger_kernel/ops/rms_norm.py#L44). The equivalent line for loading input is [`X_row = tl.load(X_ptr + col_offsets, mask=mask, other=0)`](https://github.com/linkedin/Liger-Kernel/blob/7e0f459149d298c84f162363cc6f1347494b80f2/src/liger_kernel/ops/rms_norm.py#L76C5-L76C61) This X_row is reused and it doesn't cause redundant loads.

For both the first load in inductor and liger implementation, they are all unrolled into multiple LDG.E instructions. But the loop in inductor stops the data reusing. For the second redundant load in inductor implementation, this loop is not unrolled because of  the store instruction in the end of the loop I think. 

### For input id 0
For the input id 0, which is the smallest input, its inductor compiled code is the following. 

```
# kernel path: /tmp/torchinductor_yhao/d7/cd7zxtkrvvmyfsqgvi26o2eurcvl2bael2advqf5r6rruhxxl62f.py
# Topologically Sorted Source Nodes: [pow_1, variance, add, rsqrt, hidden_states_1, mul_1], Original ATen: [aten.pow, aten.mean, aten.add, aten.rsqrt, aten.mul]
# Source node to ATen node mapping:
#   add => add
#   hidden_states_1 => mul
#   mul_1 => mul_1
#   pow_1 => pow_1
#   rsqrt => rsqrt
#   variance => mean
# Graph fragment:
#   %pow_1 : [num_users=1] = call_function[target=torch.ops.aten.pow.Tensor_Scalar](args = (%primals_1, 2), kwargs = {})
#   %mean : [num_users=1] = call_function[target=torch.ops.aten.mean.dim](args = (%pow_1, [-1], True), kwargs = {})
#   %add : [num_users=1] = call_function[target=torch.ops.aten.add.Tensor](args = (%mean, 1e-06), kwargs = {})
#   %rsqrt : [num_users=2] = call_function[target=torch.ops.aten.rsqrt.default](args = (%add,), kwargs = {})
#   %mul : [num_users=1] = call_function[target=torch.ops.aten.mul.Tensor](args = (%primals_1, %rsqrt), kwargs = {})
#   %mul_1 : [num_users=1] = call_function[target=torch.ops.aten.mul.Tensor](args = (%primals_2, %mul), kwargs = {})
triton_per_fused_add_mean_mul_pow_rsqrt_0 = async_compile.triton('triton_per_fused_add_mean_mul_pow_rsqrt_0', '''
import triton
import triton.language as tl
from triton.compiler.compiler import AttrsDescriptor

from torch._inductor.runtime import triton_helpers, triton_heuristics
from torch._inductor.runtime.triton_helpers import libdevice, math as tl_math
from torch._inductor.runtime.hints import AutotuneHint, ReductionHint, TileHint, DeviceProperties
triton_helpers.set_driver_to_gpu()

@triton_heuristics.persistent_reduction(
   size_hints=[2048, 1024],
   reduction_hint=ReductionHint.INNER,
   filename=__file__,
   triton_meta={'signature': {'in_out_ptr0': '*fp32', 'in_ptr0': '*fp32', 'in_ptr1': '*fp32', 'out_ptr0': '*fp32', 'xnumel': 'i32', 'rnumel': 'i32'}, 'device': DeviceProperties(type='cuda', index=0, cc=90, major=9, regs_per_multiprocessor=65536, max_threads_per_multi_processor=2048, multi_processor_count=132, warp_size=32), 'constants': {}, 'configs': [AttrsDescriptor(divisible_by_16=(0, 1, 2, 3, 4, 5), equal_to_1=())]},
   inductor_meta={'autotune_hints': set(), 'kernel_name': 'triton_per_fused_add_mean_mul_pow_rsqrt_0', 'mutated_arg_names': ['in_out_ptr0'], 'optimize_mem': False, 'no_x_dim': True, 'num_load': 2, 'num_reduction': 1, 'backend_hash': '1D985D9C0CD26A9BB5BAD8A051C971598DEF7D09C84A1EB5695128BF97C731E4', 'are_deterministic_algorithms_enabled': False, 'assert_indirect_indexing': True, 'autotune_local_cache': True, 'autotune_pointwise': True, 'autotune_remote_cache': None, 'force_disable_caches': False, 'dynamic_scale_rblock': True, 'max_autotune': False, 'max_autotune_pointwise': False, 'min_split_scan_rblock': 256, 'spill_threshold': 16, 'store_cubin': False}
)
@triton.jit
def triton_per_fused_add_mean_mul_pow_rsqrt_0(in_out_ptr0, in_ptr0, in_ptr1, out_ptr0, xnumel, rnumel):
   xnumel = 2048
   XBLOCK: tl.constexpr = 1
   rnumel = 1024
   RBLOCK: tl.constexpr = 1024
   xoffset = tl.program_id(0) * XBLOCK
   xindex = tl.full([1], xoffset, tl.int32)
   xmask = tl.full([RBLOCK], True, tl.int1)
   rindex = tl.arange(0, RBLOCK)[:]
   roffset = 0
   rmask = tl.full([RBLOCK], True, tl.int1)
   r1 = rindex
   x0 = xindex
   tmp0 = tl.load(in_ptr0 + (r1 + (1024*x0)), None)
   tmp10 = tl.load(in_ptr1 + (r1), None, eviction_policy='evict_last')
   tmp1 = tmp0 * tmp0
   tmp2 = tl.broadcast_to(tmp1, [RBLOCK])
   tmp4 = triton_helpers.promote_to_tensor(tl.sum(tmp2, 0))
   tmp5 = 1024.0
   tmp6 = tmp4 / tmp5
   tmp7 = 1e-06
   tmp8 = tmp6 + tmp7
   tmp9 = libdevice.rsqrt(tmp8)
   tmp11 = tmp0 * tmp9
   tmp12 = tmp10 * tmp11
   tl.debug_barrier()
   tl.store(in_out_ptr0 + (x0), tmp9, None)
   tl.store(out_ptr0 + (r1 + (1024*x0)), tmp12, None)
''', device_str='cuda')


```

The tmp0 is to load all inputs and no redundant loads when we do the computation for tmp11. However, its performance is still slightly worse than liger version. 

For line tmp0, it is compiled to a bunch of SASS instructions(CUDA binary assembly code) with the most important one LDG(load from global memory) instruction\`LDG.E.128 R8, desc\[UR6\]\[R8.64\]\`. However, for the line in liger version with same semantics \`    X\_row \= tl.load(X\_ptr \+ col\_offsets, mask=mask, other=0)\`, it is compiled to the following two LDG instructions.

```
@!P2  LDG.E.128 R8, desc[UR6][R22.64]
@!P1  LDG.E.128 R12, desc[UR6][R22.64+0x800]
```

I think these two separate instructions increase the flexibility of instruction scheduling and increase the instruction overlaps. This is the most suspicious one I can find for input0.




### Versions

main branch

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @ColinPeppler @amjames @desertfire @chauhang @aakhundov @eellison @shunting314 @Chillee @jansel