DeepSeek開(kāi)源周第三彈！DeepSeek-AI 重磅發(fā)布高效FP8 GEMM庫(kù) DeepGEMM：極致性能，代碼精簡(jiǎn)，助力V3/R1模型訓(xùn)練與推理！

簡(jiǎn)單來(lái)說(shuō)這是由 DeepSeek-AI 團(tuán)隊(duì)精心打造的 FP8 通用矩陣乘法 (GEMM) 加速庫(kù)，專為追求極致性能和代碼簡(jiǎn)潔而生

我們一起來(lái)扒一扒：

Hopper GPU 上狂飆 1350+ FP8 TFLOPS！

在算力為王的 AI 時(shí)代，矩陣乘法 (GEMM) 的效率至關(guān)重要。DeepGEMM 正是為了解決這一痛點(diǎn)！它充分利用 NVIDIA Hopper 架構(gòu) GPU的強(qiáng)大算力，在 FP8 精度下，性能 高達(dá) 1350+ TFLOPS！這意味著更快的模型訓(xùn)練速度，更流暢的推理體驗(yàn)，以及更低的計(jì)算成本！

DeepGEMM 不僅適用于傳統(tǒng)的 稠密模型，更完美支持 混合專家模型 (MoE) 的 GEMM 計(jì)算，無(wú)論是 DeepSeek-V3 還是 R1 模型，都能得到強(qiáng)力加速！

代碼極簡(jiǎn)！核心邏輯僅 300 行，堪比教程級(jí)！

你沒(méi)聽(tīng)錯(cuò)！DeepGEMM 的核心 kernel 函數(shù)代碼量?jī)H有 驚人的 ~300 行！ DeepSeek-AI 團(tuán)隊(duì)秉持著 “大道至簡(jiǎn)” 的設(shè)計(jì)理念，在保證極致性能的同時(shí)，力求代碼的 可讀性 和 可維護(hù)性。 即使是剛?cè)腴T(mén) CUDA 開(kāi)發(fā)的同學(xué)，也能輕松理解 DeepGEMM 的實(shí)現(xiàn)原理，甚至可以作為學(xué)習(xí) Hopper FP8 矩陣乘法和優(yōu)化的絕佳教程！

? 無(wú)需編譯！完全 Just-In-Time (JIT) 編譯，即裝即用！ ?

告別繁瑣的編譯過(guò)程！DeepGEMM 采用了 全 Just-In-Time (JIT) 編譯 技術(shù)，所有 kernel 都在運(yùn)行時(shí)動(dòng)態(tài)編譯，無(wú)需在安裝時(shí)進(jìn)行任何預(yù)編譯。 這意味著你可以 即裝即用 DeepGEMM，省去了大量的配置和編譯時(shí)間，讓你可以更專注于模型開(kāi)發(fā)和實(shí)驗(yàn)。

DeepGEMM 的 JIT 設(shè)計(jì)還帶來(lái)了額外的優(yōu)勢(shì)：它可以根據(jù)不同的 GEMM 形狀、block size 等參數(shù)進(jìn)行 動(dòng)態(tài)優(yōu)化，始終選擇最佳的 kernel 配置，保證在各種場(chǎng)景下都能發(fā)揮出最佳性能。

?? 硬核技術(shù)解析：DeepGEMM 的性能秘訣

DeepGEMM 在代碼簡(jiǎn)潔的同時(shí)，性能卻能比肩甚至超越一些專家調(diào)優(yōu)的庫(kù)，這背后離不開(kāi)一系列硬核技術(shù)的加持：

