DeepSeek-R2曝5月前上線！第三彈DeepGEMM 300行代碼暴擊專(zhuān)家優(yōu)化內(nèi)核

2025-02-26 10:24:51

DeepSeek開(kāi)源第三彈，是支持稠密和MoE模型的FP8計(jì)算庫(kù)——DeepGEMM，支持V3/R1訓(xùn)推。僅憑300行代碼，就超過(guò)了專(zhuān)家優(yōu)化的內(nèi)核。開(kāi)發(fā)者驚嘆：DeepSeek有最好的GPU工程師，仿佛擁有某種編譯器黑魔法！更令人興奮的是，DeepSeek-R2有望在5月前提前發(fā)布。

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

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

僅用300行代碼，DeepGEMM開(kāi)源庫(kù)就能超越專(zhuān)家精心調(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í)依賴(lài)，代碼清晰易懂
完全即時(shí)編譯，即用即跑
核心邏輯僅約300行代碼，卻在大多數(shù)矩陣規(guī)模下超越專(zhuān)家級(jí)優(yōu)化內(nèi)核
同時(shí)支持密集布局和兩種MoE布局

圖片

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

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

圖片

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

它改變了我們使用FP8 GEMM庫(kù)的方式，簡(jiǎn)單、快速、開(kāi)源。這就是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開(kāi)源兩天，前兩個(gè)項(xiàng)目爆火程度難以想象。FlashMLA已在GitHub斬獲近10k星標(biāo)，DeepEP的星標(biāo)已有5k。

圖片

DeepGEMM

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

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

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

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

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

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

性能

研究人員在配備N(xiāo)VCC 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ǔ)布局）

圖片

快速入門(mén)

要求

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]方式克隆獲取）

開(kāi)發(fā)

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

首先，通過(guò)命令克隆倉(cāng)庫(kù)及其子模塊。然后，創(chuàng)建第三方庫(kù)（CUTLASS和CuTe）的符號(hào)鏈接以便開(kāi)發(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ì)將包及其依賴(lài)項(xiàng)安裝到系統(tǒng)中以便在項(xiàng)目中使用。

python setup.py install

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

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

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

持久化線程束專(zhuān)用化

遵循CUTLASS的設(shè)計(jì)，DeepGEMM中的內(nèi)核采用線程束（warp）專(zhuān)用化技術(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ù)取

常見(jiàn)的細(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)存訪問(wèn)）集群大小
即使在不進(jìn)行自動(dòng)調(diào)優(yōu)的情況下，也能確定性地選擇最優(yōu)配置
完全展開(kāi)MMA（矩陣乘加）流水線，為編譯器提供更多優(yōu)化機(jī)會(huì)
這一特性對(duì)處理小規(guī)模矩陣運(yùn)算尤為重要
詳細(xì)信息請(qǐng)參考kernel文件中的launch_k_iterations部分

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

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