Contributions to MLX

A new operation: mx.searchsorted

NumPy's searchsorted had been an open request for MLX for a long time. A first version composed it from existing operations; the maintainers wanted a real kernel, so #4035 adds it as a native primitive with its own CPU, Metal and CUDA implementations: one binary search per value. The search runs through the same comparator sort uses on each backend, so ordering agrees with mx.sort everywhere, including where NaNs land. On an M5 Max it is 5.06 to 13.17× faster than the composed version once the search dominates, and 1.90 to 3.29× at the dispatch floor. It was validated against NumPy across 459 checks.

A quantized kernel that was quietly wrong

On M5, gather_qmm with sorted indices, the quantized matrix multiply used for mixture-of-experts layers, returned wrong results whenever the inner dimension was not a multiple of the tile size. Nothing failed; the numbers were simply off. Two bounds errors in the K tail caused it: a tile bounded by the full block instead of what remained, and a safe load zeroing the wrong axis. #4009 fixed both. For 4-bit weights at K = 160 the error went from 3.198 (97.1%) to 0.006, and the person who reported the bug confirmed the fix on their own M5 Max.

Making several Macs behave as one

MLX can spread work across machines over TCP (the ring backend) or over RDMA on Thunderbolt 5 (JACCL). Most of my changes are here, found by running it across Mac minis and fixing what broke:

• #4060: when a peer closed its connection, the remaining ranks missed it and spun at 100% CPU forever. They now report the error and exit.
• #4281: a single transfer of 2 GiB or more was rejected by the socket calls. Transfers are now split, with bandwidth unchanged.
• #4180: naming an RDMA device that did not exist led to an out-of-bounds read and a crash. It is now a clear error.
• #4191: with no usable address on the link, an uninitialised value went into the connection and surfaced much later as an opaque “errno 22”. It now fails early and says what to fix on the Thunderbolt link.
• #4209: the error for a failed queue pair dropped the one detail that explained it. With the errno kept, it showed the devices cap out at ten.
• #4210: Open MPI started without mpirun claimed a world of one and stopped automatic backend selection from ever reaching JACCL.
• #4284: a hostfile missing the RDMA pairs a topology needed was accepted, and every rank ran alone while reporting success. The launcher now rejects it.
• #4286: the build compared the macOS SDK version as a number, so a future SDK 26.10 would read as 26.1 and silently skip building JACCL.
• #4219: runnable data-parallel and tensor-parallel examples, whose results agree across one, two and four ranks.

The CUDA backend

MLX also runs on NVIDIA GPUs. #4619 exports the CUDA backend's command encoder and device classes, as the Metal backend already did, so an extension with its own GPU kernels can use them instead of failing at import on hidden symbols. Tested on an RTX 4090 with an extension that runs its own cuBLASLt matrix multiply, under both mx.grad and mx.compile.

And one in mlx-lm

In MLX's language-model package, #1739 keeps message.content present in a non-streaming chat completion when a model stops mid-reasoning, so OpenAI-style clients no longer crash looking for it.