Some recent awesome optimizations merged in MLX: Built-in sliding window support in fused SDPA Use mma/nax for QMV with small rows and large N/K, making batched decoding 2~3x faster Dequantize fp4 into float16 directly, which makes prefilling 1.4x faster for old macs Add a 2-stage nax QMM kernel that makes small prompt prefill almost 2x faster for some models 1-bit affine quantization Thin GEMM kernels making small-M/large-N and large-M/small-N 1.5x faster 1.4x faster fused SDPA by prefetching K/V And there are actually a lot more awesome changes that I'm too exhausted to list them all, the commit history of MLX is an interesting read for anyone who wants to make open source contributions.

Is this the first time Meta and Apple work together on CEO level? MLX is the preferred local inference platform for Meta, and we get 0-day support for Muse Spark Open in mlx-lm with vision support. A lot of people use third party inference engines other than mlx-lm, and they will all benefit from the optimization works for Muse Spark Open in MLX. Which is an important meaning of the mlx-lm project: to provide a good foundation for supporting new models, not to compete with open source community. For the technical details, our new blog is a great read: Local and Distributed Large LLM Inference with MLX and RDMA on Mac Studio

MLX now has fast VJP implementations for attention and linear attention, which utilize neutral accelerators. They make training/fine-tuning much faster on Macs, along with massive reduction of memory uses.

@jundotkim makes local inference awesome on Mac with oMLX, really glad to see Hugging Face accelerating the work. The community are building great projects with MLX, and they have been motivating us to keep improving it.

For linear attention a typical implementation only caches the last hidden states, i.e. the context of last token. While for speculative decoding like MTP we need to rollback at least n_draft tokens' context, which does not work with linear attention's cache. An intuitive solution is to remember at least n_draft hidden states in linear attention's cache, I experimented with a cache implementation (which is a bit frustrating to write) and I think the idea works well. The really confusing thing is, no one seems to take this approach, i.e. making a general cache for linear attention that stores hidden states in temporal manner with fixed size, instead people wire MTP inside GDN with a checkpoint of the hidden states during the draft window. Surely it works, but isn't it the ugliest possible implementation and you would have to pollute code of every model that uses MTP? I can see an answer is to avoid increased RAM usage, but is is really small keeping only n_draft context, and the intrusive changes to model implementation kill the elegant abstractions.

MLX v0.32.2 is here. It has only been a week since last release, but there are some nice optimizations benefiting both decoding and prefilling that we would like you to enjoy early:

Thoughts on mlx-lm: top-priority is making it the central registry of MLX model implementations, with tools for evaluation and profiling, and we should add vision models too. Inference engines can have their own schedulers and custom kernels, and do whatever hack to make inference ultra fast, while importing mlx-lm as a library of models. We can rely on community to contribute model implementations, but there would be a fixed procedure to verify correctness of the implementation, ideally automatically. Everything else except for critical bugs, should be irrelevant at the moment and I'm closing PRs and issues aggressively. Many people will be mad at this, and certainly I would be making mistakes closing legitimate things, but for the project to survive, and for the community to grow healthy, I don't see another way.

MLX v0.32.1 has been released, with a lot of minor fixes and performance improvements thanks to the contributors. Last 2 releases also received much more external contributions than ever.

I was working on some optimization for mlx and noticed that our rmsnorm backward was hitting significantly lower bandwidth than forward. I wrote a small note on why that’s the case and how it could be fixed.