第三天，DeepSeek發(fā)布了DeepGEMM。

這是一個(gè)支持稠密和MoE模型的FP8 GEMM（通用矩陣乘法）計(jì)算庫(kù)，可為V3/R1的訓(xùn)練和推理提供強(qiáng)大支持。

僅用300行代碼，DeepGEMM開源庫(kù)就能超越專家精心調(diào)優(yōu)的矩陣計(jì)算內(nèi)核，為AI訓(xùn)練和推理帶來(lái)史詩(shī)級(jí)的性能提升！

DeepGEMM庫(kù)具有以下特征：

• 在Hopper GPU上實(shí)現(xiàn)高達(dá)1350+ FP8 TFLOPS的算力 
• 極輕量級(jí)依賴，代碼清晰易懂 
• 完全即時(shí)編譯，即用即跑 
• 核心邏輯僅約300行代碼，卻在大多數(shù)矩陣規(guī)模下超越專家級(jí)優(yōu)化內(nèi)核 
• 同時(shí)支持密集布局和兩種MoE布局

圖片

開發(fā)者驚嘆道：才300行代碼，就能打敗專家優(yōu)化的內(nèi)核？！

要么是DeepSeek真的破解了GPU運(yùn)算的天機(jī)，要么我們就是見證了有史以來(lái)最高級(jí)的編譯器黑科技。

圖片

總之，這個(gè)DeepGEMM聽起來(lái)簡(jiǎn)直是數(shù)學(xué)界的超級(jí)英雄，比飛快的計(jì)算器還要快。

它改變了我們使用FP8 GEMM庫(kù)的方式，簡(jiǎn)單、快速、開源。這就是AI計(jì)算的未來(lái)！

圖片

同時(shí)，外媒還曝出了另一個(gè)重磅消息：原計(jì)劃在5月初發(fā)布的DeepSeek-R2，現(xiàn)在發(fā)布時(shí)間將再次提前！

圖片

在DeepSeek-R2中，將實(shí)現(xiàn)更好的編碼，還能用英語(yǔ)以外的語(yǔ)言進(jìn)行推理。

業(yè)內(nèi)人士預(yù)測(cè)，DeepSeek-R2的發(fā)布，將是AI行業(yè)的一個(gè)關(guān)鍵時(shí)刻。目前DeepSeek在創(chuàng)建高成本效益模型上的成功，已經(jīng)打破了該領(lǐng)域少數(shù)主導(dǎo)玩家的壟斷。

DeepSeek開源兩天，前兩個(gè)項(xiàng)目爆火程度難以想象。FlashMLA已在GitHub斬獲近10k星標(biāo)，DeepEP的星標(biāo)已有5k。

圖片

圖片

DeepGEMM

DeepGEMM是一個(gè)專為清晰高效的FP8通用矩陣乘法（General Matrix Multiplications，GEMMs）設(shè)計(jì)的庫(kù)，它采用了DeepSeek-V3中提出的細(xì)粒度縮放技術(shù)。

該庫(kù)支持常規(guī)矩陣乘法和混合專家模型（Mix-of-Experts，MoE）分組矩陣乘法。DeepGEMM使用CUDA編寫，無(wú)需在安裝時(shí)進(jìn)行編譯，而是通過(guò)輕量級(jí)即時(shí)編譯（Just-In-Time，JIT）模塊在運(yùn)行時(shí)編譯所有內(nèi)核。

目前，DeepGEMM僅支持NVIDIA Hopper張量核。為了解決FP8張量核在累加計(jì)算時(shí)的精度問題，該庫(kù)采用了基于CUDA核心的兩級(jí)累加（提升）技術(shù)。

雖然DeepGEMM借鑒了CUTLASS和CuTe的一些概念，但避免了過(guò)度依賴它們的模板或代數(shù)系統(tǒng)。

相反，該庫(kù)追求設(shè)計(jì)簡(jiǎn)潔，僅包含一個(gè)核心內(nèi)核函數(shù)，代碼量?jī)H約300行。這使其成為學(xué)習(xí)Hopper FP8矩陣乘法和優(yōu)化技術(shù)的理想入門資源。

