Repository: pytorch
Issue: Inductor vs. Liger Performance Track
Labels: triaged,oncall: pt2,module: inductor

### 🐛 Describe the bug

Recently, we did some benchmarking on custom operators in liger kernels compared with inductor compiled kernels. Inductor is worse on some  cases. Here is the operator and config list we need to improve. 
## List

### Format

For each operator, the data format in the following task list is:

- [ ] **Operator Name**, (20th percentile speedup, 50th percentile (median), 80th percentile)

### Speedup Calculation

The speedup numbers are computed as follows:

$$
\text{inductor\\_vs\\_liger} = \frac{\text{speedup\\_inductor}}{\text{speedup\\_liger}} = \frac{\frac{\text{latency\\_eager}}{\text{latency\\_inductor}}}{\frac{\text{latency\\_eager}}{\text{latency\\_liger}}} = \frac{\text{latency\\_liger}}{\text{latency\\_inductor}}
$$


Since each operator has multiple inputs, there are multiple speedup numbers. We use the 20th, 50th, and 80th percentiles to better represent the results. If the number is >1, it means inductor's results are faster. The GPU peak memory usage is determined using the same process.

We need to improve inductor's performance on cases that less than 1. 

fwd_bwd fp32 latency:
 - [x] cross_entropy: 1.37, 1.50, 1.60
 - [ ] embedding: 0.95, 1.43, 1.86
 - [x] fused_linear_cross_entropy: 1.11, 1.14, 1.24
 - [x] fused_linear_jsd: 1.42, 1.68, 2.32
 - [x] geglu: 1.00, 1.00, 1.00
 - [x] jsd: 1.18, 3.38, 3.67
 - [x] kl_div: 1.18, 1.21, 1.29
 - [ ] rms_norm: 0.90, 0.99, 1.29
 - [ ] rope: 0.41, 0.48, 0.58 #141265
 - [x] swiglu: 0.99, 1.00, 1.00

fwd_bwd fp32 peak gpu memory usage:
 - [ ] cross_entropy: 0.75, 0.75, 0.75
 - [ ] embedding: 0.77, 0.84, 0.93
 - [ ] fused_linear_cross_entropy: 0.49, 0.64, 0.73
 - [ ] fused_linear_jsd: 0.78, 0.79, 0.81
 - [ ] geglu: 0.53, 0.72, 0.91
 - [x] jsd: 1.09, 1.09, 1.09
 - [x] kl_div: 1.00, 1.00, 1.00
 - [x] rms_norm: 1.01, 1.01, 1.01
 - [ ] rope: 0.53, 0.54, 0.59
 - [ ] swiglu: 0.55, 0.74, 0.91

fwd fp32 latency:
 - [x] cross_entropy: 1.71, 2.28, 2.28
 - [ ] embedding, (0.83, 0.95, 1.00) https://github.com/pytorch/pytorch/issues/142250
 - [x] fused_linear_cross_entropy: 3.07, 3.25, 3.58
 - [x] fused_linear_jsd: 2.90, 3.63, 5.18
 - [x] geglu: 0.99, 1.00, 1.02
 - [x] jsd: 2.96, 15.97, 17.37
 - [ ] kl_div: 0.90, 0.92, 0.97
 - [ ] rms_norm: 0.91, 0.96, 0.98 #141916
 - [ ] rope: 0.80, 0.90, 0.95
 - [x] swiglu: 1.00, 1.00, 1.00

fwd fp32 peak gpu memory usage:
 - [x] cross_entropy: 1.00, 1.00, 1.00
 - [x] embedding: 1.00, 1.00, 1.00
 - [ ] fused_linear_cross_entropy: 0.38, 0.62, 0.91
 - [x] fused_linear_jsd: 1.19, 1.40, 1.59
 - [ ] geglu: 0.36, 0.68, 1.00
 - [x] jsd: 1.67, 1.67, 1.67
 - [x] kl_div: 1.00, 1.00, 1.00
 - [x] rms_norm: 1.00, 1.00, 1.00
 - [ ] rope: 0.76, 0.78, 0.81
 - [ ] swiglu: 0.39, 0.70, 1.00

bwd fp32 latency:
 - [x] cross_entropy: 1.07, 1.10, 1.18
 - [x] embedding: 1.13, 1.88, 3.32
 - [ ] fused_linear_cross_entropy: 0.00, 0.01, 0.01
 - [ ] fused_linear_jsd: 0.01, 0.02, 0.03
 - [x] geglu: 1.00, 1.00, 1.00
 - [ ] jsd: 0.71, 0.72, 0.73
 - [x] kl_div: 1.25, 1.29, 1.36
 - [x] rms_norm: 1.36, 1.44, 1.66
 - [ ] rope: 0.36, 0.42, 0.44
 - [x] swiglu: 0.99, 0.99, 0.99