• ? 精細(xì)粒度 Scaling (Fine-grained Scaling)： DeepGEMM 采用了 DeepSeek-V3 論文中提出的精細(xì)粒度 scaling 技術(shù)，更有效地利用 FP8 的動(dòng)態(tài)范圍，提升計(jì)算精度和性能
• ? CUDA-core 雙層累加 (Two-level Accumulation)： 為了解決 FP8 tensor core 累加精度不足的問(wèn)題，DeepGEMM 巧妙地使用了 CUDA-core 雙層累加技術(shù)，保證了計(jì)算結(jié)果的準(zhǔn)確性
• ? Persistent Warp-specialization (持久 Warp 特化)： 借鑒 CUTLASS 的設(shè)計(jì)思想，DeepGEMM 的 kernel 進(jìn)行了 warp 特化，實(shí)現(xiàn)了數(shù)據(jù)移動(dòng)、tensor-core MMA 指令和 CUDA-core promotion 的高效重疊，最大化利用硬件資源
• ? Tensor Memory Accelerator (TMA)： DeepGEMM 充分利用 Hopper 架構(gòu)引入的 TMA 特性，加速 LHS、RHS 矩陣和 scaling factor 的加載，以及輸出矩陣的存儲(chǔ)，實(shí)現(xiàn)更快的數(shù)據(jù)訪問(wèn)速度
• ? 統(tǒng)一優(yōu)化 Block Scheduler 和 Rasterization (柵格化)： DeepGEMM 采用統(tǒng)一的 block scheduler，并結(jié)合 Rasterization 技術(shù)，提升 L2 cache 的復(fù)用率，進(jìn)一步優(yōu)化性能
• ? FFMA SASS Interleaving： DeepGEMM 甚至深入到 SASS 匯編層面進(jìn)行優(yōu)化，通過(guò)調(diào)整 FFMA 指令的 interleaving 模式，提升 warp 級(jí)別的并行度，榨干硬件的每一絲潛力

DeepGEMM 雖然借鑒了 CUTLASS 和 CuTe 的一些概念，但它并沒(méi)有過(guò)度依賴于復(fù)雜的模板或代數(shù)庫(kù)，而是更加注重 簡(jiǎn)潔性 和 易用性。 這使得 DeepGEMM 不僅是一個(gè)高性能的計(jì)算庫(kù)，更是一個(gè)學(xué)習(xí) Hopper FP8 矩陣乘法和優(yōu)化的優(yōu)秀資源

?? 實(shí)測(cè)性能數(shù)據(jù)：實(shí)力說(shuō)話！ ??

DeepGEMM 的性能究竟如何？我們用數(shù)據(jù)說(shuō)話！在 DeepSeek-V3/R1 模型常用的各種 shape 上，DeepGEMM 都展現(xiàn)出了驚人的性能：

• ? Normal GEMMs for dense models (稠密模型 GEMM)

• ? Grouped GEMMs for MoE models (contiguous layout) (MoE 模型 GEMM - 連續(xù)布局)

• ? Grouped GEMMs for MoE models (masked layout) (MoE 模型 GEMM - Masked 布局)

從數(shù)據(jù)中可以看出，DeepGEMM 在各種矩陣 shape 下都表現(xiàn)出色，速度提升明顯！ ??

快速上手 DeepGEMM：只需幾步！

想要體驗(yàn) DeepGEMM 的強(qiáng)大性能？ 上手非常簡(jiǎn)單！

環(huán)境要求：

• ? NVIDIA Hopper 架構(gòu) GPU (sm_90a)
• ? Python 3.8+
• ? CUDA 12.3+ (推薦 12.8+ 獲得最佳性能)
• ? PyTorch 2.1+
• ? CUTLASS 3.6+ (可以通過(guò) Git submodule 克隆)

安裝步驟：

1. 1. 克隆 DeepGEMM 代碼庫(kù) (需要遞歸克隆 submodule)：

git clone --recursive https://github.com/deepseek-ai/DeepGEMM.git

1. 2. 創(chuàng)建 third-party 庫(kù)的符號(hào)鏈接 (CUTLASS 和 CuTe)：

python setup.py develop

1. 3. 測(cè)試 JIT 編譯：

python tests/test_jit.py

1. 4. 測(cè)試所有 GEMM 實(shí)現(xiàn) (normal, contiguous-grouped, masked-grouped)：

python tests/test_core.py

1. 5. 安裝 DeepGEMM：

python setup.py install

安裝完成后，只需在你的 Python 項(xiàng)目中 import deep_gemm 即可開(kāi)始使用！

寫(xiě)在最后：

?? 項(xiàng)目地址：

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

DeepGEMM 現(xiàn)已以MIT許可方式正式開(kāi)源！

DeepGEMM 的靈感來(lái)自 CUTLASS 項(xiàng)目