Repository: pytorch
Issue: The results of baddbmm on CPU seem to have issues under certain conditions
Labels: high priority,triaged,module: regression,module: correctness (silent),module: arm

### 🐛 Describe the bug

```python
import torch

b = 64
m = 1
n = 5
k =128

tensor1 = torch.ones(b, m, k, dtype=torch.float16)
tensor2 = torch.ones(n, b, k, dtype=torch.float16)
tensor3 = torch.ones(b, m, n, dtype=torch.float16)
# tensor2 = tensor2.transpose(0, 1).transpose(1, 2)
tensor2 = tensor2.permute(1, 2, 0)

print(f'{tensor1.shape} {tensor2.shape} {tensor3.shape}')
matmul_result_0 = torch.baddbmm(
        tensor3,
        tensor1,
        tensor2,
        beta=0.0,
        alpha=0.1
    )

matmul_result_1 = torch.baddbmm(
        tensor3,
        tensor1,
        tensor2,
        beta=0.0,
        alpha=0.2
    )

matmul_result_2 = 0.3 * tensor1 @ tensor2 + tensor3 * 0
matmul_result_3 = 0.4 * tensor1 @ tensor2 + tensor3 * 0
if(torch.equal(matmul_result_3, matmul_result_2)):
  print(f'@...... equal {matmul_result_2[0, 0, :10]}    {matmul_result_3[0, 0, :10]}')
if torch.equal(matmul_result_1, matmul_result_0) :
  print(f"equal.....{matmul_result_0[0, 0, :10]}    {matmul_result_1[0, 0, :10]}")
```
I don't understand why the results of `matmul_result_0` and `matmul_result_1` are equal, but when I use `@` for the calculation, the results are not the same. On another machine with PyTorch version `2.3.1`, the results of `matmul_result_0` and `matmul_result_1` are correct as well.

### Versions

```bash
PyTorch version: 2.4.0
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.4 LTS (aarch64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.2) 9.4.0
Clang version: Could not collect
CMake version: version 3.16.3
Libc version: glibc-2.31

Python version: 3.9.0 | packaged by conda-forge | (default, Nov 26 2020, 07:47:13)  [GCC 9.3.0] (64-bit runtime)
Python platform: Linux-5.4.0-100-generic-aarch64-with-glibc2.31
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:                    aarch64
CPU op-mode(s):                  64-bit
Byte Order:                      Little Endian
CPU(s):                          96
On-line CPU(s) list:             0-95
Thread(s) per core:              1
Core(s) per socket:              48
Socket(s):                       2
NUMA node(s):                    4
Vendor ID:                       0x48
Model:                           0
Stepping:                        0x1
CPU max MHz:                     2600.0000
CPU min MHz:                     200.0000
BogoMIPS:                        200.00
L1d cache:                       6 MiB
L1i cache:                       6 MiB
L2 cache:                        48 MiB
L3 cache:                        96 MiB
NUMA node0 CPU(s):               0-23
NUMA node1 CPU(s):               24-47
NUMA node2 CPU(s):               48-71
NUMA node3 CPU(s):               72-95
Vulnerability Itlb multihit:     Not affected
Vulnerability L1tf:              Not affected
Vulnerability Mds:               Not affected
Vulnerability Meltdown:          Not affected
Vulnerability Spec store bypass: Vulnerable
Vulnerability Spectre v1:        Mitigation; __user pointer sanitization
Vulnerability Spectre v2:        Not affected
Vulnerability Srbds:             Not affected
Vulnerability Tsx async abort:   Not affected
Flags:                           fp asimd evtstrm aes pmull sha1 sha2 crc32 atomics fphp asimdhp cpuid asimdrdm jscvt fcma dcpop asimddp asimdfhm

Versions of relevant libraries:
[pip3] numpy==1.26.4
[pip3] torch==2.4.0
[pip3] torchaudio==2.4.0
[pip3] torchvision==0.19.0
[conda] numpy                     1.26.4                   pypi_0    pypi
[conda] torch                     2.4.0                    pypi_0    pypi
[conda] torchaudio                2.4.0                    pypi_0    pypi
[conda] torchvision               0.19.0                   pypi_0    pypi
```

cc @ezyang @gchanan @zou3519 @kadeng @msaroufim @malfet @snadampal @milpuz01