bwd fp32 peak gpu memory usage:
 - [ ] cross_entropy: 0.75, 0.75, 0.75
 - [ ] embedding: 0.76, 0.82, 0.92
 - [ ] fused_linear_cross_entropy: 0.51, 0.59, 0.69
 - [ ] fused_linear_jsd: 0.64, 0.67, 0.69
 - [ ] geglu: 0.86, 0.87, 0.89
 - [x] jsd: 1.00, 1.00, 1.00
 - [x] kl_div: 1.00, 1.00, 1.00
 - [x] rms_norm: 1.01, 1.01, 1.01
 - [ ] rope: 0.60, 0.62, 0.66
 - [ ] swiglu: 0.87, 0.89, 0.90

fwd_bwd bf16 latency:
 - [ ] cross_entropy: 0.94, 0.99, 1.07
 - [x] fused_linear_cross_entropy: 1.56, 1.95, 2.78
 - [x] fused_linear_jsd: 6.08, 8.14, 11.68
 - [x] geglu: 1.00, 1.01, 1.01
 - [x] jsd: 1.29, 3.82, 3.95
 - [x] kl_div: 1.01, 1.04, 1.12
 - [x] rms_norm: 0.98, 1.12, 1.87
 - [ ] rope: 0.31, 0.44, 0.53
 - [x] swiglu: 0.99, 1.00, 1.00

fwd_bwd bf16 peak gpu memory usage:
 - [ ] cross_entropy: 0.89, 0.89, 0.89
 - [ ] fused_linear_cross_entropy: 0.42, 0.54, 0.71
 - [ ] fused_linear_jsd: 0.63, 0.84, 0.85
 - [ ] geglu: 0.53, 0.72, 0.91
 - [x] jsd: 1.08, 1.08, 1.08
 - [x] kl_div: 1.00, 1.00, 1.00
 - [ ] rms_norm: 0.76, 0.76, 0.76
 - [ ] rope: 0.54, 0.56, 0.63
 - [ ] swiglu: 0.55, 0.74, 0.91

fwd bf16 latency:
 - [x] cross_entropy: 1.88, 2.43, 2.60
 - [ ] embedding: 0.21, 0.60, 0.98
 - [x] fused_linear_cross_entropy: 4.66, 5.84, 8.48
 - [x] fused_linear_jsd: 11.18, 17.52, 28.28
 - [x] geglu: 1.01, 1.02, 1.02
 - [x] jsd: 3.77, 28.18, 28.68
 - [ ] kl_div: 0.48, 0.53, 0.65
 - [ ] rms_norm: 0.71, 0.75, 0.76
 - [ ] rope: 0.60, 0.87, 0.92
 - [x] swiglu: 1.00, 1.00, 1.00

fwd bf16 peak gpu memory usage:
 - [x] cross_entropy: 1.00, 1.00, 1.00
 - [x] embedding: 1.00, 1.00, 1.00
 - [ ] fused_linear_cross_entropy: 0.40, 0.64, 0.92
 - [ ] fused_linear_jsd: 0.27, 1.00, 1.74
 - [ ] geglu: 0.36, 0.68, 1.00
 - [x] jsd: 1.37, 1.37, 1.37
 - [x] kl_div: 1.00, 1.00, 1.00
 - [ ] rms_norm: 0.67, 0.67, 0.67
 - [ ] rope: 0.77, 0.80, 0.85
 - [ ] swiglu: 0.39, 0.70, 1.00

bwd bf16 latency:
 - [ ] cross_entropy: 0.82, 0.84, 0.86
 - [ ] fused_linear_cross_entropy: 0.02, 0.04, 0.07
 - [ ] fused_linear_jsd: 0.05, 0.08, 0.16
 - [x] geglu: 0.99, 1.00, 1.00
 - [ ] jsd: 0.87, 0.87, 0.89
 - [x] kl_div: 1.12, 1.16, 1.22
 - [x] rms_norm: 1.13, 1.39, 1.62
 - [ ] rope: 0.32, 0.37, 0.39
 - [x] swiglu: 0.99, 0.99, 1.00

