Repository: pytorch
Issue: Performance for `TransformerEncoderLayer` and `TransformerDecoderLayer` drops severely if pytorch version changed from `pt1.10` to `pt2.0`
Labels: module: performance,module: nn,triaged

### 🐛 Describe the bug

while testing the performance for transformer blocks, we found that the performance severely drops while updating pytorch version from `pt1.10v1` to `pt2.0v2_cuda11.8`。
the performance for `TransformerEncoderLayer` and `TransformerDecoderLayer` on `pt1.10v1` is the following
```bash
Time per iteration of encoder: 4.612 ms
Time per iteration of decoder: 5.848 ms
```
while for `pt2.0v2_cuda11.8`, the results drop dramastically.
```bash
Time per iteration of encoder: 9.414 ms
Time per iteration of decoder: 13.649 ms
```



Reproducing using the following benchmark codes:
```python
import torch
import torch.nn as nn
import time


def test_transformer_encoder_block(d_model, n_head, ffn_hidden, batch_size, seq_len, num_iterations):
    transformer_encoder_block = nn.TransformerEncoderLayer(d_model=d_model, nhead=n_head, dim_feedforward=ffn_hidden, dropout=0.1, batch_first=True).cuda()
    input_tensor = torch.rand(batch_size, seq_len, d_model).cuda()
    start_time = time.time()
    for i in range(num_iterations):
        output_tensor = transformer_encoder_block(input_tensor)
    end_time = time.time()
    print("Time per iteration of encoder: {:.3f} ms".format((end_time - start_time) / num_iterations * 1000))


def test_transformer_decoder_block(d_model, n_head, ffn_hidden, batch_size, tgt_len, memory_len, num_iterations):
    transformer_decoder_block = nn.TransformerDecoderLayer(d_model=d_model, nhead=n_head, batch_first=True).cuda()
    tgt_tensor = torch.rand(batch_size, tgt_len, d_model).cuda()
    memory_tensor = torch.rand(batch_size, memory_len, d_model).cuda()
    start_time = time.time()
    for i in range(num_iterations):
        output_tensor = transformer_decoder_block(tgt_tensor, memory_tensor)
    end_time = time.time()
    print("Time per iteration of decoder: {:.3f} ms".format((end_time - start_time) / num_iterations * 1000))

# test
test_transformer_encoder_block(d_model=512, n_head=8, ffn_hidden=2048, batch_size=32, seq_len=512, num_iterations=1000)
test_transformer_decoder_block(d_model=512, n_head=8, ffn_hidden=2048, batch_size=32, tgt_len=512, memory_len=512, num_iterations=1000)
```

some perf results(top 10; cuda time sorted) on`transformer_encoder_block` and  `transformer_decoder_block` from the above code snippet are shown below,

