M15: native linalg via mx::linalg::* - #27
Merged
Merged
Conversation
Bind lu, svd, qr, cholesky, eigh, solve, and triangular_solve to their MLX linalg counterparts, replacing via_binary round-trips. MLX's linalg primitives are CPU-only (no GPU kernels as of 0.25), so the NIFs force a CPU stream regardless of the caller's stream. Residual fallbacks: - qr with mode: :complete (MLX only supports reduced) - determinant (Nx default calls lu, which is now native) triangular_solve handles left_side: false and transform_a: :transpose by composing native transpose + native solve_triangular.
- Remove dead rank-1 branch in mat_transpose_axes (never called) - Avoid NIF round-trip in SVD: compute full-matrices shape from t.shape instead of querying Native.shape(ref) - Restructure triangular_solve into 4 explicit cases to eliminate double-transpose when transform_a: :transpose + left_side: false - Replace manual lower_triangular_matrix generator with Nx.tril - Extract make_well_conditioned/1 to DRY diagonal-dominance in 3 property tests - Add missing test: triangular_solve with left_side: false + transform_a: :transpose
Two non-deterministic crash modes (SIGABRT from LAPACK SVD convergence failure, SIGSEGV with no trace) observed during M15 development. First seen after M14 but M15 may be the actual trigger. Documents what's known and what investigation is needed.
SIGABRT: the SVD property test was the only linalg test not applying make_well_conditioned/1, so random ill-conditioned matrices could trigger LAPACK convergence failure (sgesvdx_ error code 5) which abort()s the VM. SIGSEGV: all linalg ops force cpu_stream() but input tensors may have pending lazy ops on the caller's GPU stream. Add mx::eval() on inputs before each cross-stream handoff so the CPU primitive never races the GPU for the same buffer. 20 consecutive full-suite runs pass cleanly (previously ~1 in 4 crashed).
SIGABRT fixes: - make_well_conditioned/1 multiplier raised from n*10 to n*10+20. The old value landed on the diagonal-dominance boundary for n=2, allowing f32 rounding to produce singular pivots that abort via MLX's uncatchable StreamThread exception path. SIGSEGV (exit 139) is pre-existing — reproduces at ~4/10 on main with no linalg tests. Likely an MLX Metal driver issue, not M15- related. Updated PLAN.md to document this.
Linalg NIFs take a WorkerThread as first parameter and dispatch via run_sync, consistent with the rest of the NIF surface. Inside run_sync, a CPU default stream is used (MLX linalg ops are CPU-only). MLX handles cross-stream dependencies natively.
Add @describetag :linalg to all linalg describe blocks so they can be included/excluded via mix test --only/--exclude linalg.
ausimian
force-pushed
the
m15-native-linalg
branch
from
April 17, 2026 10:33
406d1ec to
89295e4
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
lu,svd,qr,cholesky,eigh,solve,triangular_solve) offvia_binaryfallback onto native MLXmx::linalg::*primitivesrun_synccallback (adapted to M14.5 worker-thread model)qrwithmode: :completefalls back tovia_binary(MLX only supports reduced QR);determinantuses Nx's default impl which now calls the nativeluTest plan
Nx.BinaryBackendwith well-conditioned random inputsmix precommit(compile, format, credo, tests) — all green--only conformance) — 17 tests pass*_fullsuites (Qwen3, Qwen3 quantized, ViT, Whisper, MNIST training, fast-kernels) — 10 tests pass