【DeepSeek】DeepSeek-R2曝5月前上線！第三彈DeepGEMM 300行程式碼暴擊專家最佳化核心
DeepSeek開源第三彈，是支援稠密和MoE模型的FP8計算庫——DeepGEMM，支援V3/R1訓推。僅憑300行程式碼，就超過了專家最佳化的核心。開發者驚嘆：DeepSeek有最好的GPU工程師，彷彿擁有某種編譯器黑魔法！更令人興奮的是，DeepSeek-R2有望在5月前提前發佈。
第三天，DeepSeek發佈了DeepGEMM。
這是一個支援稠密和MoE模型的FP8 GEMM（通用矩陣乘法）計算庫，可為V3/R1的訓練和推理提供強大支援。
僅用300行程式碼，DeepGEMM開源庫就能超越專家精心調優的矩陣計算核心，為AI訓練和推理帶來史詩級的性能提升！
DeepGEMM庫具有以下特徵：
- 在Hopper GPU上實現高達1350+ FP8 TFLOPS的算力
- 極輕量級依賴，程式碼清晰易懂
- 完全即時編譯，即用即跑
- 核心邏輯僅約300行程式碼，卻在大多數矩陣規模下超越專家級最佳化核心
- 同時支援密集佈局和兩種MoE佈局
開發者驚嘆道：才300行程式碼，就能打敗專家最佳化的核心？！
要麼是DeepSeek真的破解了GPU運算的天機，要麼我們就是見證了有史以來最高級的編譯器黑科技。
總之，這個DeepGEMM聽起來簡直是數學界的超級英雄，比飛快的計算器還要快。
它改變了我們使用FP8 GEMM庫的方式，簡單、快速、開源。這就是AI計算的未來！
同時，外媒還曝出了另一個重磅消息：原計畫在5月初發佈的DeepSeek-R2，現在發佈時間將再次提前！
在DeepSeek-R2中，將實現更好的編碼，還能用英語以外的語言進行推理。
業內人士預測，DeepSeek-R2的發佈，將是AI行業的一個關鍵時刻。目前DeepSeek在建立高成本效益模型上的成功，已經打破了該領域少數主導玩家的壟斷。
DeepSeek開源兩天，前兩個項目爆火程度難以想像。FlashMLA已在GitHub斬獲近10k星標，DeepEP的星標已有5k。
DeepGEMM
DeepGEMM是一個專為清晰高效的FP8通用矩陣乘法（General Matrix Multiplications，GEMMs）設計的庫，它採用了DeepSeek-V3中提出的細粒度縮放技術。
該庫支援常規矩陣乘法和混合專家模型（Mix-of-Experts，MoE）分組矩陣乘法。DeepGEMM使用CUDA編寫，無需在安裝時進行編譯，而是通過輕量級即時編譯（Just-In-Time，JIT）模組在執行階段編譯所有核心。
目前，DeepGEMM僅支援NVIDIA Hopper張量核。為瞭解決FP8張量核在累加計算時的精度問題，該庫採用了基於CUDA核心的兩級累加（提升）技術。
雖然DeepGEMM借鑑了CUTLASS和CuTe的一些概念，但避免了過度依賴它們的範本或代數系統。
相反，該庫追求設計簡潔，僅包含一個核心核心函數，程式碼量僅約300行。這使其成為學習Hopper FP8矩陣乘法和最佳化技術的理想入門資源。
儘管採用輕量級設計，DeepGEMM在處理各種矩陣形狀時的性能都能夠達到甚至超越經專家調優的庫。
性能
研究人員在配備NVCC 12.8的H800上測試了DeepSeek-V3/R1推理過程中，可能使用的所有矩陣形狀（包括預填充和解碼階段，但不包括張量平行計算）。
所有性能提升指標均與基於CUTLASS 3.6內部精心最佳化的實現進行對比計算得出。
DeepGEMM在某些矩陣形狀下的表現還不夠理想，如果你對此感興趣，可以提交最佳化相關的Pull Request（拉取請求）。
稠密模型的常規GEMM
下表展示了不同矩陣維度（M、N、K）下DeepGEMM庫的性能資料，結果顯示在某些組態（如 M=128, N=2112, K=7168）下實現了高達 2.4 倍的加速，反映了DeepGEMM在最佳化GPU矩陣計算方面的效率和靈活性。
MoE模型的分組GEMM（使用連續儲存佈局）
MoE模型的分組GEMM（使用掩碼儲存佈局）
快速入門
要求
- NVIDIA Hopper架構GPU（需支援sm_90a計算能力）
- Python v3.8或更高版本
- CUDA v12.3及以上版本（強烈建議使用v12.8或更新版本以獲得最佳性能）
- PyTorch v2.1及以上版本
- CUTLASS v3.6或更高版本 （可通過Git子模組[submodule]方式克隆獲取）
開發
下面程式碼是DeepGEMM項目的安裝和測試指南。
首先，通過命令克隆倉庫及其子模組。然後，建立第三方庫（CUTLASS和CuTe）的符號連結以便開發。接著，測試JIT編譯功能。最後，測試所有GEMM實現。
安裝
下面程式碼使用指令碼安裝Python包，會將包及其依賴項安裝到系統中以便在項目中使用。
接下來，在你的Python項目中匯入deep_gemm，就可以開始使用啦！
最佳化技術
注意：下面用🐳標記的是，CUTLASS中未包含的技術。
持久化執行緒束專用化
遵循CUTLASS的設計，DeepGEMM中的核心採用執行緒束（warp）專用化技術，實現了資料移動、張量核心MMA（矩陣乘累加）指令和CUDA核心提升操作的重疊執行。下圖簡要說明了這個過程：
TMA執行緒主要負責資料載入（Data load）和任務分發（TMA issue），用黃色和藍色表示。數學執行緒則交替執行WGMA（Wavefront Matrix Multiply-Accumulate）計算（綠色）和資料提升（Promotion，黃色），展示了一種平行計算策略，其中資料載入與矩陣計算和最佳化操作協同工作，以提高效率和性能。
Hopper TMA特性
張量記憶體加速器（Tensor Memory Accelerator，TMA）是Hopper架構引入的新硬體特性，用於實現更快速的非同步資料移動。具體來說，在以下方面使用TMA：
- LHS（左矩陣）、LHS縮放因子和RHS（右矩陣）的TMA載入
- 輸出矩陣的TMA儲存
- LHS矩陣的TMA多播
- TMA描述符預取
常見的細節最佳化
- 使用stmatrixPTX指令
- 針對不同執行緒束組的暫存器數量精確控制
- 最大化指令重疊，如TMA 儲存與非TMA RHS 縮放因子載入的重疊🐳
統一且經過最佳化的塊調度器
- 所有非分組和分組核心使用同一調度器
- 採用光柵化技術提高L2快取重用率
完全JIT設計 🐳
DeepGEMM採用完全即時編譯（JIT）設計，無需在安裝時編譯。所有核心在執行階段通過輕量級JIT實現進行編譯。這種方法具有以下優勢：
- GEMM（通用矩陣乘法）形狀、塊大小和流水線階段數被視為編譯時常數
- 有效節省暫存器空間
- 使編譯器能夠進行更多最佳化
- 能夠自動選擇塊大小、執行緒組數量、最優流水線階段和TMA（張量記憶體訪問）叢集大小
- 即使在不進行自動調優的情況下，也能確定性地選擇最優組態
- 完全展開MMA（矩陣乘加）流水線，為編譯器提供更多最佳化機會
- 這一特性對處理小規模矩陣運算尤為重要
- 詳細資訊請參考kernel檔案中的launch_k_iterations部分
總的來說，JIT顯著提升了小形狀的計算性能，這與Triton編譯器採用的方法類似。
非對齊塊大小🐳
對於某些形狀，採用2的冪次對齊的塊大小可能導致SM利用率不足。
例如，當M=256，N=7168時，傳統的塊大小分配BLOCK_M=128，BLOCK_N=128隻能利用 (256/128) * (7168/128) = 112個SM（總共132個）。
為解決這個問題，團隊為諸如112這樣的非對齊塊大小提供了支援，使得 (256/128) * (7168/112) = 128個SM能夠充分工作。將這種技術與細粒度縮放結合需要精心最佳化，但最終能帶來顯著的性能提升。
FFMA SASS交錯最佳化🐳
團隊發現CUTLASS FP8核心在NVCC 12.2和12.3版本之間存在性能差異。
通過比對編譯後的SASS程式碼，可以發現在一系列FADD指令中有一個位按交錯模式翻轉。
參考開源CUDA彙編器實現後，團隊確定這個位控制著讓出（yield）操作，可能用於增強執行緒束級平行性（推測是通過讓出當前執行緒束使其他執行緒束得以執行）。
為此，團隊開發了專門的指令碼來修改編譯後二進制中的FFMA指令。除了修改讓出位，還調整了重用位（當執行緒束被讓出時停用暫存器重用）。
這種最佳化通過創造更多MMA指令和提升類FFMA指令重疊的機會，顯著提高了細粒度縮放FP8 GEMM的性能（在某些情況下提升超過10%）。
參考資料：
https://x.com/deepseek_ai/status/1894553164235640933 (新智元)