for pt1.10, 
```
-------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  
                                                   Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg     Self CUDA   Self CUDA %    CUDA total  CUDA time avg    # of Calls  
-------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  
                                               aten::mm         2.77%     170.550ms        30.93%        1.907s     317.758us     927.361ms        24.99%     927.361ms     154.560us          6000  
                                   aten::_fused_dropout         2.82%     174.105ms         6.18%     380.759ms      95.190us     681.397ms        18.36%     681.397ms     170.349us          4000  
void at::native::(anonymous namespace)::fused_dropou...         0.00%       0.000us         0.00%       0.000us       0.000us     681.397ms        18.36%     681.397ms     170.349us          4000  
void cutlass::Kernel<cutlass_80_tensorop_s1688gemm_1...         0.00%       0.000us         0.00%       0.000us       0.000us     645.084ms        17.38%     645.084ms     129.017us          5000  
                                              aten::bmm         2.01%     123.741ms         2.42%     149.069ms      74.534us     468.478ms        12.63%     468.478ms     234.239us          2000  
                                aten::native_layer_norm         3.03%     186.898ms         3.78%     233.016ms     116.508us     445.250ms        12.00%     445.250ms     222.625us          2000  
void at::native::unrolled_elementwise_kernel<at::nat...         0.00%       0.000us         0.00%       0.000us       0.000us     431.756ms        11.64%     431.756ms      61.679us          7000  
                                             aten::add_         1.01%      62.239ms         1.58%      97.348ms      16.225us     366.450ms         9.88%     366.450ms      61.075us          6000  
void at::native::(anonymous namespace)::RowwiseMomen...         0.00%       0.000us         0.00%       0.000us       0.000us     363.030ms         9.78%     363.030ms     181.515us          2000  
                                         aten::_softmax         0.21%      13.244ms         0.33%      20.213ms      20.213us     312.455ms         8.42%     312.455ms     312.455us          1000  
-------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  
Self CPU time total: 6.164s
Self CUDA time total: 3.711s

Time per iteration of encoder: 29.544 ms
-------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  
                                                   Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg     Self CUDA   Self CUDA %    CUDA total  CUDA time avg    # of Calls  
-------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  
                                               aten::mm         4.08%     227.057ms         5.32%     296.055ms      29.605us        1.248s        21.14%        1.248s     124.819us         10000  
                                   aten::_fused_dropout         4.79%     266.467ms        10.42%     579.759ms      96.626us        1.116s        18.91%        1.116s     186.039us          6000  
void at::native::(anonymous namespace)::fused_dropou...         0.00%       0.000us         0.00%       0.000us       0.000us        1.116s        18.91%        1.116s     186.039us          6000  
void cutlass::Kernel<cutlass_80_tensorop_s1688gemm_1...         0.00%       0.000us         0.00%       0.000us       0.000us     969.153ms        16.42%     969.153ms     107.684us          9000  
                                              aten::bmm         4.31%     239.975ms         5.26%     292.485ms      73.121us     937.281ms        15.88%     937.281ms     234.320us          4000  
                                aten::native_layer_norm         4.04%     225.016ms         5.27%     293.339ms      97.780us     660.988ms        11.20%     660.988ms     220.329us          3000  
void at::native::unrolled_elementwise_kernel<at::nat...         0.00%       0.000us         0.00%       0.000us       0.000us     652.791ms        11.06%     652.791ms      54.399us         12000  
                                         aten::_softmax         0.47%      26.312ms         0.71%      39.339ms      19.669us     622.398ms        10.54%     622.398ms     311.199us          2000  
void (anonymous namespace)::softmax_warp_forward<flo...         0.00%       0.000us         0.00%       0.000us       0.000us     622.398ms        10.54%     622.398ms     311.199us          2000  
void at::native::(anonymous namespace)::RowwiseMomen...         0.00%       0.000us         0.00%       0.000us       0.000us     538.506ms         9.12%     538.506ms     179.502us          3000  
-------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  
Self CPU time total: 5.565s
Self CUDA time total: 5.904s

Time per iteration of decoder: 48.706 ms
```


for pt2.0,

