Repository: pytorch
Issue: very strange speed of torch.bmm with specific tensor shape
Labels: module: performance,module: cuda,triaged,module: linear algebra

### 🐛 Describe the bug

I test the below code in a A100 GPU.
It outputs
avg_time: 1.97
avg_time: 0.115
avg_time: 2.09

The speed is much faster when the row number is 256.
I know Cuda performance depends on tensor shape, but it should not have so much gap.


```python
import torch
import time
torch.backends.cudnn.benchmark = True

num_iter = 1000

def run(a, b, half=True):
    a = a.cuda()
    b = b.cuda()
    if half:
        a = a.half()
        b = b.half()
    
    # cudnn warmup
    for _ in range(50):
        _ = torch.bmm(a, b)
    
    torch.cuda.synchronize()
    t0 = time.time()
    
    for _ in range(num_iter):
        _ = torch.bmm(a, b)
    
    torch.cuda.synchronize()
    t1 = time.time()
    
    avg_time = (t1 - t0) / num_iter * 1000
    print(f'avg_time: {avg_time:.3g}')
    return avg_time

if __name__ == '__main__':
    print('Pytorch version\t:', torch.__version__)
    print('CUDA version\t:', torch.version.cuda)

    bs = 128
    is_half = True


    n = 255
    inp0 = torch.randn([bs, n, 1024])
    inp1 = torch.randn([bs, 1024, n])
    run(inp0, inp1, is_half)

    n = 256
    inp0 = torch.randn([bs, n, 1024])
    inp1 = torch.randn([bs, 1024, n])
    run(inp0, inp1, is_half)
    
    n = 257
    inp0 = torch.randn([bs, n, 1024])
    inp1 = torch.randn([bs, 1024, n])
    run(inp0, inp1, is_half)
```

### Versions

Collecting environment information...
PyTorch version: 1.12.1
Is debug build: False
CUDA used to build PyTorch: 11.3
ROCM used to build PyTorch: N/A

OS: CentOS Linux 7 (Core) (x86_64)
GCC version: (GCC) 7.3.0
Clang version: Could not collect
CMake version: version 2.8.12.2
Libc version: glibc-2.17

Python version: 3.10.4 (main, Mar 31 2022, 08:41:55) [GCC 7.5.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.2.152
GPU models and configuration: 
GPU 0: A100-SXM-80GB
GPU 1: A100-SXM-80GB
GPU 2: A100-SXM-80GB
GPU 3: A100-SXM-80GB
GPU 4: A100-SXM-80GB
GPU 5: A100-SXM-80GB
GPU 6: A100-SXM-80GB
GPU 7: A100-SXM-80GB

Nvidia driver version: 460.32.03
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

Versions of relevant libraries:
[pip3] numpy==1.23.1
[pip3] torch==1.12.1
[pip3] torch-tb-profiler==0.4.0
[pip3] torchaudio==0.12.1
[pip3] torchsummary==1.5.1
[pip3] torchvision==0.13.1
[conda] blas                      1.0                         mkl    defaults
[conda] cudatoolkit               11.3.1               h2bc3f7f_2    defaults
[conda] ffmpeg                    4.3                  hf484d3e_0    pytorch
[conda] mkl                       2021.4.0           h06a4308_640    defaults
[conda] mkl-service               2.4.0           py310h7f8727e_0    defaults
[conda] mkl_fft                   1.3.1           py310hd6ae3a3_0    defaults
[conda] mkl_random                1.2.2           py310h00e6091_0    defaults
[conda] numpy                     1.23.1          py310h1794996_0    defaults
[conda] numpy-base                1.23.1          py310hcba007f_0    defaults
[conda] pytorch                   1.12.1          py3.10_cuda11.3_cudnn8.3.2_0    pytorch
[conda] pytorch-mutex             1.0                        cuda    pytorch
[conda] torch-tb-profiler         0.4.0                    pypi_0    pypi
[conda] torchaudio                0.12.1              py310_cu113    pytorch
[conda] torchsummary              1.5.1                    pypi_0    pypi
[conda] torchvision               0.13.1              py310_cu113    pytorch

cc @VitalyFedyunin @ngimel @jianyuh @nikitaved @pearu @mruberry @walterddr @IvanYashchuk @xwang233 @Lezcano