DeepGEMM – DeepSeek 開源的 FP8 通用矩陣乘法庫

DeepGEMM是什么

DeepGEMM是由DeepSeek開發(fā)的開源庫，專為高效和簡潔的FP8矩陣乘法（GEMM）而設(shè)計。目前，該庫僅兼容NVIDIA Hopper架構(gòu)的張量核心。DeepGEMM不僅支持普通的GEMM操作，還支持混合專家（MoE）模型中的分組矩陣乘法。它基于即時編譯（JIT）技術(shù)，允許在運行時進行動態(tài)優(yōu)化，無需事先進行編譯。通過細(xì)粒度縮放和CUDA核心的雙級累加機制，DeepGEMM有效解決了FP8精度不足的問題，并利用Hopper的Tensor Memory Accelerator（TMA）特性大幅提升數(shù)據(jù)傳輸效率。其核心代碼簡約，僅約300行，便于學(xué)習(xí)和優(yōu)化，且在多種矩陣形狀下的性能達(dá)到或超過專家級優(yōu)化庫的水平。

DeepGEMM的主要功能

• 高效FP8矩陣乘法（GEMM）：專為FP8（8位浮點數(shù)）矩陣乘法優(yōu)化的庫，采用細(xì)粒度縮放技術(shù)，顯著提升計算性能與精度。
• 支持普通和分組GEMM：
  • 普通GEMM：適合常規(guī)矩陣乘法操作。
  • 分組GEMM：優(yōu)化混合專家（MoE）模型中的分組矩陣乘法，支持連續(xù)布局和掩碼布局，提升多專家共享形狀的計算效率。
• 即時編譯（JIT）設(shè)計：所有內(nèi)核在運行時動態(tài)編譯，避免安裝時編譯，根據(jù)矩陣形狀和塊大小等參數(shù)進行優(yōu)化，提升性能并節(jié)約寄存器。
• Hopper架構(gòu)優(yōu)化：專為NVIDIA Hopper架構(gòu)設(shè)計，充分利用TMA特性，包括加載、存儲、多播和描述符預(yù)取，大幅提高數(shù)據(jù)傳輸效率。
• 細(xì)粒度縮放和雙級累加：通過細(xì)粒度縮放技術(shù)和CUDA核心的雙級累加機制，解決FP8計算的精度問題，將FP8結(jié)果提升至更高精度格式（如BF16），確保計算精度。
• 輕量級設(shè)計：核心代碼簡潔，易于理解和擴展，避免復(fù)雜的模板或代數(shù)結(jié)構(gòu)依賴，降低學(xué)習(xí)和優(yōu)化的門檻。

產(chǎn)品官網(wǎng)

• GitHub倉庫：https://github.com/deepseek-ai/DeepGEMM

DeepGEMM的性能表現(xiàn)

• 普通GEMM（非分組）性能
  • 最高加速比：在特定矩陣形狀下，DeepGEMM能夠?qū)崿F(xiàn)高達(dá)2.7倍的加速，大幅提升矩陣乘法效率。
  • 計算性能：在大規(guī)模矩陣計算中，DeepGEMM的計算性能超過1000 TFLOPS，接近Hopper架構(gòu)GPU的理論峰值。

• 分組GEMM（MoE模型）性能
  • 加速比：在分組GEMM中，DeepGEMM的加速比為1.1至1.2倍，顯著提升MoE模型的訓(xùn)練和推理效率。
  • 內(nèi)存帶寬優(yōu)化：利用TMA特性，DeepGEMM在內(nèi)存帶寬的利用上表現(xiàn)卓越，接近硬件性能極限。
    • 連續(xù)布局（Contiguous Layout）

• • • 掩碼布局（Masked Layout）

DeepGEMM的系統(tǒng)要求

• 硬件要求：
  • GPU架構(gòu)：必須支持NVIDIA Hopper架構(gòu)，具體要求為支持sm_90a的GPU。推薦使用H800或H100等專為FP8計算和Tensor Core優(yōu)化的Hopper架構(gòu)GPU。
  • CUDA兼容性：需支持CUDA 12.3或更高版本，推薦使用CUDA 12.8或更高版本以獲得最佳性能。
• 軟件要求：
  • 操作系統(tǒng)推薦：建議使用Linux操作系統(tǒng)（如Ubuntu、CentOS等），以便于CUDA和PyTorch的更好支持。
  • Python版本：Python 3.8或更高版本。
  • CUDA工具包：CUDA 12.3或更高版本。CUDA版本需與GPU架構(gòu)相匹配，推薦使用12.8或更高版本以充分發(fā)揮Hopper架構(gòu)的優(yōu)勢。
  • PyTorch：PyTorch 2.1或更高版本。
  • CUTLASS庫：CUTLASS 3.6或更高版本。
• 其他要求：
  • 標(biāo)準(zhǔn)編譯工具（如gcc、make等）。
  • torch.utils.cpp_extension模塊，用于CUDA擴展。

DeepGEMM的應(yīng)用場景

• 大規(guī)模AI模型推理：加速高維矩陣乘法，提升推理速度。
• 混合專家（MoE）模型：優(yōu)化分組矩陣乘法，增強計算效率。
• 低精度計算：通過細(xì)粒度縮放解決FP8精度問題，確保高精度輸出。
• 高性能計算：基于Hopper架構(gòu)特性，提升矩陣運算效率。
• 深度學(xué)習(xí)框架優(yōu)化：作為底層優(yōu)化庫，加速模型的訓(xùn)練和推理。