【DeepSeek】DeepSeek開源第三彈！極致榨乾GPU，FP8訓推秘籍公開

內建JIT編譯，像教學一樣幹淨！

智東西2月26日報導，剛剛，DeepSeek開源周第三彈發佈——DeepGEMM，一個支援密集和MoE GEMM的FP8 GEMM庫，為V3/R1訓練和推理提供動力。

• ⚡ Hopper GPU上性能高達1350+ FP8 TFLOPS
• ✅ 沒有過多的依賴，像教學一樣幹淨
• ✅ 完全JIT即時編譯（安裝不用預編譯）
• ✅ 極簡設計：核心邏輯約為300行 - 在大多數矩陣大小上都優於專家調整的kernels
• ✅ 支援密集（Dense）佈局和兩種MoE佈局

GitHub：

https://github.com/deepseek-ai/DeepGEMM

眼尖的網友已經在項目貢獻者名單中捕捉到了一個“Liang”，並在DeepSeek推文評論區發問：“是梁文鋒（DeepSeek創始人）嗎？”

DeepGEMM是一個專為乾淨、高效的FP8通用矩陣乘法（GEMM）而設計的庫，具有細粒度擴展功能，如DeepSeek-V3中所述。它支援普通和混合專家（MoE）分組GEMM。該庫用CUDA編寫，在安裝過程中無需編譯，而是使用輕量級即時（JIT）模組在執行階段編譯所有kernel。

根據DeepSeek曬出的資料，普通GEMM（密集模型）中矩陣運算可提速多達2.7倍，分組GEMM（MoE模型）中連續性佈局、掩碼佈局下可提速多達1.2倍。

目前，DeepGEMM僅支援輝達Hopper Tensor Core。為瞭解決不精確的FP8 Tensor Core累積問題，它採用了CUDA核心兩級累積（提升）。

雖然它利用了CUTLASS和CuTe的一些概念，但它避免了對其範本或代數的過度依賴。相反，該庫的設計非常簡單，只有一個核心kernel函數，包含大約300行程式碼。這使其成為學習Hopper FP8矩陣乘法和最佳化技術的乾淨且易於訪問的資源。

儘管DeepGEMM設計輕量，但其性能卻與各種矩陣形狀的專家調整庫相當或超過後者。

DeepSeek在搭載NVCC 12.8的H800上測試了DeepSeek-V3/R1推理中可能使用的所有形狀（包括預填充和解碼，但沒有張量平行性）。所有加速指標都是與其基於CUTLASS 3.6的內部精心最佳化的實現進行比較計算的。

DeepGEMM在有些形狀上的表現並不是很好，因此DeepSeek歡迎開發者來最佳化PR。在普通GEMM（密集模型）中，矩陣運算最高提速達到2.7倍。

在分組GEMM（MoE模型）中，連續性佈局、掩碼佈局下速度可提升1.1倍~1.2倍。

DeepGEMM一發佈，DeepSeek的推文評論區好評如潮。有人為輝達股票發愁：

有人熱情誇讚新程式碼庫和DeepSeek工程師：

DeepSeek分享了清晰的上手指南，需要Hopper架構GPU、必須支援sm_90a，要求是Python 3.8、CUDA 12.3、PyTorch 2.1、CUTLASS 3.6或更新版本。DeepSeek強烈推薦CUDA 12.8或更高的版本以獲得最佳性能。

開發：

安裝：

將deep_gemm匯入Python項目，就可以開始享用了。

這個程式碼庫僅包含GEMM kernel。它要求LHS擴展因子進行TMA對齊和轉置，並且僅支援NT格式（非轉置LHS和轉置RHS）。對於轉置或其他FP8轉換操作，需單獨實現或將它們融合到先前的kernel中。雖然該庫提供了一些簡單的PyTorch實用函數，但這些函數可能會導致性能下降。DeepSeek的主要重點是最佳化GEMM kernels本身。

除了kernel外，該程式碼庫還提供了一些實用函數和環境變數。

DeepSeek用🐳表示CUTLASS中排除的技術。按照CUTLASS設計，DeepGEMM中的核心經過了warp專門化，可實現重疊資料移動、張量核心MMA指令和CUDA核心提升。下圖是說明此過程的簡化圖：

1、Hopper TMA功能

張量記憶體加速器（TMA）是Hopper架構引入的一項新硬體功能，旨在實現更快、非同步的資料移動。具體來說，DeepSeek利用TMA來實現以下目的：

• LHS、LHS擴展因子和RHS矩陣的TMA負載
• TMA儲存輸出矩陣
• TMA multicast組播（LHS矩陣獨有）
• TMA描述符預取

2、常見細節最佳化

• 使用stmatrix PTX指令
• 針對不同的warpgroups定製的暫存器計數控制
• 儘可能重疊，例如重疊TMA儲存和非TMA RHS擴展因子載入🐳

3、統一最佳化的塊調度器

• 一個調度程序適用於所有非分組和分組核心
• 光柵化以增強L2快取重用

4、完全JIT設計🐳

DeepGEMM採用完全即時編譯（JIT）設計，安裝時無需編譯。所有核心均使用輕量級JIT實現在執行階段進行編譯。這種方法具有以下幾個優點：

• GEMM形狀、塊大小和管道階段數被視為編譯時常數
• 保存暫存器
• 編譯器可能會做更多最佳化
• 自動選擇塊大小、warpgroups數量、最佳流程階段和TMA叢集大小
• 但如果沒有自動調整，最佳方案就會被確定地選擇
• 全面展開MMA流程，為編譯器提供更多最佳化機會
• 對於小形狀非常重要
• 詳情請參閱launch_k_iterations kernel檔案

總體而言，JIT顯著提高了小形狀的性能，類似於Triton編譯器的方法。

5、塊大小不對齊

對於某些形狀，與2的冪對齊的塊大小可能會導致SM未得到充分利用。例如，對於M=256, N=7168，典型的塊大小分配會BLOCK_M=128, BLOCK_N=128導致只有(256 / 128) * (7168 / 128) = 112132個SM得到利用。

為瞭解決這個問題，DeepSeek支援未對齊的塊大小（如 112），使(256 / 128) * (7168 / 112) = 128SM能夠在這種場景中工作。在細粒度擴展的同時實施此技術需要仔細最佳化，但最終可以提高性能。 (智東西)