Optimizing LLM Training in Swift: From 2.8 Gflop/s to 1.1 Tflop/s

A developer documents the process of optimizing handwritten matrix multiplication code in Swift for training a GPT-2 style LLM on Apple Silicon. Starting from a naive Swift implementation that achieved only 2.8 Gflop/s, the author applied a series of optimizations: replacing Array with MutableSpan to avoid copy-on-write overhead, using Relaxed.multiplyAdd from Swift Numerics for fused-multiply-add instructions, manual loop unrolling with InlineArray, multithreading via DispatchQueue.concurrentPerform, reverse-engineered AMX instructions, and finally custom Metal GPU compute shaders. The final implementation reached 1.1 Tflop/s, a 382x improvement. The reference model is Andrej Karpathy's llm.c, a plain C implementation of GPT-2 with 124 million parameters. The author notes that Apple's Accelerate framework and other libraries would be more efficient for production use, and that the AMX unit is undocumented and subject to breaking changes. The article is part one of a series exploring neural network training in Swift on Apple Silicon.