從 Gflop/s 到 Tflop/s：在 Swift 中優化矩陣乘法
從 Gflop/s 到 Tflop/s：在 Swift 中優化矩陣乘法
訓練大型語言模型 (LLM) 本質上是一場大規模的矩陣乘法練習。其核心過程是執行數兆次的 z += x * y 重複迴圈。對於在 Apple Silicon 上工作的開發者來說，挑戰通常在於如何在 Swift 的高階安全性與這些工作負載所需的原始效能之間取得平衡。
在最近的一次探索中，開發者 zdw 著手使用 Swift 重寫 Andrej Karpathy 的 llm.c（一個與 GPT-2 相容的純 C 實作）。目標不僅是為了達到與 C 的同等效能，更是為了利用 M 系列晶片上所有可用的工具——從 CPU 和 SIMD 指令集，到「秘密」的 AMX 協處理器，以及透過 Metal 使用的 GPU——將 Swift 推向極限。
起點：效能差距
當將核心的 matmul_forward 函數從 C 轉換為基礎 Swift 時，初步結果非常懸殊。儘管在 Release 配置下執行並移除了 runtime asserts，基礎 Swift 實作仍比純 C 版本慢了 15 到 20 倍。
| Model | Tokens/s | Training iterations/s | Training vs llm.c |
|---|---|---|---|
| llm.c | 0.926 | 0.175 | 100% |
| Basic Swift | 0.054 | 0.014 | 7.3% |
這代表了大約 2.8 Gflop/s 的效能——這個數字在 1999 年可能令人印象深刻，但對於現代 LLM 工作負載來說是無法接受的。
主要的罪魁禍首被確定為 _ArrayBuffer.beginCOWMutation()。Swift 的寫入時複製 (Copy-on-Write, COW) 唯一性檢查產生了巨大的開銷，即使在陣列是唯一的情況下也是如此。
縮小差距：Swift 層級的優化
為了突破 COW 的瓶頸，第一步是採用 MutableSpan（在 Swift 6.2 中引入），它提供了一種近乎零開銷的方式來存取記憶體。雖然這提高了訓練速度，但前向傳播 (forward pass) 仍然很慢，因為 Swift 缺乏與 C 的 -ffast-math 標記直接對應的功能，而該標記可以啟用融合乘加 (Fused Multiply-Add, FMA) 指令。
利用鬆散數學與 SIMD
透過使用 Swift-Numerics 函式庫及其 Relaxed.multiplyAdd 函數，實作可以終於利用 fmla (SIMD 向量化 FMA) 指令。僅此一項改動就讓每秒 token 數增加了近 10 倍。
迴圈展開與內聯陣列
為了匹配優化的 C 實作（透過跨越迴圈來鼓勵編譯器進行展開），作者利用了 InlineArray (Swift 6.2)。這允許了堆疊分配的緩衝區，避免了在迴圈內進行堆積分配陣列的高昂成本。在此階段，「快速 Swift」已達到與 C 同等的效能，甚至在每秒訓練迭代次數上略微超過了它（llm.c 的 106.6%）。
規模化：多執行緒與 AMX
雖然解決了單執行緒效能問題，但下一個飛躍需要利用所有可用的 CPU 核心。使用 DispatchQueue.concurrentPerform 允許將工作負載分配到 M3 Max 的 16 個核心上。然而，這也讓程式碼變得相當「雜亂」，需要使用 withUnsafeMutableBufferPointer 和 @unchecked Sendable 包裝器來規避 Swift 的並行安全性檢查。
「秘密」武器：AMX
除了標準的 SIMD，Apple Silicon 還包含 AMX (Apple Matrix Coprocessor)。雖然 Apple 官方僅透過 Accelerate 框架來公開提供此功能，但透過逆向工程的指令集如 AMX_MATFP 允許直接操作 16x16 的分塊 (tiles)。
警告： 不建議在生產環境中直接使用 AMX 指令，因為它們是未公開的，且可能面臨二進位相容性中斷。Accelerate 框架仍是建議的路徑。
實作 AMX 指令將訓練效能推升至原始 llm.c 實作的 958.8%。
最終疆域：Metal 與 GPU
為了達到 Tflop/s 級別，工作負載被移至 GPU 並使用 Metal 進行處理。轉換過程涉及在 Metal/C++ 中編寫計算核心 (compute kernel) 以及在 Swift 中編寫調用層。
- Basic Metal: 簡單的 kernel 提供了一點比 AMX 更大的提升。
- Threaded Metal: 優化
threadsPerThreadgroup帶來了顯著的跳躍（llm.c的 2204.6%）。 - Tiled Metal: 透過實作分塊 (tiling) kernel 來改善記憶體局部性（減少遍歷長列的需求），效能終於突破了 1 Tflop/s 的障礙。
最終效能比較
| Model | Tokens/s | Training iterations/s | Training vs llm.c |
|---|---|---|---|
| llm.c | 0.926 | 0.175 | 100% |
| Multithreaded Swift | 4.356 | 1.014 | 558.5% |
| AMX | 5.884 | 1.678 | 958.8% |
| Tiled Metal | 11.123 | 5.351 | 3057.7% |
技術洞察與反論點
FMA 爭議
A社群討論中提出的一個關鍵點涉及 -ffast-math 的使用。雖然作者使用了它來啟用 FMA，但一些專家認為 -ffast-math 過於寬泛，可能導致不理想的數值精確度問題。對於那些特別想要 FMA 而不想承擔完整 fast-math 風險的人，建議的替代方案是 -ffp-contract=fast。
GPU 軟體護城河
從「Basic Metal」到「Tiled Metal」的困難度凸顯了為什麼像 NVIDIA 的 CUDA 這樣的軟體生態系統能保持如此強大的主導地位。頂尖的 GPU 效能不僅僅取決於硬體，還取決於擁有一個針對特定數據形狀高度優化的 kernel 函式庫。
結論
從最初 2.8 Gflop/s 的簡單實作開始，作者最終達到了 1.1 Tflop/s——效能提升了 382 倍。這段旅程強調了，雖然 Swift 能夠達到甚至超越 C 的速度，但代價往往是失去語言原有的優雅，因為程式碼會逐漸演變成使用不安全指標與手動記憶體管理。對於生產環境的應用程式，教訓很明顯：使用已建立的框架（Accelerate, CoreML, MPSGraph）來處理這些經過多年優化的 kernel。