bwd bf16 peak gpu memory usage:
 - [ ] cross_entropy: 0.89, 0.89, 0.89
 - [ ] fused_linear_cross_entropy: 0.51, 0.59, 0.69
 - [ ] fused_linear_jsd: 0.68, 0.69, 0.70
 - [ ] geglu: 0.86, 0.87, 0.89
 - [x] jsd: 1.04, 1.04, 1.04
 - [x] kl_div: 1.00, 1.00, 1.00
 - [ ] rms_norm: 0.81, 0.81, 0.81
 - [ ] rope: 0.61, 0.64, 0.70
 - [ ] swiglu: 0.87, 0.89, 0.90

## Reproduce
Install tritonbench dependencies.
```
git clone https://github.com/pytorch-labs/tritonbench/
git submodule update --init --recursive
python install.py
```
Install liger kernels.
```
git clone https://github.com/linkedin/Liger-Kernel.git
cd Liger-Kernel
pip install -e .
# or if using transformers
pip install -e .[transformers]
```
or
```
pip install liger-kernel-nightly
```
Benchmark operators.
```
python run.py --op rope --mode fwd_bwd --precision fp32 --metrics latency,gpu_peak_mem,speedup,mem_footprint
```
Use `--op` to specify the operator to benchmark. Use `--mode` to specify the mode to
benchmark. Use `--precision` to specify the precision to benchmark. Use `--metrics`
to specify the metrics to benchmark. You can also add `--ncu_rep` to `--metrics` to save the nsight
compute profiling reports for this operator.

By default, all inputs will be benchmarked, but inductor is not slower on all inputs. You can use `--input-id` to specify the
index of the input to benchmark. The following is an exampled output.
```
% python run.py --op rope --mode fwd_bwd --precision fp32 --metrics latency,gpu_peak_mem,speedup,mem_footprint
  0%|                                                                                                                                                                                                                                                              | 0/8 [00:00<?, ?it/s]`LlamaRotaryEmbedding` can now be fully parameterized by passing the model config through the `config` argument. All other arguments will be removed in v4.46
100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 8/8 [00:08<00:00,  1.11s/it]
       (H, T)    apply_rotary_pos_emb-latency    apply_rotary_pos_emb-gpu_peak_mem    liger_rotary_pos_emb-latency    liger_rotary_pos_emb-speedup    liger_rotary_pos_emb-gpu_peak_mem    liger_rotary_pos_emb-mem_footprint    inductor_rotary_pos_emb_full_op-latency    inductor_rotary_pos_emb_full_op-speedup    inductor_rotary_pos_emb_full_op-gpu_peak_mem    inductor_rotary_pos_emb_full_op-mem_footprint
-------------  ------------------------------  -----------------------------------  ------------------------------  ------------------------------  -----------------------------------  ------------------------------------  -----------------------------------------  -----------------------------------------  ----------------------------------------------  -----------------------------------------------
 (8192, 1024)                        0.554368                             0.312484                        0.107328                         5.16518                             0.205537                               1.52033                                   0.151904                                    3.64946                                        0.364921                                         0.856306
 (8192, 2048)                        1.09846                              0.591413                        0.163776                         6.70711                             0.37752                                1.56657                                   0.297376                                    3.69386                                        0.696287                                         0.849381
 (8192, 4096)                        2.14995                              1.14927                         0.320192                         6.71457                             0.721486                               1.59292                                   0.529696                                    4.05884                                        1.35902                                          0.845662
 (8192, 8192)                        4.23626                              2.26499                         0.630336                         6.72063                             1.40942                                1.60704                                   1.13363                                     3.73689                                        2.68449                                          0.843733
(8192, 16384)                        8.51472                              4.49642                         1.25142                          6.80403                             2.78528                                1.61435                                   2.51034                                     3.39186                                        5.33542                                          0.84275
  (512, 2048)                        0.29776                              0.068436                        0.125696                         2.36889                             0.055083                               1.24242                                   0.226336                                    1.31557                                        0.075006                                         0.912407
 (2048, 2048)                        0.292704                             0.173031                        0.073984                         3.95631                             0.11957                                1.44711                                   0.188608                                    1.55192                                        0.199262                                         0.86836
 (8192, 2048)                        1.09587                              0.591413                        0.163776                         6.69129                             0.37752                                1.56657                                   0.264064                                    4.15002                                        0.696287                                         0.849381
 ```

### Versions

nightly

cc @ezyang @chauhang @penguinwu @voznesenskym @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @ColinPeppler @amjames @desertfire @aakhundov @eellison @shunting314 @Chillee @xuzhao9