一個輸入維度 4096、輸出維度 4096 的線性層,將一個(batch_size × 4096)輸入矩陣乘以一個(4096 × 4096)權重矩陣,產生一個(batch_size × 4096)輸出。每個輸出元素是一個輸入行和一個權重列的點積:4096 次乘法和 4095 次加法。單一樣本就是 4096 × 4096 ≈ 1680 萬次乘加運算。這只是一層。一個 32 層的 Transformer 每層要做數十次這樣的運算。
矩陣乘法是「令人尷尬的平行」:每個輸出元素都可以獨立計算。CPU 依序計算(每個元素快,但串行)。GPU 同時計算數千個(每個元素稍慢,但大規模平行)。一塊 NVIDIA H100 執行約 1000 TFLOP/s 的 FP16 矩陣乘法——大約每秒 1 千兆次乘加運算。這種平行性正是深度學習變得可行的全部原因。
GEMM(通用矩陣乘法)如此核心,以至於硬體供應商對其進行瘋狂最佳化。CUDA 核心就是為矩陣乘法設計的。Tensor Core(NVIDIA)在單一時脈週期內執行 4×4 矩陣乘法。整個記憶體階層(暫存器、共享記憶體、L1/L2 快取、HBM)都是為了保持資料流向矩陣乘法單元而組織的。當人們說 AI 推論是「記憶體頻寬瓶頸」時,他們的意思是硬體的乘法速度比讀取矩陣的記憶體速度還快。