DeepSeek开源DeepGEMM：FP8矩阵乘法新突破

摘要： DeepSeek近日开源了其FP8通用矩阵乘法库DeepGEMM，该库专为NVIDIA Hopper架构设计，旨在高效简洁地执行FP8矩阵乘法（GEMM）操作。DeepGEMM支持普通和混合专家（MoE）分组的GEMM，采用即时编译（JIT）技术，并利用Hopper的Tensor Memory Accelerator（TMA）特性，显著提升数据传输效率和计算精度。该库核心代码仅约300行，易于学习和优化，在多种矩阵形状上均达到或超过专家级优化库的性能。

北京，[日期] – 在人工智能领域，模型规模的不断扩大对计算性能提出了更高的要求。为了应对这一挑战，DeepSeek公司宣布开源其FP8通用矩阵乘法库DeepGEMM，为AI研究人员和开发者提供了一个高效、灵活的工具，以加速大规模AI模型推理和混合专家（MoE）模型的训练。

FP8：低精度计算的未来

随着AI模型复杂性的增加，对计算资源的需求也呈指数级增长。FP8（8位浮点）作为一种低精度计算格式，能够在降低内存占用和提高计算吞吐量的同时，保持可接受的精度水平，因此受到了广泛关注。然而，FP8的精度限制也给实际应用带来了挑战。

DeepGEMM：解决精度难题，优化Hopper架构

DeepGEMM正是为了解决FP8精度不足的问题而生。它采用了以下关键技术：

细粒度缩放： 通过对输入数据进行细粒度缩放，DeepGEMM能够充分利用FP8的动态范围，提高计算精度。
CUDA核心双级累加： DeepGEMM在CUDA核心内部使用双级累加机制，将FP8计算结果提升到更高精度的格式（如BF16），进一步确保计算精度。
即时编译（JIT）： DeepGEMM基于JIT技术，在运行时动态编译内核，根据矩阵形状、块大小等参数进行优化，无需安装时编译，节省寄存器并提升性能。
Hopper架构优化： DeepGEMM专为NVIDIA Hopper架构设计，充分利用Tensor Memory Accelerator（TMA）特性，包括TMA加载、存储、多播和描述符预取，显著提升数据传输效率。

性能卓越，应用广泛

DeepGEMM的性能表现令人印象深刻。在普通GEMM（非分组）操作中，DeepGEMM在某些矩阵形状下能达到2.7倍的加速比，计算性能超过1000 TFLOPS，接近Hopper架构GPU的理论峰值。在分组GEMM（MoE模型）中，DeepGEMM也能实现1.1-1.2倍的加速比，显著提升MoE模型的训练和推理效率。

DeepGEMM的应用场景包括：

大规模AI模型推理： 加速高维矩阵乘法，提升推理速度。
混合专家（MoE）模型： 优化分组矩阵乘法，提高计算效率。
低精度计算： 通过细粒度缩放解决FP8精度问题，保持高精度输出。
高性能计算： 基于Hopper架构特性，提升矩阵运算效率。
深度学习框架优化： 作为底层优化库，加速模型训练和推理。

轻量级设计，易于学习和扩展

DeepGEMM的核心代码仅约300行，避免了复杂模板或代数结构的依赖，降低了学习和优化的难度。这使得DeepGEMM不仅易于集成到现有项目中，也为研究人员提供了一个理想的平台，用于探索和优化FP8矩阵乘法。

系统要求

要使用DeepGEMM，需要满足以下系统要求：

硬件： 支持NVIDIA Hopper架构（sm_90a）的GPU，如H800或H100。
CUDA： CUDA 12.3或更高版本（推荐12.8或更高版本）。
软件： Linux操作系统（如Ubuntu、CentOS等），Python 3.8或更高版本，PyTorch 2.1或更高版本，CUTLASS 3.6或更高版本。

开源地址

DeepGEMM的项目地址为：https://github.com/deepseek-ai/DeepGEMM

结论

DeepSeek开源DeepGEMM，为FP8矩阵乘法提供了一个高效、灵活的解决方案，有望加速AI模型推理和MoE模型训练，推动人工智能技术的发展。DeepGEMM的轻量级设计和卓越性能，使其成为AI研究人员和开发者的理想选择。

参考文献