Some Kernel+Inference Engineering Thoughts
Today, I looked into why ANEs aren't used heavily on Macs for LLMs. I also got through a few more minutes of Baseten's video, which I'm almost done with, and I'm nearly finished with chapter 1 of Basten's inference engineering book.
ANEs on Macs are typically only good for fixed-size matrices, so there's a potential speedup for linear attention on Macs. This also explains their lack of use in Mac-based engines. People have shown that they're slower for certain models, though I forgot whether those models were linear.
I learned that expert parallelism divides the experts in each layer across GPUs. This means GPU:0 could contain experts 1 and 2 for several layers.
I also relearned tensor parallelism. It splits a model's weights for each layer across multiple GPUs, then concatenates the outputs through an all-gather to get the final output.
Apparently, FFNs, aka the MoEs, make up most of Kimi K3. I never realized that!
I'm planning to use what I've learned to optimize the recent Flint model on MLX even further. It's been an amazing model so far!
follow me on twitter for more: https://x.com/dj_sidequests/

Comments
Post a Comment