torus

Rust + Metal LLM inference engine [In progress]

----> repo <----

Introduction

The choice of stack was influenced by the following considerations:

  1. Learning GPU-agnostic programming tricks while still enjoying the comfort of development on a personal laptop.
  2. Targeting a (relatively) less contested niche in the field of inference engines. There's an abundance of high-quality projects aiming at a more sensible stack: Python/Triton; C++/CUDA.
  3. Just a matter of personal taste influenced by author's previous experiences with both Rust and Metal (but not at the same time!).

The scope is relatively small. Stuff that I want to pay extra attention to:

  1. Metal kernels profiling
  2. CPU bottlenecks profiling
  3. Proper graph representation in Rust
  4. CUDA graph alternatives in Metal
  5. Possible benchmarking and comparison to mlx-lm implementation

LLM usage disclaimer: for educational purposes, agentic LLM usage across this project is kept at zero. The only exception is using Mercury autocomplete to write Rust code faster while maintaining full control and understanding of it.

21.08 update

  1. Actually metal-rs is deprecated so I am moving to objc2-meta which is maintained more actively.
  2. Created a small MLP to compare against since I have both relu/gelu and linear kernels.
Thoughts? Leave a comment