盡管采用輕量級(jí)設(shè)計(jì)，DeepGEMM在處理各種矩陣形狀時(shí)的性能都能夠達(dá)到甚至超越經(jīng)專家調(diào)優(yōu)的庫(kù)。

性能

研究人員在配備NVCC 12.8的H800上測(cè)試了DeepSeek-V3/R1推理過(guò)程中，可能使用的所有矩陣形狀（包括預(yù)填充和解碼階段，但不包括張量并行計(jì)算）。

所有性能提升指標(biāo)均與基于CUTLASS 3.6內(nèi)部精心優(yōu)化的實(shí)現(xiàn)進(jìn)行對(duì)比計(jì)算得出。

DeepGEMM在某些矩陣形狀下的表現(xiàn)還不夠理想，如果你對(duì)此感興趣，可以提交優(yōu)化相關(guān)的Pull Request（拉取請(qǐng)求）。

稠密模型的常規(guī)GEMM

下表展示了不同矩陣維度（M、N、K）下DeepGEMM庫(kù)的性能數(shù)據(jù)，結(jié)果顯示在某些配置（如 M=128, N=2112, K=7168）下實(shí)現(xiàn)了高達(dá) 2.4 倍的加速，反映了DeepGEMM在優(yōu)化GPU矩陣計(jì)算方面的效率和靈活性。

圖片

MoE模型的分組GEMM（使用連續(xù)存儲(chǔ)布局）

圖片

MoE模型的分組GEMM（使用掩碼存儲(chǔ)布局）

圖片

快速入門

要求

• NVIDIA Hopper架構(gòu)GPU（需支持sm_90a計(jì)算能力）
• Python v3.8或更高版本
• CUDA v12.3及以上版本（強(qiáng)烈建議使用v12.8或更新版本以獲得最佳性能）
• PyTorch v2.1及以上版本
• CUTLASS v3.6或更高版本 （可通過(guò)Git子模塊[submodule]方式克隆獲?。?/span>

開發(fā)

下面代碼是DeepGEMM項(xiàng)目的安裝和測(cè)試指南。

首先，通過(guò)命令克隆倉(cāng)庫(kù)及其子模塊。然后，創(chuàng)建第三方庫(kù)（CUTLASS和CuTe）的符號(hào)鏈接以便開發(fā)。接著，測(cè)試JIT編譯功能。最后，測(cè)試所有GEMM實(shí)現(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包，會(huì)將包及其依賴項(xiàng)安裝到系統(tǒng)中以便在項(xiàng)目中使用。

python setup.py install

接下來(lái)，在你的Python項(xiàng)目中導(dǎo)入deep_gemm，就可以開始使用啦！

優(yōu)化技術(shù)

注意：下面用??標(biāo)記的是，CUTLASS中未包含的技術(shù)。

持久化線程束專用化

遵循CUTLASS的設(shè)計(jì)，DeepGEMM中的內(nèi)核采用線程束（warp）專用化技術(shù)，實(shí)現(xiàn)了數(shù)據(jù)移動(dòng)、張量核心MMA（矩陣乘累加）指令和CUDA核心提升操作的重疊執(zhí)行。下圖簡(jiǎn)要說(shuō)明了這個(gè)過(guò)程：

TMA線程主要負(fù)責(zé)數(shù)據(jù)加載（Data load）和任務(wù)分發(fā)（TMA issue），用黃色和藍(lán)色表示。數(shù)學(xué)線程則交替執(zhí)行WGMA（Wavefront Matrix Multiply-Accumulate）計(jì)算（綠色）和數(shù)據(jù)提升（Promotion，黃色），展示了一種并行計(jì)算策略，其中數(shù)據(jù)加載與矩陣計(jì)算和優(yōu)化操作協(xié)同工作，以提高效率和性能。

圖片

Hopper TMA特性

張量?jī)?nèi)存加速器（Tensor Memory Accelerator，TMA）是Hopper架構(gòu)引入的新硬件特性，用于實(shí)現(xiàn)更快速的異步數(shù)據(jù)移動(dòng)。具體來(lái)說(shuō)，在以下方面使用TMA：

