DeepSeek 開源活動來到了第三天，新專案如約而至。
此次開源庫名為 DeepGEMM，GitHub 地址在文末。
DeepSeek 表示，這是一個支援密集和混合專家（MoE，Mixture of Experts）通用矩陣乘法（GEMM，General Matrix Multiplication）的 FP8 GEMM 庫，為 V3/R1 的訓練和推理提供支援。
在 Hopper GPU 上最高可達 1350+FP8 TFLOPS。其他優點包括：
沒有過多的依賴，像教程一樣簡潔完全即時編譯核心邏輯約為 300 行，但在大多數矩陣大小上均優於專家調優的核心支援密集佈局和兩種 MoE 佈局
（來源：DeepSeek）
據 DeepSeek 介紹，DeepGEMM 是一個專門設計的、乾淨且高效的工具庫，它的核心任務是進行一種叫做 GEMM 的數學運算，這是 AI 模型訓練和執行中非常常見的一種計算。
DeepGEMM 的特別之處在於，它使用了一種叫做 FP8 的超高效、低精度計算方式，能讓運算速度更快，同時佔用更少的記憶體。這種方式在 DeepSeek-V3 中被提出，並且支援精細的縮放調整（fine-grained scaling），讓計算更靈活。
除了普通的矩陣乘法，DeepGEMM 還能處理混合專家矩陣乘法。
（來源：DeepSeek）
目前，DeepGEMM 只支援英偉達 Hopper 架構的張量核心。張量核心是 GPU 裡的一種特殊硬體，專門用來加速矩陣運算。不過，Hopper 的張量核心在做 FP8 計算時，會出現累加（accumulation，就是把結果一點點加起來的過程）不夠精確的問題。
為了解決這個問題，DeepGEMM 採用了兩級累加（two-level accumulation）的辦法。它藉助 CUDA 核心來做更精確的累加，確保結果不會因為硬體限制而出錯。
DeepGEMM 借鑑了一些來自 CUTLASS 和 CuTe 的概念。這兩個也是 NVIDIA 的高效能矩陣運算庫。不過，DeepGEMM 沒有完全依賴它們的複雜模板或數學系統。
相反，該庫的設計非常簡單，只有一個核心核心函式，包含大約 300 行程式碼。這使其成為學習 Hopper FP8 矩陣乘法和最佳化技術的、乾淨且易於訪問的資源。
儘管 DeepGEMM 設計輕量，但它的效能足以媲美那些由專家精心調優的庫，甚至在某些矩陣形狀（matrix shapes，指矩陣的大小和結構）上表現得更好。
在效能方面，DeepSeek 在搭載 NVCC 12.8 的 H800 上測試了 DeepSeek-V3/R1 推理中可能使用的所有形狀（包括預填充和解碼，但沒有張量並行性）。
從效能對比表格上可見，DeepGEMM 最高能加速 2.7 倍。
圖 | 密集模型加速資料（來源：DeepSeek）
圖 | MoE 模型加速資料（來源：DeepSeek）
所有加速指標都是與 DeepSeek 基於 CUTLASS 3.6 的內部精心最佳化的實現進行比較計算的。
鑑於 DeepGEMM 在某些矩陣形狀上的表現並不是很好，DeepSeek 也邀請各路大神幫助最佳化這個開源庫。
在最佳化方面，DeepSeek 實現了一些 CUTLASS 設計之外的最佳化。這些也是 DeepGEMM 最具創新的地方。
（來源：DeepSeek）
首先是完全 JIT 設計，安裝時無需編譯。所有核心均使用輕量級 JIT 實現在執行時進行編譯。這種方法的優點包括：
-
GEMM 形狀、塊大小和管道階段數被視為編譯時常量
-
自動選擇塊大小、warpgroups 數量、最佳管道階段和 TMA 叢集大小
-
全面展開 MMA 管道，為編譯器提供更多最佳化機會
（來源：DeepSeek）
其次是支援非對齊的塊大小。在做矩陣乘法時，矩陣會被分成一個個小塊（block）來計算。塊的大小通常是固定的，比如 128×128（因為 128 是 2 的冪，計算機喜歡這種數字）。但有時候，矩陣的尺寸跟這些塊大小不完全匹配，就會浪費計算資源。
舉個例子：
假設矩陣的行數（M）是 256，列數（N）是 7168。如果用常見的塊大小 BLOCK_M=128（行）和 BLOCK_N=128（列），那麼行方向可以分成 256÷128=2 個塊，列方向可以分成 7168÷128=56 個塊。總共用到的 SM（GPU 的計算單元）是 2×56=112 個。
GPU 通常有更多 SM（比如 128 個），但這裡只用了 112 個，剩下的就閒著了，沒充分利用。
DeepGEMM 的辦法是支援非對齊的塊大小，也就是不一定非要是 2 的冪（比如 128），可以根據矩陣尺寸靈活調整。
還是上面的例子，如果把 BLOCK_N 改成 112（而不是 128，不是 2 的冪），行方向還是 256÷128=2 個塊，列方向變成了 7168÷112=64 個塊。總共用到的 SM 變成 2×64=128 個。
這下正好把 128 個 SM 全用上，沒有浪費。
（來源：DeepSeek）
最後一點是，DeepSeek 透過觀察和修改底層程式碼（SASS 彙編指令）實現了效能最佳化。
簡單來說，DeepSeek 發現新版編譯器（NVCC 12.3）裡有些底層程式碼變了，效能變得更好。深入研究後，他們覺得這是因為執行緒排程（yield）的方式變了，能讓 GPU 同時幹更多工作。
於是，他們模仿這個變化，寫了個指令碼修改自己的程式碼（調整 FFMA 指令的某些位），讓 MMA 指令（矩陣乘累加）和提升指令更高效地“重疊在一起執行”。
結果是 DeepGEMM 在某些場景下快了 10% 以上，特別適合那種需要靈活調整的 FP8 矩陣乘法。
最後，DeepSeek 致謝了 CUTLASS 專案，稱其是 DeepGEMM 的靈感來源。
參考資料：
https://x.com/deepseek_ai/status/1894553164235640933
https://github.com/deepseek-ai/DeepGEMM
排版：劉雅坤