SIGN IN SIGN UP

[MPS] Speed up QR with blocked householder updates (#198732)

First reported here https://github.com/kornia/kornia/issues/4522

QR decomposition runs reaaallly slow now on main. Embarassingly slow, batch size 1 **512 x 512 QR at 6.976042 SECONDS**...

Perf, comparisons against before on MPS are deliberately excluded for obvious reason:

| Shape | CPU (ms) | MPS (ms) | CPU / MPS |
|---|---:|---:|---:|
| 128 x 128 | 0.219 | 0.298 | 0.73x |
| 384 x 384 | 1.740 | 0.993 | 1.75x |
| 512 x 512 | 4.083 | 1.465 | 2.78x |
| 1024 x 1024 | 18.796 | 4.954 | 3.79x |
| 2048 x 2048 | 97.228 | 19.503 | 4.98x |
| 1152 x 128 | 1.959 | 0.897 | 2.18x |
| 2048 x 512 | 15.830 | 6.025 | 2.63x |
| 16384 x 512 | 106.555 | 32.437 | 3.29x |
| 64 x 9 x 8 | 0.067 | 0.031 | 2.13x |
| 256 x 3 x 3 | 0.066 | 0.020 | 3.35x |
| 2048 x 9 x 8 | 1.504 | 0.219 | 6.87x |
| 4096 x 9 x 8 | 2.957 | 0.413 | 7.17x |
| 16384 x 9 x 8 | 11.555 | 1.566 | 7.38x |
| 65536 x 9 x 8 | 46.528 | 6.225 | 7.47x |
| 2048 x 3 x 3 | 0.403 | 0.040 | 9.98x |
| 4096 x 3 x 3 | 0.781 | 0.067 | 11.67x |
| 16384 x 3 x 3 | 3.066 | 0.222 | 13.83x |
| 65536 x 3 x 3 | 11.598 | 0.827 | 14.02x |

In terms of numerics, they are better than the ones currently present on main

| Case | Metric | CPU (LAPACK) | MPS main | MPS this PR |
|---|---|---:|---:|---:|
| 512 x 512 | backward error | 8.0e-07 | 4.1e-06 | 7.7e-07 |
| 512 x 512 | orthogonality | 1.9e-05 | 1.8e-04 | 1.9e-05 |
| 1024 x 64 | backward error | 2.2e-07 | 2.2e-06 | 3.2e-07 |
| 64 x 32 x 32 | backward error | 2.4e-07 | 3.5e-07 | 3.7e-07 |
| 128 x 128, scaled by 1e-10 | backward error | 4.0e-07 | 0.70 | 5.0e-07 |
| 128 x 128, scaled by 1e20 | backward error | 4.0e-07 | NaN | 5.1e-07 |
| 256 x 256, condition number 1e6 | backward error | 4.3e-07 | 3.9e-05 | 4.7e-07 |
| R only, 2048 x 2048 | error in R | 1.7e-06 | 9.9e-07 | 1.4e-06 |
Pull Request resolved: https://github.com/pytorch/pytorch/pull/198732
Approved by: https://github.com/malfet
I
Isalia20 committed
4d19930b18fadc326ffe2a9618b8b457f3794043
Parent: e4112a1
Committed by PyTorch MergeBot <pytorchmergebot@users.noreply.github.com> on 10/1/2026, 5:47:32 AM