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

DeepGEMM是什么

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

DeepGEMM的主要功能

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

產品官網

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

DeepGEMM的性能表現

• 普通GEMM（非分組）性能
  • 最高加速比：在特定矩陣形狀下，DeepGEMM能夠實現高達2.7倍的加速，大幅提升矩陣乘法效率。
  • 計算性能：在大規模矩陣計算中，DeepGEMM的計算性能超過1000 TFLOPS，接近Hopper架構GPU的理論峰值。

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

• • • 掩碼布局（Masked Layout）

DeepGEMM的系統要求

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

DeepGEMM的應用場景

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