Repository: pytorch
Issue: Performance regression in torch.nn.Linear
Labels: module: performance,module: nn,module: cuda,triaged,module: regression

### 🐛 Describe the bug

We identified a significant performance regression in `torch.nn.Linear` for tall-and-skinny inputs, especially on non-MIG A100 cards. Here's a simple reproducer that measures the run time of torch.nn.Linear's forward pass (which uses torch.functional.linear) against a simple matmul-and-add implementation, across a wide range of input sizes.

```python
import timeit
import torch


def main():
    device = torch.device("cuda:0")
    torch.manual_seed(0)
 
    lin = torch.nn.Linear(16, 16, device=device)
    xlist = [torch.randn(2**e, 16, device=device) for e in range(8, 24)]

 
    def run_lin(x):
        for _ in range(100):
            lin(x)
        torch.cuda.synchronize() 

    def run_mm(x):
        for _ in range(100):
            x @ lin.weight.T + lin.bias
        torch.cuda.synchronize()
 
    number = 100
    for x in xlist:
        t_lin, t_mm = [
            timeit.timeit(
                "run(x)",
                setup="torch.cuda.synchronize()",
                number=number,
                globals={"run": run, "x": x, "torch": torch},
            )
            / number
            for run in [run_lin, run_mm]
        ]

        print(
            f"{str(tuple(x.shape)):14s} "
            f"{t_lin * 1000:6.2f} ms, "
            f"{t_mm * 1000:6.2f} ms"
        )
 

if __name__ == "__main__":
    print(torch.__path__[0])
    print(torch.__version__)
    main()
```

The expectation here is that the nn.Linear calculation should run faster than matmul-and-add, because nn.Linear uses a fused kernel via torch.functional.linear. And this is indeed what happens under PyTorch 1.11.0+cu113. However, under PyTorch 2.1.2, the nn.Linear calculation becomes significantly slower for certain tall-and-skinny input sizes than under PyTorch 1.11.0+cu113. And on A100 80G cards, 2.1.2's nn.Linear is also much slower than the naïve matmul-and-add.

Each row in the tables below shows input shape, running time of nn.Linear, and running time of matmul-and-add. The slowdown (old vs new) is most dramatic (~3x) on A100 80G for input shapes in a certain range (as marked) that we happen to care about.

### 1.11.0+cu113
```
(256, 16)        2.39 ms,   2.60 ms
(512, 16)        2.37 ms,   2.62 ms
(1024, 16)       2.32 ms,   2.56 ms
(2048, 16)       2.37 ms,   2.61 ms
(4096, 16)       2.29 ms,   2.55 ms
(8192, 16)       2.92 ms,   3.23 ms
(16384, 16)      2.97 ms,   3.21 ms
(32768, 16)      2.88 ms,   3.18 ms
(65536, 16)      2.99 ms,   3.30 ms              <--- here
(131072, 16)     2.99 ms,   3.30 ms
(262144, 16)     5.64 ms,   7.41 ms
(524288, 16)    13.82 ms,  12.03 ms
(1048576, 16)   23.59 ms,  24.13 ms              <--- till here          
(2097152, 16)   47.89 ms,  46.01 ms
(4194304, 16)   91.92 ms,  91.43 ms
(8388608, 16)  184.56 ms, 182.64 ms
```

### 2.1.2+cu121
```
(256, 16)        2.30 ms,   3.04 ms
(512, 16)        2.17 ms,   2.93 ms
(1024, 16)       2.23 ms,   3.00 ms
(2048, 16)       2.27 ms,   3.15 ms
(4096, 16)       2.21 ms,   2.99 ms
(8192, 16)       2.17 ms,   2.98 ms
(16384, 16)      2.18 ms,   2.95 ms
(32768, 16)      3.09 ms,   2.92 ms
(65536, 16)      5.81 ms,   2.94 ms               <--- here
(131072, 16)    11.04 ms,   4.53 ms
(262144, 16)    21.69 ms,   6.24 ms
(524288, 16)    44.38 ms,  12.09 ms
(1048576, 16)   86.68 ms,  25.51 ms               <--- till here
(2097152, 16)   46.75 ms,  47.30 ms
(4194304, 16)   93.70 ms,  91.80 ms
(8388608, 16)  184.15 ms, 184.72 ms
```

### Versions

Issue exists from 1.13 -> nightly.

Issue does not exist in 1.11.

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