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