$ the-wire · showcase
MPS complex solves land, Arm round() matches torch, dynamo fixes repro fidelity
By RepoJournal · Filed · About PyTorch
PyTorch MPS now supports complex triangular solves with a custom Metal kernel, executorch's Arm backend fixes round() to banker's rounding, and dynamo repros preserve inference flags and FP64 softmax arguments.
MPS gains complex support for triangular solves [1]. MPSMatrixSolveTriangular is float-only and always triggers a device-to-host sync, so a new custom Metal kernel handles complex64. The implementation adds a conjugate flag: triangular_solve and solve_triangular use plain transpose, while _cholesky_solve_helper_mps uses the adjoint, which makes complex cholesky_solve correct and stays a no-op for real. The PR also fixes a non-contiguous out handling bug for float inputs. Separately, the svd_kernel_mps CPU gate for small inputs (numel < 8192) was removed [2] because linalg.svd always runs _linalg_check_errors, forcing a sync regardless of backend; routing on-device tensors to CPU was a net loss. Both changes are in PyTorch core [1] [2].
Two dynamo debug fixes restore repro fidelity. After-AOT repros now preserve is_inference [3]. Previously, when an inference compilation failed with repro_after="aot", the generated repro dropped the flag and recompiled with compile_fx_inner's default is_inference=False, potentially masking the original failure. Also, FP64 softmax repro arguments are fixed [4]: cast_dtype_args_to_fp64 now disables half_to_float after softmax and log-softmax inputs are promoted to FP64, covering default and out overloads with a CPU regression test. A test fix accommodates NumPy 2.5.0's removal of row_stack from the top-level namespace, aligning torch._numpy's exported symbols [5].
In executorch, the Arm backend's round() decomposition now matches torch round semantics [6]. DecomposeRoundPass lowered aten.round using half-away-from-zero, but torch.round is half-to-even (banker's rounding). For input 0.5, the delegated output was 1 instead of 0; for 2.5, it was 3 instead of 2; for -2.5, it was -3 instead of -2. The fix aligns delegated output with the reference. Separately, vendor and example scripts now use dynamic CPU counts for parallel builds [7]. Across 65 sites in 37 files, the hardcoded `cmake --build -j` is replaced with `-j$(( $(nproc ...) + 1 ))`, using `nproc` on Linux and `sysctl -n hw.ncpu` on macOS, degrading to `-j1` if neither exists.
In pytorch/rl, several bugfixes landed. Replay buffer transport and transport_options are now exposed on ReplayBufferConfig with the same defaults as ReplayBuffer.__init__, removing a config/class parity gap [8]. Cross-format checkpoint recovery is stabilized on Windows: rotation now uses filesystem modification nanoseconds before filename ordering, so interrupted replacements keep the newly written checkpoint even when two manifests have identical creation timestamps [9]. On older torch (before 2.4), the legacy TCPStore backend is now selected, since CUDA wheels built with torch 2.1 may not include libuv despite the use_libuv argument [10].
Action items
- → Update Arm backend users to pick up the round() fix; verify outputs for values exactly at .5 pytorch/executorch [immediate]
- → If you use after-AOT repros for inference failures, re-test after this dynamo change to see previously missed failures pytorch/pytorch [plan]
- → If you use MPS complex cholesky_solve or triangular_solve, test after the Metal kernel lands pytorch/pytorch [plan]
References
- [1] [MPS] Add naive triangular_solve for complex numbers (#195861) ↗ pytorch/pytorch
- [2] [MPS] Remove svd_kernel_mps small-input CPU gate (#195950) ↗ pytorch/pytorch
- [3] [dynamo] Preserve is_inference in after-AOT repros (#193565) ↗ pytorch/pytorch
- [4] [dynamo] Fix FP64 softmax repro arguments (#194595) ↗ pytorch/pytorch
- [5] Fix test_exported_objects for NumPy >= 2.5.0 (row_stack removal) (#195047) ↗ pytorch/pytorch
- [6] Arm backend: fix round() decomposition to round-half-to-even ↗ pytorch/executorch
- [7] Use dynamic CPU count for -j in vendor/example scripts and docs ↗ pytorch/executorch
- [8] [BugFix] Expose replay buffer transport in Hydra config ↗ pytorch/rl
- [9] [BugFix] Stabilize cross-format checkpoint recovery ↗ pytorch/rl
- [10] [BugFix] Use legacy TCPStore on older torch ↗ pytorch/rl