• LHS（左矩陣）、LHS縮放因子和RHS（右矩陣）的TMA加載
• 輸出矩陣的TMA存儲(chǔ)
• LHS矩陣的TMA多播
• TMA描述符預(yù)取

常見的細(xì)節(jié)優(yōu)化

• 使用stmatrixPTX指令
• 針對(duì)不同線程束組的寄存器數(shù)量精確控制
• 最大化指令重疊，如TMA 存儲(chǔ)與非TMA RHS 縮放因子加載的重疊??

統(tǒng)一且經(jīng)過(guò)優(yōu)化的塊調(diào)度器

• 所有非分組和分組內(nèi)核使用同一調(diào)度器
• 采用光柵化技術(shù)提高L2緩存重用率

完全JIT設(shè)計(jì) ??

DeepGEMM采用完全即時(shí)編譯（JIT）設(shè)計(jì)，無(wú)需在安裝時(shí)編譯。所有內(nèi)核在運(yùn)行時(shí)通過(guò)輕量級(jí)JIT實(shí)現(xiàn)進(jìn)行編譯。這種方法具有以下優(yōu)勢(shì)：

• GEMM（通用矩陣乘法）形狀、塊大小和流水線階段數(shù)被視為編譯時(shí)常量

          有效節(jié)省寄存器空間

           使編譯器能夠進(jìn)行更多優(yōu)化

• 能夠自動(dòng)選擇塊大小、線程組數(shù)量、最優(yōu)流水線階段和TMA（張量?jī)?nèi)存訪問）集群大小
• 即使在不進(jìn)行自動(dòng)調(diào)優(yōu)的情況下，也能確定性地選擇最優(yōu)配置
• 完全展開MMA（矩陣乘加）流水線，為編譯器提供更多優(yōu)化機(jī)會(huì)
• 這一特性對(duì)處理小規(guī)模矩陣運(yùn)算尤為重要
• 詳細(xì)信息請(qǐng)參考kernel文件中的launch_k_iterations部分

總的來(lái)說(shuō)，JIT顯著提升了小形狀的計(jì)算性能，這與Triton編譯器采用的方法類似。

非對(duì)齊塊大小??

對(duì)于某些形狀，采用2的冪次對(duì)齊的塊大小可能導(dǎo)致SM利用率不足。

例如，當(dāng)M=256，N=7168時(shí)，傳統(tǒng)的塊大小分配BLOCK_M=128，BLOCK_N=128只能利用 (256/128) * (7168/128) = 112個(gè)SM（總共132個(gè)）。

為解決這個(gè)問題，團(tuán)隊(duì)為諸如112這樣的非對(duì)齊塊大小提供了支持，使得 (256/128) * (7168/112) = 128個(gè)SM能夠充分工作。將這種技術(shù)與細(xì)粒度縮放結(jié)合需要精心優(yōu)化，但最終能帶來(lái)顯著的性能提升。

FFMA SASS交錯(cuò)優(yōu)化??

團(tuán)隊(duì)發(fā)現(xiàn)CUTLASS FP8內(nèi)核在NVCC 12.2和12.3版本之間存在性能差異。

通過(guò)比對(duì)編譯后的SASS代碼，可以發(fā)現(xiàn)在一系列FADD指令中有一個(gè)位按交錯(cuò)模式翻轉(zhuǎn)。

參考開源CUDA匯編器實(shí)現(xiàn)后，團(tuán)隊(duì)確定這個(gè)位控制著讓出（yield）操作，可能用于增強(qiáng)線程束級(jí)并行性（推測(cè)是通過(guò)讓出當(dāng)前線程束使其他線程束得以執(zhí)行）。

為此，團(tuán)隊(duì)開發(fā)了專門的腳本來(lái)修改編譯后二進(jìn)制中的FFMA指令。除了修改讓出位，還調(diào)整了重用位（當(dāng)線程束被讓出時(shí)禁用寄存器重用）。

這種優(yōu)化通過(guò)創(chuàng)造更多MMA指令和提升類FFMA指令重疊的機(jī)會(huì)，顯著提高了細(xì)粒度縮放FP8 GEMM的性能（在某些情況下提升超過(guò)10%）。

參考資料：https://x.com/deepseek_ai/status/1894553164235640933