```
-------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  
                                                   Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg     Self CUDA   Self CUDA %    CUDA total  CUDA time avg    # of Calls  
-------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  
                                            aten::addmm         1.01%      93.315ms        18.35%        1.701s     566.910us        4.336s        47.79%        4.501s       1.500ms          3000  
                                 ampere_sgemm_128x64_tn         0.00%       0.000us         0.00%       0.000us       0.000us        3.239s        35.69%        3.239s       1.619ms          2000  
                                 ampere_sgemm_128x32_tn         0.00%       0.000us         0.00%       0.000us       0.000us        1.970s        21.72%        1.970s       1.970ms          1000  
                                              aten::bmm         0.38%      35.092ms        20.38%        1.890s     944.766us        1.491s        16.44%        1.548s     774.210us          2000  
                                               aten::mm         0.64%      59.369ms         1.03%      95.178ms      95.178us        1.399s        15.42%        1.429s       1.429ms          1000  
                                ampere_sgemm_128x128_nn         0.00%       0.000us         0.00%       0.000us       0.000us     966.682ms        10.65%     966.682ms     966.682us          1000  
                                   aten::native_dropout         0.72%      66.984ms        14.44%        1.339s     334.644us     689.672ms         7.60%     911.642ms     227.911us          4000  
void at::native::(anonymous namespace)::fused_dropou...         0.00%       0.000us         0.00%       0.000us       0.000us     689.672ms         7.60%     689.672ms     172.418us          4000  
                                       cudaLaunchKernel        83.84%        7.771s        83.84%        7.771s     353.247us     616.177ms         6.79%     616.177ms      28.008us         22000  
                                 ampere_sgemm_32x128_tn         0.00%       0.000us         0.00%       0.000us       0.000us     525.755ms         5.79%     525.755ms     525.755us          1000  
-------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  
Self CPU time total: 9.270s
Self CUDA time total: 9.073s

Time per iteration of encoder: 21.734 ms
-------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  
                                                   Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg     Self CUDA   Self CUDA %    CUDA total  CUDA time avg    # of Calls  
-------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  
                                            aten::addmm         0.88%     120.506ms         9.52%        1.305s     326.309us        4.862s        35.12%        4.948s       1.237ms          4000  
                                 ampere_sgemm_128x64_tn         0.00%       0.000us         0.00%       0.000us       0.000us        4.151s        29.99%        4.151s       1.384ms          3000  
                                              aten::bmm         0.47%      64.278ms        10.83%        1.484s     370.974us        2.977s        21.50%        3.019s     754.725us          4000  
                                               aten::mm         0.41%      55.889ms         6.99%     957.567ms     319.189us        2.821s        20.38%        2.861s     953.504us          3000  
                                 ampere_sgemm_128x32_tn         0.00%       0.000us         0.00%       0.000us       0.000us        1.973s        14.25%        1.973s       1.973ms          1000  
                                ampere_sgemm_128x128_nn         0.00%       0.000us         0.00%       0.000us       0.000us        1.929s        13.94%        1.929s     964.737us          2000  
                                 ampere_sgemm_32x128_tn         0.00%       0.000us         0.00%       0.000us       0.000us        1.560s        11.27%        1.560s     519.968us          3000  
                                   aten::native_dropout         0.66%      90.220ms         6.30%     862.789ms     143.798us        1.130s         8.16%        1.283s     213.773us          6000  
void at::native::(anonymous namespace)::fused_dropou...         0.00%       0.000us         0.00%       0.000us       0.000us        1.130s         8.16%        1.130s     188.261us          6000  
                                ampere_sgemm_128x128_tn         0.00%       0.000us         0.00%       0.000us       0.000us        1.047s         7.57%        1.047s     523.683us          2000  
-------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  
Self CPU time total: 13.704s
Self CUDA time total: 13.844s

Time per iteration of decoder: 38.392 ms

```
all mainly executed ops seems suffer performance degrade, which is weried.



btw, this issue also occurs on `pt1.12` according to our experiments.




### Versions

pt1.10:

