A course on LLM serving using MLX for system engineers. The codebase is solely (almost!) based on MLX array/matrix APIs without any high-level neural network APIs, so that we can build the model serving infrastructure from scratch and dig into the optimizations.
The goal is to learn the techniques behind efficiently serving a large language model (e.g., Qwen2 models).
Why MLX: nowadays it's easier to get a macOS-based local development environment than setting up an NVIDIA GPU.
Why Qwen2: this was the first LLM I've interacted with -- it's the go-to example in the vllm documentation. I spent some time looking at the vllm source code and built some knowledge around it.
The tiny-llm book is available at https://skyzh.github.io/tiny-llm/. You can follow the guide and start building.
You may join skyzh's Discord server and study with the tiny-llm community.
Week 1 is complete. Week 2 is in progress.
| Week + Chapter | Topic | Code | Test | Doc |
|---|---|---|---|---|
| 1.1 | Attention | β | β | β |
| 1.2 | RoPE | β | β | β |
| 1.3 | Grouped Query Attention | β | β | β |
| 1.4 | RMSNorm and MLP | β | β | β |
| 1.5 | Load the Model | β | β | β |
| 1.6 | Generate Responses (aka Decoding) | β | β | β |
| 1.7 | Sampling | β | β | β |
| 2.1 | Key-Value Cache | β | π§ | π§ |
| 2.2 | Quantized Matmul and Linear - CPU | β | π§ | π§ |
| 2.3 | Quantized Matmul and Linear - GPU | β | π§ | π§ |
| 2.4 | Flash Attention 2 - CPU | β | π§ | π§ |
| 2.5 | Flash Attention 2 - GPU | β | π§ | π§ |
| 2.6 | Continuous Batching | β | π§ | π§ |
| 2.7 | Chunked Prefill | β | π§ | π§ |
| 3.1 | Paged Attention - Part 1 | π§ | π§ | π§ |
| 3.2 | Paged Attention - Part 2 | π§ | π§ | π§ |
| 3.3 | MoE (Mixture of Experts) | π§ | π§ | π§ |
| 3.4 | Speculative Decoding | π§ | π§ | π§ |
| 3.5 | Prefill-Decode Separation (requires two Macintosh devices) | π§ | π§ | π§ |
| 3.6 | Parallelism | π§ | π§ | π§ |
| 3.7 | AI Agent / Tool Calling | π§ | π§ | π§ |
Other topics not covered: quantized/compressed kv cache, prefix/prompt cache; sampling, fine tuning; smaller kernels (softmax, silu, etc)