Repository: pytorch
Issue: Softmax function slows down for data with large range
Labels: module: performance,module: cuda,triaged

### 🐛 Describe the bug

For two tensors with the same shape, but different initial range, `torch.softmax` has significantly different launch latency. 

I have took a look at [the source code](https://github.com/pytorch/pytorch/blob/2b52db9c953d063db7b46c12f4df35b47aca4381/aten/src/ATen/native/cuda/PersistentSoftmax.cuh#L69), but can't get any clue. 

For people who do not know how `softmax` was implemented, please let me briefly explain the core part (in case dim < 1024):
step 1 :  each thread of block loads data into buffer from global memory (VRAM) for selected size.
step 2 :  calculate max value among the target dimension using warp reduction. 
step 3 :  calculate `::exp(value - max_value) -- compliled as __expf if --fast_math flag is passed` for each data point in the buffer and gather the sum among the target dimension using warp reduction. 
step 4 : apply `::exp(value - max_value) / sum` for each data point in the buffer and save to output tensor. 
P.S. `exp(x - max_value) / sum(exp(x - max_value)) = exp(x) * exp(-max_value) / exp(-max_value) * sum(exp(x)) = exp(x) / sum(exp(x))` we get mathematically identical result with much higer numerical stability by subtracting max value. 

My first guess is max_value part, because __expf(x) may have different latency when the value varies. So I implemented the softmax kernel myself with CUDA C, and I found launch time have no difference at all with my own implementation (and faster than pytorch version). Now I am really confused about this. 

Any help is welcome. 

Test code: 
```python
import random
import numpy as np
import torch
start = torch.cuda.Event(enable_timing=True)
end = torch.cuda.Event(enable_timing=True)


def main():
    seed = 8282
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.set_device(0)
    torch.backends.cudnn.benchmark = True
    torch.backends.cuda.matmul.allow_tf32 = False
    torch.backends.cudnn.allow_tf32 = False

    batch_size = 8
    N = 100

    src_seq_len = 512
    tgt_seq_len = 512
    num_heads = 8

    dtype = torch.float16

    x = torch.zeros((batch_size, num_heads, src_seq_len, tgt_seq_len))
    x = x.to(0, dtype=dtype)

    scale = 1
    x.uniform_(-scale, scale)

    cost = 0
    with torch.inference_mode():
        for _ in range(N):
            start.record()
            out = torch.softmax(x, dim=-1)
            end.record()
            torch.cuda.synchronize()
            cost += start.elapsed_time(end)
    print(cost)

    scale = 100
    x.uniform_(-scale, scale)

    cost = 0
    with torch.inference_mode():
        for _ in range(N):
            start.record()
            out = torch.softmax(x, dim=-1)
            end.record()
            torch.cuda.synchronize()
            cost += start.elapsed_time(end)
    print(cost)

if __name__ == '__main__':
    main()
```

Result of above code:
```bash
10.987520195543766
33.92409551143646
```

### Versions

PyTorch version: 1.13.0a0+08820cb
Is debug build: False
CUDA used to build PyTorch: 11.7
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.4 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.1) 9.4.0
Clang version: Could not collect
CMake version: version 3.23.2
Libc version: glibc-2.31

Python version: 3.8.13 | packaged by conda-forge | (default, Mar 25 2022, 06:04:10)  [GCC 10.3.0] (64-bit runtime)
Python platform: Linux-5.4.0-135-generic-x86_64-with-glibc2.10
Is CUDA available: True
CUDA runtime version: 11.7.99
CUDA_MODULE_LOADING set to:
GPU models and configuration:
GPU 0: NVIDIA A100-PCIE-40GB
GPU 1: NVIDIA A100-PCIE-40GB
GPU 2: NVIDIA A100-PCIE-40GB
GPU 3: NVIDIA A100-PCIE-40GB

Nvidia driver version: 515.86.01
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.4.1
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.4.1
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.4.1
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.4.1
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.4.1
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.4.1
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.4.1
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.22.4
[pip3] pytorch-quantization==2.1.2
[pip3] torch==1.13.0a0+08820cb
[pip3] torch-model-archiver==0.7.0
[pip3] torch-tensorrt==1.2.0a0
[pip3] torchserve==0.7.0
[pip3] torchtext==0.13.0a0
[pip3] torchvision==0.14.0a0
[conda] mkl                       2020.4             h726a3e6_304    conda-forge
[conda] mkl-include               2020.4             h726a3e6_304    conda-forge
[conda] numpy                     1.22.4           py38h99721a1_0    conda-forge
[conda] pytorch-quantization      2.1.2                    pypi_0    pypi
[conda] torch                     1.13.0a0+08820cb          pypi_0    pypi
[conda] torch-model-archiver      0.7.0                    pypi_0    pypi
[conda] torch-tensorrt            1.2.0a0                  pypi_0    pypi
[conda] torchserve                0.7.0                    pypi_0    pypi
[conda] torchtext                 0.13.0a0                 pypi_0    pypi
[conda] torchvision               0.14.0a0                 pypi_0    pypi


cc @ngimel