Repository: pytorch
Issue: Increasing batch size makes network forward 1000 times slower
Labels: module: cudnn,module: nn,module: cuda,triaged

### 🐛 Describe the bug

I have a two layer network. The input is a 2D array of token ids, first layer is an embedding layer that replaces each pixel the respective embedding, the second layer does a convolution. Find the network definition below.
```python
class Model(nn.Module):
    def __init__(self, voc_dim, emb_dim, out_channels):
        super().__init__()
        self.embedding = nn.Embedding(
                num_embeddings=voc_dim,
                embedding_dim=emb_dim,
                padding_idx=0,
                sparse=True,
            )
        self.conv = nn.Conv2d(
                in_channels=emb_dim,
                out_channels=out_channels,
                kernel_size=3,
                stride=1,
                dilation=1,
                padding=1,
                bias=False,
            )
    def forward(self, x):
        # x is (batch_size, 672, 512), embeddings (batch_size, 768, 672, 512)
        embeddings = self.embedding(x).permute(0, 3, 1, 2)
        out = self.conv(embeddings)
        return out

network = Model(voc_dim=100000, emb_dim=768, out_channels=64)
network.cuda()
network.train()
```
For batch sizes of 4 and 8, inference with mixed precision, forward pass runs smoothly. But for batch size 16, forward becomes ~1000 times slower.  Batch size 16 does not fit in the memory without mixed precision. Forward pass is timed with the following code.
```python
def forward_speed(batch_size, n=5):
    forward_pass_timed = []
    for i in range(n):
        start = torch.cuda.Event(enable_timing=True)
        end = torch.cuda.Event(enable_timing=True)
        inp = torch.rand(batch_size, 672, 512).long().cuda()
        start.record()
        with torch.autocast(device_type="cuda", dtype=torch.float16):
            x = network(inp)
        end.record()
        torch.cuda.synchronize()
        forward_pass_timed.append(start.elapsed_time(end))
    print(f'Batch size - {batch_size}, Mean of {n} runs in ms {round(np.mean(forward_pass_timed), 2)}, STD {round(np.std(forward_pass_timed), 2)}')
```
The output of the function for batch sizes 4,8,16.
<img width="604" alt="Screenshot 2023-08-02 at 23 42 12" src="https://github.com/pytorch/pytorch/assets/127740186/c894d48b-1cee-408a-b85b-40878e71fd57">

Interestingly, when I measure just one layer i.e only embedding or only convolution, they are not slow, this effect is seen only when they are called together in a forward pass. GPU is A100 40GB.

### Versions

PyTorch version: 2.0.0+cu117
Is debug build: False
CUDA used to build PyTorch: 11.7
ROCM used to build PyTorch: N/A

OS: Red Hat Enterprise Linux 8.4 (Ootpa) (x86_64)
GCC version: (GCC) 8.5.0 20210514 (Red Hat 8.5.0-18)
Clang version: Could not collect
CMake version: version 3.26.3
Libc version: glibc-2.28

Python version: 3.10.10 (main, Mar 21 2023, 18:45:11) [GCC 11.2.0] (64-bit runtime)
Python platform: Linux-4.18.0-305.19.1.el8_4.x86_64-x86_64-with-glibc2.28
Is CUDA available: True
CUDA runtime version: 11.0.221
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA A100-PCIE-40GB
Nvidia driver version: 495.29.05
cuDNN version: Could not collect
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:           AuthenticAMD
CPU family:          25
Model:               1
Model name:          AMD EPYC 7543 32-Core Processor
Stepping:            1
CPU MHz:             1638.617
CPU max MHz:         2800.0000
CPU min MHz:         1500.0000
BogoMIPS:            5589.65
Virtualization:      AMD-V
L1d cache:           32K
L1i cache:           32K
L2 cache:            512K
L3 cache:            32768K
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 mmx fxsr sse sse2 ht syscall nx mmxext fxsr_opt pdpe1gb rdtscp lm constant_tsc rep_good nopl nonstop_tsc cpuid extd_apicid aperfmperf pni pclmulqdq monitor ssse3 fma cx16 pcid sse4_1 sse4_2 movbe popcnt aes xsave avx f16c rdrand lahf_lm cmp_legacy svm extapic cr8_legacy abm sse4a misalignsse 3dnowprefetch osvw ibs skinit wdt tce topoext perfctr_core perfctr_nb bpext perfctr_llc mwaitx cpb cat_l3 cdp_l3 invpcid_single hw_pstate sme ssbd mba sev ibrs ibpb stibp vmmcall sev_es fsgsbase bmi1 avx2 smep bmi2 invpcid cqm rdt_a rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 xsaves cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local clzero irperf xsaveerptr wbnoinvd amd_ppin arat npt lbrv svm_lock nrip_save tsc_scale vmcb_clean flushbyasid decodeassists pausefilter pfthreshold v_vmsave_vmload vgif umip pku ospke vaes vpclmulqdq rdpid overflow_recov succor smca

Versions of relevant libraries:
[pip3] numpy==1.23.5
[pip3] torch==2.0.0
[pip3] torchvision==0.15.2
[pip3] triton==2.0.0
[conda] numpy                     1.23.5                   pypi_0    pypi
[conda] torch                     2.0.0                    pypi_0    pypi
[conda] torchvision               0.15.2                   pypi_0    pypi
[conda] triton                    2.0.0                    pypi_0    pypi

cc @csarofeen @ptrblck @xwang233 @albanD @mruberry @jbschlosser @walterddr @mikaylagawarecki