DeepSeek AI 最近釋出的 DeepGEMM 庫憑藉其創新的 FP8 矩陣運算最佳化技術引起了技術社群的關注。雖然該庫為通用矩陣乘法(GEMMs)提供了多項效能改進,但其中 FFMA SASS 交錯技術特別令技術專家印象深刻,在某些情況下效能提升超過 10%。
FFMA SASS 交錯技術背後的奧秘
DeepGEMM 團隊發現在 NVCC 12.2 和 12.3 編譯器版本之間,CUTLASS FP8 核心的效能有所提升。透過對編譯後的 SASS(流式彙編)程式碼進行仔細分析,他們發現 FADD 指令中的一個特定位以交錯模式被翻轉。這個位控制著屈服功能,本質上允許當前執行緒束讓出執行權,透過讓其他執行緒束工作,可能增強執行緒束級並行性。
基於這一發現,團隊開發了一個指令碼來修改編譯後二進位制檔案中的 FFMA(融合浮點乘加)指令。他們不僅操作了屈服位,還翻轉了重用位,因為如果執行緒束被讓出,暫存器就無法被重用。這種看似微小的修改創造了更多機會,使矩陣乘累加(MMA)指令與提升 FFMA 指令重疊,從而帶來顯著的效能提升。
「我認為這真的令人驚歎。」
關鍵 AI 基礎設施的專業最佳化
社群討論強調,雖然這類效能最佳化在效能至關重要的矩陣數學中很常見,但其他 AI 公司尚未廣泛將其應用於這個特定問題。正如一位評論者指出,大多數 AI 參與者依賴高效能 GEMM 操作,但通常會採用 CUTLASS 或 cuBLAS 等標準實現,而不是利用未記錄的功能。
這種級別的最佳化展示了 AI 公司為了從昂貴的 GPU 叢集中榨取每一點效能所願意付出的努力。即使是 10% 的效能提升,在大規模運營時也能轉化為顯著的成本節約。正如討論中指出的,當公司在 GPU 基礎設施上投資數億時,這樣的收益可能足以支付許多人的薪水。
DeepGEMM 主要特點和要求
-
效能提升:與最佳化的 CUTLASS 3.6 實現相比,速度提升高達 2.7 倍
-
最佳化技術:
- 永續性執行緒束專業化
- Hopper TMA(張量記憶體加速器)特性
- 統一塊排程器與光柵化
- 完全 JIT 設計
- 非對齊塊大小
- FFMA SASS 交錯
-
硬體要求:
- 支援 sm_90a 的 Hopper 架構 GPU
- Python 3.8+
- CUDA 12.3+(推薦 12.8+)
- PyTorch 2.1+
- CUTLASS 3.6+
行業影響和可訪問性
DeepGEMM 的開源釋出似乎在戰略上定位為使整個行業受益,特別是為 AI 模型提供服務的大型提供商。該庫需要 Hopper 架構 GPU(支援 sm_90a)並專為 DeepSeek-V3 等場景設計,支援普通和混合專家(MoE)分組 GEMMs。
一些社群成員已經嘗試在 RTX 5080 等消費級硬體上測試該庫,遇到了與共享記憶體容量相關的限制。該庫明確為 NVIDIA Hopper 張量核心設計,使其主要適用於企業級 AI 基礎設施,而非消費級應用。
DeepGEMM 的技術深度突顯了 AI 基礎設施最佳化日益增長的複雜性。隨著 AI 模型在規模和複雜性上不斷增長,這些看似微小的硬體指令級最佳化對於推動當前硬體可能性邊界的組織來說變得越來越有價值。
參考:DeepGEMM: clean and efficient FP8 GEMM kernels with fine-grained scaling