Build a large language model inference engine optimized for Apple Silicon (M-series chips) from scratch, starting with GPT-2 Small (124M), while developing a detailed understanding of every part of the Transformer architecture.
- Implement the full GPT-2 inference workflow from scratch across four progressive phases.
- Make effective use of Apple Silicon's unified memory architecture (UMA) and Metal Performance Shaders.
- Manually implement the BPE tokenizer, multi-head attention, LayerNorm, residual connections, and KV Cache.
- Load the official GPT-2 weights and generate text, ultimately reaching 60+ tok/s.
| Phase | Tech Stack | Highlights | Average Speed |
|---|---|---|---|
| Phase 0 | HuggingFace Transformers + MPS | Use existing libraries to validate the baseline | ~45 tok/s |
| Phase 1 | Pure numpy, CPU | Handwritten operators without acceleration | ~3 tok/s |
| Phase 2 | PyTorch nn.Module + MPS | Handwritten model structure with GPU acceleration | ~80 tok/s |
| Phase 3 | C++ + Accelerate + Metal | Handwritten inference engine with KV Cache | ~63 tok/s |
Note: The numbers above come from
benchmark.pymeasurements on Apple Silicon M-series hardware, averaged over three runs.
.
├── model/ # Model weights and config; download manually, see DOWNLOAD.md
│ ├── config.json
│ ├── vocab.json
│ ├── merges.txt
│ ├── tokenizer.json
│ ├── tokenizer_config.json
│ └── model.safetensors # ~548 MB, excluded by .gitignore
│
├── phase0/
│ └── run.py # transformers + torch MPS, about 20 lines
│
├── phase1-numpy/
│ ├── loader.py # Handwritten safetensors parser
│ ├── model.py # Pure numpy GPT-2 forward pass
│ ├── generate.py # Top-k sampling and autoregressive generation
│ └── run.py # Main entry point
│
├── phase2-torch/
│ ├── loader.py # safetensors to torch.Tensor
│ ├── model.py # nn.Module GPT-2 with F.scaled_dot_product_attention
│ ├── generate.py # @torch.inference_mode with MPS synchronization
│ └── run.py # MPS warmup and speed test
│
├── phase3-cpp/
│ ├── CMakeLists.txt # CMake, arm64, Accelerate + Metal + MPS
│ ├── PRD.md # Design document
│ ├── main.cpp # CLI entry point
│ ├── include/ # Headers: common / loader / tokenizer / ops / model
│ ├── src/
│ │ ├── loader.cpp # Binary safetensors parsing for F32/F16/BF16
│ │ ├── tokenizer.cpp # BPE tokenizer with byte-to-unicode mapping
│ │ ├── ops_cpu.cpp # cblas_sgemm, layer_norm, gelu, softmax
│ │ ├── ops_metal.mm # MPSMatrixMultiplication and runtime shader compilation
│ │ └── model.cpp # KV Cache prefill/decode and generate()
│ └── third_party/
│ └── nlohmann/json.hpp # Single-header JSON library
│
├── benchmark.py # Benchmark all four phases and generate REPORT.md
├── REPORT.md # Latest benchmark results
├── DOWNLOAD.md # Weight download instructions
└── .gitignore
git clone <repo-url>
cd Apple-Silicon-LLM-Engine-from-Scratch
# See DOWNLOAD.md to download GPT-2 Small weights into model/Phase 0 (transformers)
conda activate base
cd phase0 && python run.pyPhase 1 (pure numpy)
cd phase1-numpy && python run.pyPhase 2 (PyTorch + MPS)
cd phase2-torch && python run.pyPhase 3 (C++ + Metal)
cd phase3-cpp
cmake -B build -DCMAKE_BUILD_TYPE=Release
cmake --build build -j$(sysctl -n hw.logicalcpu)
./build/gpt2 --prompt "The meaning of life is" --max 100python benchmark.py
# Test all phases and write the results to REPORT.mdGPT-2 Small:
| Parameter | Value |
|---|---|
| Parameters | 124M |
| Layers | 12 |
| Attention heads | 12 |
| Hidden size | 768 |
| Context length | 1024 tokens |
| Vocabulary size | 50257 |
| License | MIT |
- safetensors parsing: A handwritten binary parser with F32 / F16 / BF16 support.
- QKV layout:
qkv[pos, 3*768]stores[Q|K|V]interleaved in each row, rather than as three independent blocks. - KV Cache: Prefill processes the full prompt, while decode computes only one token per step (O(T) vs O(T²)).
- Metal UMA:
storageMode=SharedMTLBuffer allows CPU/GPU zero-copy sharing of physical memory. - Runtime shader compilation: MSL source is embedded as strings and compiled with
newLibraryWithSource:, without requiring thexcrun metaltoolchain. - BPE tokenizer: Byte-to-unicode mapping with greedy merges, supporting original GPT-2 encode/decode behavior.
- macOS 13+ (Ventura or newer)
- Apple Silicon (M1 / M2 / M3 / M4 series)
- Python 3.10+ with a conda environment
- Xcode Command Line Tools for compiling Phase 3
MIT