Optimizing Swift Matrix Multiplication for LLM Training on Apple Silicon
Replace Array with MutableSpan and compile with -remove-runtime-asserts to cut COW overhead, raising training iterations from 0.014 to 0.042 per second.
Replace Array with MutableSpan in all matrix multiplication functions and compile with -remove-runtime-asserts to eliminate COW overhead.
Summary
The article explores how to accelerate a handwritten Swift implementation of matrix multiplication for training a small GPT‑2‑compatible model, using the plain C reference from Andrej Karpathy’s llm.c as a benchmark. The C version, compiled with -O3, achieves one training iteration every 7 seconds and an inference speed of less than one token per second, while the naïve Swift translation runs 15–20× slower, delivering only 0.014 training iterations per second and 2.8 Gflop/s. The bottleneck in Swift is identified as _ArrayBuffer.beginCOWMutation() due to runtime array bounds checks, which the author mitigates by compiling with -remove-runtime-asserts and replacing Array with the zero‑overhead MutableSpan introduced in Swift 6.2. After applying MutableSpan to the output array and all intermediate buffers, training iterations jump to 0.042 s⁻¹, a 3× improvement, though the forward pass remains dominated by the inner multiplication loop. The piece concludes that while Swift can approach C performance with careful memory handling, the true speed gains come from avoiding COW overhead and leveraging the new MutableSpan API.
Key changes
- Swift 6.2 adds MutableSpan with zero overhead for mutable buffers
- Using MutableSpan for the output array eliminates _ArrayBuffer.beginCOWMutation() cost
- Compiling with -remove-runtime-asserts removes array bounds checks
- Forward pass hotspot remains value += inp[bt*C+i]*weight[o*C+i]
- Training iterations improve 3× after MutableSpan integration
- Basic Swift is 15–20× slower than C at -O3
- C implementation runs one iteration every 7 s and 0.2 trillion FLOPs per iteration
- Measured Swift performance is 2.8 Gflop/s