DeepSeek开源高效FP8矩阵乘法库DeepGEMM

2月26日消息，DeepSeek公司在开源周期间继续推进其技术开放计划。在早前宣布开源MLA解码核FlashMLA以及DeepEP两款代码库之后，该公司于第三天正式开放了DeepGEMM代码库。

据介绍，DeepGEMM是一款专注于实现高效FP8通用矩阵乘法（GEMM）的工具库，其设计目标是为V3/R1训练和推理任务提供强大的支持。该库同时兼容普通矩阵运算和专家混合（MoE）分组的GEMM运算，并采用CUDA语言编写。在安装过程中无需进行编译操作，而是通过运行时的轻量级即时编译模块动态生成所有内核代码。

目前，DeepGEMM仅适用于英伟达Hopper架构的硬件环境。为解决FP8张量核心累加过程中可能出现的精度问题，该库采用了两级累加方法以充分利用CUDA核心性能。此外，DeepGEMM的设计极其精简，整个代码库的核心功能集中在一个内核函数中，总代码量约为300行。

尽管体积小巧，但DeepGEMM在不同矩阵形状下的计算性能表现优异，能够与经过深度优化的专业库相媲美甚至超越。根据DeepSeek团队的测试数据，在H800 GPU上使用NVCC 12.8编译器对可能用于DeepSeek-V3/R1推理的各种矩阵形状（包括预填充和解码阶段，但不涉及张量并行）进行了全面评估。测试结果显示，DeepGEMM的计算性能最高可达1358 TFLOPS，内存带宽峰值为2668 GB/s。相比基于CUTLASS 3.6的优化实现，其性能提升幅度最高可达2.7倍。在分组GEMM（MoE模型）的连续性布局和掩码布局下，DeepGEMM的性能提升也可达到1.2倍以上。

需要注意的是，使用DeepGEMM需要满足以下环境要求：硬件需支持Hopper架构（sm_90a），操作系统需安装Python 3.8或更高版本，CUDA版本需为12.3及以上（推荐使用12.8），同时需配备PyTorch 2.1及以上版本以及CUTLASS 3.6及以上版本。