Systems · ML inference · 2026
bare-metal-llm
GPT-2 inference on hand-written Metal kernels
A from-scratch GPT-2 124M inference engine on an Apple M2. Tokenizer, safetensors loader, 14 compute kernels and a real KV cache, with no ML framework in the inference path. Matches an MLX reference token for token.
- Hand-written kernels
- 14
- Greedy output vs MLX
- Exact match
- Throughput
- 19.3 tok/s
Why
Ollama, MLX and llama.cpp make running a model on Apple Silicon look effortless. I wanted to know what they are actually doing: memory layout, kernel dispatch, and how a forward pass gets onto the GPU.
The rule was simple. No MLX, no llama.cpp, no PyTorch in the inference path. MLX is used only offline, as a correctness reference and a benchmark baseline.
Picking the model
GPT-2 124M over Qwen2.5-0.5B, because it needs strictly fewer kernel types to reach a working forward pass: LayerNorm instead of RMSNorm, learned position embeddings instead of RoPE, and plain multi-head attention instead of grouped-query attention. Fewer moving parts meant every one of them could be verified.
What got built
Fourteen Metal kernel functions: matmul, linear and its transposed variant for tied embeddings, layernorm, softmax, causal attention (stateless and KV-cached), embedding lookup, QKV split and merge, residual add, and GELU. Each one is tested on its own against a NumPy reference.
A safetensors parser and a byte-level BPE tokenizer, both written from the file formats. The tokenizer matches the reference tokenizers library exactly on every test string.
A real KV cache, so each decoding step only computes the new token and attends against what is cached.
The bug worth telling
GELU started producing NaN on real weights. tanh() was overflowing: once an activation got large enough, the implicit exp(2x) passed float32 range and became Inf/Inf. The formula was right and the output was still wrong. I only found it by running real weights and bisecting layer by layer.
Results
Across three test prompts, next-token argmax and full 15-token greedy generation match the MLX reference exactly. Logits differ by about 1e-4 because of float32 accumulation order, and that never once changed a token.
MLX runs at 103 tok/s and this engine at 19.3 tok/s. That gap is expected. These kernels are the simplest correct version of each op (no tiling, no fusion, no quantization), while MLX picks shape-tuned kernels and uses simdgroup_matrix intrinsics.
What it taught me
Unified memory changes the mental model. Both MLX and llama.cpp allocate everything as shared storage, so there is no upload or download step. The GPU and CPU stop being two different places data lives.
Command buffer submission is expensive. Encoding the whole forward pass into one command buffer was the biggest single lesson from reading ggml-metal and MLX source.
Production dispatch code picks between kernel variants per shape, using empirically tuned constants. That is a later problem than correctness, and knowing not to chase it early was useful too.
Stack: Swift · Metal Shading Language · Python (validation)