Tutorial 01: Foundations — MLX, Metal, and the Inference Engine Landscape¶
This tutorial sets the stage. Before writing a single line of model code, you need to understand what MLX is, how it talks to Apple Silicon, and what an inference engine actually does at a high level.
1. What Is an Inference Engine?¶
When you train a large language model, you end up with billions of floating-point numbers (weights) organized into mathematical layers. Inference is the process of feeding a sequence of tokens into those layers and getting the next token out — repeatedly — until the model finishes generating a response.
An inference engine does three things:
- Loads the model weights from disk into GPU-accessible memory.
- Runs the forward pass (matrix multiplications, attention, normalization) on the GPU as fast as possible.
- Manages the generation loop — sampling tokens, maintaining state (the KV cache), and serving multiple users concurrently.
On Apple Silicon, the GPU is not a separate card — it shares memory with the CPU. This changes everything about how we design the engine.
2. Why MLX?¶
MLX is Apple's array computation library. Think of it as NumPy, but the arrays live in unified memory that both the CPU and GPU can access without copying.
The Core Idea¶
import mlx.core as mx
# This creates an array. On Apple Silicon, it lives in unified memory.
# The GPU can read it directly — no PCIe transfer, no explicit upload.
a = mx.array([1.0, 2.0, 3.0])
b = mx.array([4.0, 5.0, 6.0])
# This runs on the GPU. MLX figures out the best Metal kernel to use.
c = mx.matmul(a.reshape(1, 3), b.reshape(3, 1))
print(c) # [[32.0]]
How MLX Differs from PyTorch¶
| Aspect | PyTorch | MLX |
|---|---|---|
| Memory model | CPU arrays + GPU arrays (CUDA), explicit .to(device) |
Unified memory, no transfer needed |
| Default execution | Eager (runs immediately) | Lazy (builds a compute graph, runs on mx.eval()) |
| GPU backend | CUDA (NVIDIA only) | Metal (Apple Silicon only) |
| Autograd | Dynamic computation graph | Functional transforms (mx.grad) |
The lazy evaluation is the most important difference. When you write:
a = mx.array([1.0, 2.0])
b = mx.array([3.0, 4.0])
c = a + b # Nothing happens yet — this just records the operation
d = c * 2 # Still nothing — just another node in the graph
mx.eval(c, d) # NOW the GPU runs both operations in one pass
MLX batches operations together and executes them when you call mx.eval(). This
lets the framework fuse kernels and reduce GPU round-trips.
3. How Metal Fits In¶
Metal is Apple's GPU programming language. MLX uses Metal under the hood:
Your Python code
↓
MLX array operations (mx.matmul, mx.softmax, etc.)
↓
MLX compiler (generates Metal compute shaders)
↓
Metal API (talks to the GPU)
↓
Apple Silicon GPU hardware
You never need to write Metal shaders to use MLX — it provides high-level ops that compile down to Metal automatically. But when you need maximum performance on a specific kernel (like paged attention or quantized matrix multiply), you can write custom Metal shaders and call them from Python through MLX's extension system.
This is exactly what our inference engine does:
# ext.py — lazy loader for custom Metal kernels
class _NativeKernelProxy:
def _ensure_loaded(self):
# Try to import the compiled Metal extension
from anllm import anllm_ext as module
_module = module
def __getattr__(self, item):
# Any attribute access (e.g., native.rope) loads the extension first
module = self._ensure_loaded()
return getattr(module, item)
native = _NativeKernelProxy()
When the Metal extension is not compiled yet, every operation falls back to the pure-Python MLX path. This graceful degradation is a core design principle: kernel dispatch degrades to correctness, never the reverse.
4. Setting Up a Working Environment¶
To follow these tutorials, you need:
# Python 3.11+ (MLX requires it)
python3 --version
# Install MLX
pip install mlx
# Verify it works
python3 -c "import mlx.core as mx; print(mx.array([1,2,3]) ** 2)"
# Should print: array([1, 4, 9], dtype=int32)
If you have Apple Silicon (M1/M2/M3/M4), MLX will automatically use the GPU. On Intel Macs, MLX falls back to CPU — the code still works, just slower.
5. The Data Flow of an LLM Inference¶
Here is the complete lifecycle of a single token prediction:
1. Tokenize input text → [token_id_1, token_id_2, ..., token_id_n]
2. Embed tokens → [batch, seq_len, hidden_size] (lookup table)
↓
3. For each transformer layer:
a. Normalize the input → RMSNorm
b. Compute attention → Q,K,V projections + dot-product attention
c. Add residual connection → x = x + attention_output
d. Normalize again → RMSNorm
e. Feed-forward network → gate_proj + up_proj + down_proj
f. Add residual connection → x = x + ffn_output
↓
4. Final normalization → RMSNorm
5. Project to vocabulary → [batch, seq_len, vocab_size] (logits)
6. Sample next token → argmax or temperature sampling
Each of these steps maps to a specific file in our engine:
| Step | File | Key Class/Function |
|---|---|---|
| Embed | core/embedding.py |
Embedding |
| Normalize | core/norms.py |
RMSNorm |
| Attention | core/attention.py |
scaled_dot_product_attention_grouped |
| FFN | core/layers.py |
silu, linear |
| Quantized ops | core/quantize.py |
QuantizedWeights, quantized_linear |
| Full model | core/model.py |
Qwen3ForCausalLM |
We will build each of these, from math to code, in the tutorials that follow.
6. MLX Primitives You Will Use Everywhere¶
Before we start building, here are the MLX operations that appear in every transformer implementation:
import mlx.core as mx
# --- Matrix multiply (the workhorse of neural networks) ---
# y = xW^T + b (a linear layer)
x = mx.array([[1.0, 2.0, 3.0]]) # shape: (1, 3)
W = mx.zeros((4, 3)) # shape: (4, 3) — 3 inputs, 4 outputs
y = mx.matmul(x, W.T) # shape: (1, 4)
# --- Reshape and transpose (tensor plumbing) ---
a = mx.arange(24).reshape(2, 3, 4) # batch=2, seq=3, dim=4
b = a.transpose(0, 2, 1) # swap seq and dim: (2, 4, 3)
# --- Reductions (softmax, layer norm, etc.) ---
logits = mx.array([[1.0, 2.0, 3.0]])
probs = mx.softmax(logits, axis=-1) # [0.09, 0.24, 0.67]
mean = mx.mean(a, axis=-1, keepdims=True)
# --- Elementwise ops (activations, gating) ---
x = mx.array([-1.0, 0.0, 1.0])
silu = x * (1 / (1 + mx.exp(-x))) # SiLU activation
# --- Slicing and indexing (for KV cache, attention masks) ---
seq = mx.arange(10)
window = seq[3:7] # [3, 4, 5, 6]
# --- Evaluation control ---
a = mx.array([1.0, 2.0])
b = a * 2 # lazy — nothing computed yet
c = a + 3 # lazy — nothing computed yet
mx.eval(b, c) # GPU runs both operations now
7. What We Will Build¶
Over the next 17 tutorials, we will build this, piece by piece:
Tutorial 01 → This file (MLX + Metal foundations)
Tutorial 02 → Linear algebra for transformers
Tutorial 03 → Embeddings and tokenization
Tutorial 04 → RMSNorm (math → code → kernel)
Tutorial 05 → Rotary Position Embeddings (RoPE)
Tutorial 06 → Scaled dot-product attention
Tutorial 07 → Multi-head and grouped-query attention
Tutorial 08 → Feed-forward networks and SiLU gating
Tutorial 09 → Assembling the full transformer block and model
Tutorial 10 → Quantization (4-bit weights)
Tutorial 11 → Mixture of Experts (MoE)
Tutorial 12 → KV cache and autoregressive generation
Tutorial 13 → Paged attention and memory management
Tutorial 14 → Sampling strategies
Tutorial 15 → Continuous batching and scheduling
Tutorial 16 → Writing custom Metal kernels
Tutorial 17 → Building an OpenAI-compatible API server
Tutorial 18 → Putting it all together
By the end, you will have built a complete, working inference engine from scratch. Every line of code will be explained — no black boxes.
Summary¶
- MLX is Apple's lazy-evaluation array library for Apple Silicon. Arrays live in unified memory; operations compile to Metal shaders automatically.
- Metal is the GPU backend. MLX abstracts it away, but you can write custom Metal kernels for performance-critical paths.
- Inference = load weights + run forward pass + sample tokens + manage state.
- Our engine follows a four-layer architecture:
core/(math),cache/(KV state),engine/(generation loop),server/(HTTP API).