```
Collecting environment information...
PyTorch version: 1.10.0a0+git36449ea
Is debug build: False
CUDA used to build PyTorch: 11.0
ROCM used to build PyTorch: N/A

OS: CentOS Linux 7 (Core) (x86_64)
GCC version: (GCC) 7.5.0
Clang version: Could not collect
CMake version: version 3.22.1
Libc version: glibc-2.10

Python version: 3.7.3 (default, Mar 27 2019, 22:11:17)  [GCC 7.3.0] (64-bit runtime)
Python platform: Linux-3.10.0-957.el7.x86_64-x86_64-with-centos-7.6.1810-Core
Is CUDA available: True
CUDA runtime version: 11.0.221
CUDA_MODULE_LOADING set to: 
GPU models and configuration: 
GPU 0: NVIDIA A100-SXM4-80GB
GPU 1: NVIDIA A100-SXM4-80GB
GPU 2: NVIDIA A100-SXM4-80GB
GPU 3: NVIDIA A100-SXM4-80GB
GPU 4: NVIDIA A100-SXM4-80GB
GPU 5: NVIDIA A100-SXM4-80GB
GPU 6: NVIDIA A100-SXM4-80GB
GPU 7: NVIDIA A100-SXM4-80GB

Nvidia driver version: 525.85.12
cuDNN version: Probably one of the following:
/usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn.so.8.1.1
/usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8.1.1
/usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn_adv_train.so.8.1.1
/usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8.1.1
/usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8.1.1
/usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8.1.1
/usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn_ops_train.so.8.1.1
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
Byte Order:            Little Endian
CPU(s):                128
On-line CPU(s) list:   0-127
Thread(s) per core:    2
Core(s) per socket:    32
Socket(s):             2
NUMA node(s):          2
Vendor ID:             GenuineIntel
CPU family:            6
Model:                 106
Model name:            Intel(R) Xeon(R) Platinum 8358P CPU @ 2.60GHz
Stepping:              6
CPU MHz:               3400.000
CPU max MHz:           3400.0000
CPU min MHz:           800.0000
BogoMIPS:              5200.00
Virtualization:        VT-x
L1d cache:             48K
L1i cache:             32K
L2 cache:              1280K
L3 cache:              49152K
NUMA node0 CPU(s):     0-31,64-95
NUMA node1 CPU(s):     32-63,96-127
Flags:                 fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush dts acpi mmx fxsr sse sse2 ss ht tm pbe syscall nx pdpe1gb rdtscp lm constant_tsc art arch_perfmon pebs bts rep_good nopl xtopology nonstop_tsc aperfmperf eagerfpu pni pclmulqdq dtes64 ds_cpl vmx smx est tm2 ssse3 sdbg fma cx16 xtpr pdcm pcid dca sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand lahf_lm abm 3dnowprefetch epb cat_l3 intel_pt ssbd mba ibrs ibpb stibp ibrs_enhanced tpr_shadow vnmi flexpriority ept vpid fsgsbase tsc_adjust bmi1 hle avx2 smep bmi2 erms invpcid rtm cqm rdt_a avx512f avx512dq rdseed adx smap avx512ifma clflushopt clwb avx512cd sha_ni avx512bw avx512vl xsaveopt xsavec xgetbv1 cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local dtherm ida arat pln pts avx512vbmi umip pku ospke avx512_vbmi2 gfni vaes vpclmulqdq avx512_vnni avx512_bitalg avx512_vpopcntdq spec_ctrl intel_stibp flush_l1d arch_capabilities

Versions of relevant libraries:
[pip3] flake8==5.0.4
[pip3] numpy==1.19.2
[pip3] onnx==1.6.0
[pip3] torch==1.10.0a0+git36449ea
[pip3] torchaudio==0.10.0
[pip3] torchvision==0.11.1
[conda] blas                      1.0                         mkl    https://repo.anaconda.com/pkgs/main
[conda] magma-cuda112             2.5.2                         1    pytorch
[conda] mkl                       2020.2                      256    https://repo.anaconda.com/pkgs/main
[conda] mkl                       2022.1.0                  <pip>
[conda] mkl-include               2022.1.0           h06a4308_224    https://repo.anaconda.com/pkgs/main
[conda] mkl-service               2.3.0            py37he8ac12f_0    https://repo.anaconda.com/pkgs/main
[conda] mkl_fft                   1.3.0            py37h54f3939_0    https://repo.anaconda.com/pkgs/main
[conda] mkl_random                1.1.1            py37h0573a6f_0    https://repo.anaconda.com/pkgs/main
[conda] numpy                     1.16.4                    <pip>
[conda] numpy                     1.19.2           py37h54aff64_0    https://repo.anaconda.com/pkgs/main
[conda] numpy-base                1.19.2           py37hfa32c7d_0    https://repo.anaconda.com/pkgs/main
[conda] torch                     1.10.0a0+git36449ea           <pip>
[conda] torchaudio                0.10.0                    <pip>
[conda] torchvision               0.11.1                    <pip>
```

