DeepGEMM - DeepSeek 开源的 FP8 通用矩阵乘法库

发布于 2025-02-26 · DeepGEMM是DeepSeek开源的为高效简洁的FP8矩阵乘法（GEMM）设计的库，目前仅支持NVIDIA Hopper张量核心。DeepGEMM支持普通和混合专家（MoE）分组的GEMM操作，基于即时编译（JIT）技术，无需安装时编译，支持在运行时动态优化。

DeepGEMM是什么

DeepGEMM是DeepSeek开源的 FP8 矩阵乘法（GEMM）库，设计上追求高效和代码简洁，目前只支持 NVIDIA Hopper 张量核心。它支持普通 GEMM 和混合专家（MoE）分组 GEMM，基于即时编译（JIT）技术，安装时不用编译，运行时还能动态优化。精度上用细粒度缩放加 CUDA 核心双级累加，补上 FP8 精度不足的短板；数据传输则借 Hopper 的 Tensor Memory Accelerator（TMA）提速。核心代码只有约 300 行，读起来、改起来都不费劲。也正因为够轻，它在多种矩阵形状上达到甚至超过了专家级优化库。

DeepGEMM的主要功能

高效FP8矩阵乘法（GEMM）：DeepGEMM 是专为 FP8（8 位浮点）矩阵乘法做的库，支持细粒度缩放，矩阵运算的性能和精度都有提升。
支持普通和分组GEMM：
- 普通GEMM：用于常规的矩阵乘法。
- 分组GEMM：处理混合专家（MoE）模型里的分组矩阵乘法，连续布局（contiguous layout）和掩码布局（masked layout）都支持，多个专家共享同一形状时更划算。
即时编译（JIT）设计：所有内核都在运行时用 JIT 动态编译，安装阶段不用编译。编译时按矩阵形状、块大小等参数调优，还能省下寄存器。
Hopper架构优化：针对 NVIDIA Hopper 架构设计，把 Tensor Memory Accelerator（TMA）的加载、存储、多播和描述符预取都用上了，数据传输效率提升明显。
细粒度缩放和双级累加：FP8 精度不够，DeepGEMM 用细粒度缩放来补，再靠 CUDA 核心的双级累加，把 FP8 的计算结果提到更高精度的格式（如 BF16），精度就守住了。
轻量级设计：核心代码只有约 300 行，理解和扩展都容易。它也避开了对复杂模板或代数结构的依赖，学习和优化的门槛随之降低。

DeepGEMM的项目地址

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 这类 Hopper 卡，它们本来就是为 FP8 计算和 Tensor Core 优化的。
- CUDA兼容性：需要 CUDA 12.3 或更高版本，想拿到最好的性能，建议上 CUDA 12.8 或更高版本。
软件要求：
- 操作系统推荐：用 Linux（比如 Ubuntu、CentOS），CUDA 和 PyTorch 在 Linux 上的支持更完整。
- Python版本：Python 3.8或更高版本。
- CUDA工具包：CUDA 12.3 或更高版本。版本还得和 GPU 架构匹配，想用满 Hopper 的特性，建议 12.8 或更高。
- PyTorch：PyTorch 2.1或更高版本。
- CUTLASS库：CUTLASS 3.6或更高版本。
其他要求：
- 常规的编译工具（如gcc、make等）。
- torch.utils.cpp_extension模块，用于CUDA扩展。