DeepSeek AI 發布 DeepGEMM：一個支持密集和 MoE GEMM 的 FP8 GEMM 庫，為 V3/R1 訓練和推理提供動力

有效的矩陣乘法在現代深度學習和高效能計算中仍然是非常重要的部分。隨著模型變得越來越複雜，傳統的通用矩陣乘法（General Matrix Multiplication, GEMM）方法常常面臨記憶體帶寬限制、數值精度和硬體利用不當等挑戰。這些問題因為新興的混合精度格式（例如 FP8）而變得更加複雜，這需要小心處理以避免計算不準確。最近，NVIDIA 的 Hopper 張量核心架構的進步為提高性能創造了機會，但只有當軟體設計能充分利用這些能力時，才能實現。這樣的背景下，需要一些工具來解決性能瓶頸，同時保持設計的簡單性和透明性。

DeepSeek AI 發布的 DeepGEMM 是一個提升 FP8 GEMM 操作的周到方法。DeepGEMM 專門設計用於高效且清晰的 FP8 矩陣乘法，並具有精細的縮放功能，支持標準和混合專家（Mix-of-Experts, MoE）分組 GEMM。這個庫是用 CUDA 編寫的，因其使用輕量級的即時編譯（Just-In-Time, JIT）模組進行運行時內核編譯而脫穎而出。這樣的設計選擇意味著在安裝過程中不需要冗長的編譯時間，使其能輕鬆整合到現有項目中。DeepGEMM 專為 NVIDIA Hopper 張量核心量身打造，確保它能利用現代硬體的能力，同時解決 FP8 累加不精確等固有挑戰。

技術細節與優勢

DeepGEMM 的核心是結合精細縮放和 FP8 算術來平衡速度和數值準確性。為了解決 FP8 張量核心累加的問題，這個庫通過 CUDA 核心使用兩級累加策略，通常稱為提升。這種方法在計算過程中最小化錯誤，而不犧牲性能。實現上相當簡潔，單個核心內核函數約有 300 行代碼。這種簡單性不僅有助於理解基本原理，還方便社群進一步改進。

DeepGEMM 的靈感來自於已建立的庫，如 CUTLASS 和 CuTe，但它故意避免對複雜模板或代數框架的重度依賴。相反，重點是提供一個乾淨且易於訪問的代碼庫，專注於優化正常和分組配置的 GEMM 操作。對於 MoE 模型的分組 GEMM 支持以兩種形式實現：連續和遮罩佈局。每種形式都經過精心結構設計，以適應每個專家的不同標記數量，反映現代推理和訓練任務的實際需求。

性能洞察與考量

DeepGEMM 倉庫提供的性能數據清楚地顯示了其效率的提升。在 NVIDIA H800 GPU 和 NVCC 12.8 的測試中，顯示在不同矩陣維度下，DeepGEMM 的速度提升與經過精心優化的 CUTLASS 基礎實現相比相當。例如，正常的 GEMM 操作顯示速度提升因子範圍約為 1.4 倍到 2.7 倍，具體取決於矩陣形狀。在 MoE 模型的分組 GEMM 中，連續和遮罩佈局也顯示出一致的改進，雖然提升幅度較小，約為 1.1 倍到 1.2 倍。

這些性能提升是幾個深思熟慮的設計決策的結果。這個庫的 JIT 編譯策略允許根據特定的 GEMM 形狀和硬體配置動態優化內核參數，例如區塊大小、管道階段數和工作組。此外，利用 Hopper 的張量記憶體加速器（Tensor Memory Accelerator, TMA）有助於優化數據移動，這是實現現代 GPU 架構高性能的重要因素。倉庫還詳細說明了幾個實用函數，幫助開發者對齊張量維度和配置共享記憶體，確保這個庫能順利整合到更大的系統中。

結論

DeepGEMM 代表了一種適度且有效的方法來應對 FP8 GEMM 計算的挑戰。通過專注於精度和性能，這個庫為尋求在 NVIDIA Hopper 張量核心上優化矩陣乘法的研究人員和實踐者提供了一個優雅的解決方案。它的設計強調清晰和可訪問性，這在簡潔的代碼庫和通過運行時 JIT 編譯消除預編譯步驟中得以體現。無論是標準 GEMM 還是 MoE 模型所需的更專門的分組 GEMM，DeepGEMM 都提供了一個實用且文檔完善的平台，以提高計算效率。

對於那些希望改善深度學習流程或了解現代 GPU 優化技術的人來說，DeepGEMM 是一個有價值的資源。這個倉庫在 MIT 許可下發布，並得到了開發者社群的支持，邀請進一步探索和改進。

查看 GitHub 倉庫。所有研究的功勞都歸於這個項目的研究人員。此外，隨時在 Twitter 上關注我們，別忘了加入我們的 80,000 多名機器學習 SubReddit。

🚨 推薦閱讀 – LG AI Research 發布 NEXUS：一個先進的系統，整合代理 AI 系統和數據合規標準，以解決 AI 數據集中的法律問題。