pt2.0:
```
Collecting environment information...
PyTorch version: 2.0.0+cu118
Is debug build: False
CUDA used to build PyTorch: 11.8
ROCM used to build PyTorch: N/A

OS: CentOS Linux 7 (Core) (x86_64)
GCC version: (GCC) 10.2.0
Clang version: Could not collect
CMake version: version 3.26.3
Libc version: glibc-2.17

Python version: 3.10.9 (main, Jan 11 2023, 15:21:40) [GCC 11.2.0] (64-bit runtime)
Python platform: Linux-3.10.0-957.el7.x86_64-x86_64-with-glibc2.17
Is CUDA available: True
CUDA runtime version: 11.8.89
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: 
GPU 0: NVIDIA A100-SXM4-80GB
GPU 1: NVIDIA A100-SXM4-80GB
GPU 2: NVIDIA A100-SXM4-80GB
GPU 3: NVIDIA A100-SXM4-80GB
GPU 4: NVIDIA A100-SXM4-80GB
GPU 5: NVIDIA A100-SXM4-80GB
GPU 6: NVIDIA A100-SXM4-80GB
GPU 7: NVIDIA A100-SXM4-80GB

Nvidia driver version: 525.85.12
cuDNN version: Probably one of the following:
/usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn.so.8.1.1
/usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8.1.1
/usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn_adv_train.so.8.1.1
/usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8.1.1
/usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8.1.1
/usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8.1.1
/usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn_ops_train.so.8.1.1
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
Byte Order:            Little Endian
CPU(s):                128
On-line CPU(s) list:   0-127
Thread(s) per core:    2
Core(s) per socket:    32
Socket(s):             2
NUMA node(s):          2
Vendor ID:             GenuineIntel
CPU family:            6
Model:                 106
Model name:            Intel(R) Xeon(R) Platinum 8358P CPU @ 2.60GHz
Stepping:              6
CPU MHz:               3400.000
CPU max MHz:           3400.0000
CPU min MHz:           800.0000
BogoMIPS:              5200.00
Virtualization:        VT-x
L1d cache:             48K
L1i cache:             32K
L2 cache:              1280K
L3 cache:              49152K
NUMA node0 CPU(s):     0-31,64-95
NUMA node1 CPU(s):     32-63,96-127
Flags:                 fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush dts acpi mmx fxsr sse sse2 ss ht tm pbe syscall nx pdpe1gb rdtscp lm constant_tsc art arch_perfmon pebs bts rep_good nopl xtopology nonstop_tsc aperfmperf eagerfpu pni pclmulqdq dtes64 ds_cpl vmx smx est tm2 ssse3 sdbg fma cx16 xtpr pdcm pcid dca sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand lahf_lm abm 3dnowprefetch epb cat_l3 intel_pt ssbd mba ibrs ibpb stibp ibrs_enhanced tpr_shadow vnmi flexpriority ept vpid fsgsbase tsc_adjust bmi1 hle avx2 smep bmi2 erms invpcid rtm cqm rdt_a avx512f avx512dq rdseed adx smap avx512ifma clflushopt clwb avx512cd sha_ni avx512bw avx512vl xsaveopt xsavec xgetbv1 cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local dtherm ida arat pln pts avx512vbmi umip pku ospke avx512_vbmi2 gfni vaes vpclmulqdq avx512_vnni avx512_bitalg avx512_vpopcntdq spec_ctrl intel_stibp flush_l1d arch_capabilities

Versions of relevant libraries:
[pip3] flake8==6.1.0
[pip3] numpy==1.26.2
[pip3] onnx==1.15.0
[pip3] onnxoptimizer==0.3.13
[pip3] onnxruntime==1.16.3
[pip3] torch==2.0.0+cu118
[pip3] torchvision==0.15.1a0
[pip3] triton==2.1.0
[conda] mkl                       2024.0.0                 pypi_0    pypi
[conda] mkl-include               2024.0.0                 pypi_0    pypi
[conda] numpy                     1.26.2                   pypi_0    pypi
[conda] torch                     2.0.0+cu118              pypi_0    pypi
[conda] torchvision               0.15.1a0                 pypi_0    pypi
[conda] triton                    2.1.0                    pypi_0    pypi

```

@albanD @jbschlosser @mikaylagawarecki @ezyang 


cc @msaroufim @albanD @mruberry @jbschlosser @walterddr @mikaylagawarecki @bhosmer @cpuhrsch @erichan1 @drisspg