Repository: pytorch
Issue: bmm operator in bfloat16 has low TFLOPS for some tensor shapes with CUDA 11.6
Labels: triaged,module: cublas

### 🐛 Describe the bug

bmm operator has much lower TFLOPS on tensor shapes A=(384, 1024, 1024) and B=(384, 1024, 140) in bfloat16. 
```
shape: (384, 1024, 1024, 140, torch.float16), TFLOP: 0.11, cost: 0.002092, TFLOPS : 53.89
shape: (512, 1024, 1024, 128, torch.float16), TFLOP: 0.14, cost: 0.001037, TFLOPS : 132.48
shape: (384, 1024, 1024, 140, torch.bfloat16), TFLOP: 0.11, cost: 0.009232, TFLOPS : 12.21
shape: (512, 1024, 1024, 128, torch.bfloat16), TFLOP: 0.14, cost: 0.001036, TFLOPS : 132.71
```
The reproducible script:
```
import time

import torch

def benchmark(n, k, m, l, dtype):
    warmup = 2
    number = 10

    a = torch.rand(n, k, m, dtype=dtype).cuda()
    b = torch.rand(n, m, l, dtype=dtype).cuda()
    for i in range(warmup):
        c = torch.bmm(a, b)

    torch.cuda.synchronize()
    tic = time.time()
    for i in range(number):
        c = torch.bmm(a, b)
    torch.cuda.synchronize()
    toc = time.time()

    total_flops = 2 * n * k * m * l
    cost = (toc - tic) / number
    shape = (n, k, m, l, dtype)

    print(f"shape: {shape}, TFLOP: {total_flops / 1e12:.2f}, "
          f"cost: {cost:3f}, "
          f"TFLOPS : {total_flops / cost / 1e12:.2f}""")

benchmark(384, 1024, 1024, 6720//48, torch.float16)
benchmark(512, 1024, 1024, 8192//64, torch.float16)
benchmark(384, 1024, 1024, 6720//48, torch.bfloat16)
benchmark(512, 1024, 1024, 8192//64, torch.bfloat16)
```

### Versions

Collecting environment information...
PyTorch version: 1.12.0a0+67ece03
Is debug build: False
CUDA used to build PyTorch: 11.6
ROCM used to build PyTorch: N/A

OS: Amazon Linux 2 (x86_64)
GCC version: (GCC) 7.3.1 20180712 (Red Hat 7.3.1-15)
Clang version: 11.1.0 (Amazon Linux 2 11.1.0-1.amzn2.0.2)
CMake version: version 3.18.2
Libc version: glibc-2.2.5

Python version: 3.7.10 (default, Jun  3 2021, 00:02:01)  [GCC 7.3.1 20180712 (Red Hat 7.3.1-13)] (64-bit runtime)
Python platform: Linux-4.14.276-211.499.amzn2.x86_64-x86_64-with-glibc2.2.5
Is CUDA available: True
CUDA runtime version: 11.6.55
GPU models and configuration: 
GPU 0: NVIDIA A100-SXM4-40GB
GPU 1: NVIDIA A100-SXM4-40GB
GPU 2: NVIDIA A100-SXM4-40GB
GPU 3: NVIDIA A100-SXM4-40GB
GPU 4: NVIDIA A100-SXM4-40GB
GPU 5: NVIDIA A100-SXM4-40GB
GPU 6: NVIDIA A100-SXM4-40GB
GPU 7: NVIDIA A100-SXM4-40GB

Nvidia driver version: 510.39.01
cuDNN version: Probably one of the following:
/usr/lib64/libcudnn.so.8.4.1
/usr/lib64/libcudnn_adv_infer.so.8.4.1
/usr/lib64/libcudnn_adv_train.so.8.4.1
/usr/lib64/libcudnn_cnn_infer.so.8.4.1
/usr/lib64/libcudnn_cnn_train.so.8.4.1
/usr/lib64/libcudnn_ops_infer.so.8.4.1
/usr/lib64/libcudnn_ops_train.so.8.4.1
/usr/local/cuda-10.1/targets/x86_64-linux/lib/libcudnn.so.7.6.5
/usr/local/cuda-10.2/targets/x86_64-linux/lib/libcudnn.so.7.6.5
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] numpy==1.21.6
[pip3] pytorch-ignite==0.4.2
[pip3] pytorch-lightning==1.6.5
[pip3] torch==1.12.0a0+gite67ece03
[pip3] torchaudio==0.10.1+6f539cf
[pip3] torchmetrics==0.9.3
[pip3] torchtext==0.10.0a0+0d670e0
[pip3] torchvision==0.8.1
[conda] Could not collect


cc @csarofeen @ptrblck @xwang233