@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.

The thing I hate most about aging is that, it becomes much harder to get a perfect sleep. Especially in a long trip and lack of good sleep makes me exhausted even when the trip is supposed to be relaxing and I have to end it at some time, 3 weeks is my current limitation.

In the chatgpt moment I realized AI is the future but it took me more than one year to quit my job and start studying machine learning, because I thought it was too hard to learn from scratch. This was my biggest mistake, so many things could have been done in that year.