【有啥问啥】DeepSeek开源DeepGEMM：释放FP8矩阵乘法加速的潜力

DeepGEMM

DeepSeek开源DeepGEMM：释放FP8矩阵乘法加速的潜力

在人工智能和深度学习领域，通用矩阵乘法（GEMM）是模型训练和推理的核心计算操作。随着模型规模的不断扩大，对GEMM运算效率的需求也日益增长。为了应对这一挑战，DeepSeek 公司开源了 DeepGEMM，一个专为高效FP8（8位浮点）通用矩阵乘法设计的库。本文将深入探讨 DeepGEMM 的技术特性、优势以及应用前景，帮助读者全面了解这一强大的加速工具。

什么是GEMM？为何重要？

通用矩阵乘法（GEMM）指的是计算如下形式的矩阵运算：

$C = α * o p (A) @ o p (B) + β * C$

其中：

$A$ 和 $B$ 是输入矩阵。
$C$ 是输出矩阵。
$α$ 和 $β$ 是标量系数。
$o p ()$ 代表可选的矩阵操作，例如转置或共轭转置。
$@$ 表示矩阵乘法。

GEMM 广泛应用于各种计算密集型应用，尤其是在深度学习领域，它构成了卷积神经网络（CNN）、循环神经网络（RNN）、Transformer 模型等的核心计算层。因此，GEMM 的效率直接影响着 AI 模型的训练和推理速度。

DeepGEMM：专为FP8而生的高效矩阵乘法库

DeepGEMM 是 DeepSeek 推出的一款开源库，专注于实现高效的 FP8 通用矩阵乘法。其设计目标是为 DeepSeek V3 及其他 AI 模型的训练和推理任务提供强大的加速支持。

核心特性和优势：

极致的FP8性能: DeepGEMM 采用 FP8 数据类型进行计算。FP8 是一种低精度浮点格式，相比传统的 FP32 或 FP16，FP8 具有以下显著优势：
- 更高的计算吞吐量: FP8 运算能够充分利用现代 GPU (尤其是 NVIDIA Hopper 架构) 上的 Tensor Core 加速单元，实现更高的计算速度。
- 更低的内存占用: FP8 数据格式仅占用 8 位，显著减少了内存带宽需求和存储空间，从而降低了硬件成本，并提升了数据传输效率。
细粒度缩放 (Fine-grained Scaling): DeepGEMM 采用了 DeepSeek-V3 中提出的细粒度缩放技术。这项技术能够更精细地控制 FP8 计算过程中的数值范围，最大程度地保留计算精度，同时避免溢出或下溢问题，确保 FP8 计算的稳定性和可靠性。
支持普通 GEMM 和混合专家 (MoE) 分组 GEMM: DeepGEMM 不仅支持标准的矩阵乘法，还针对混合专家模型 (MoE) 进行了优化，能够高效处理 MoE 模型中特有的分组 GEMM 运算。MoE 模型是近年来大型模型发展的重要趋势，DeepGEMM 对 MoE 的支持使其能够更好地服务于未来更大规模的 AI 模型。
CUDA 编写，即时编译 (JIT): DeepGEMM 使用 CUDA 语言编写，充分利用 NVIDIA GPU 的硬件特性。更重要的是，DeepGEMM 采用轻量级的即时编译 (JIT) 模块，在运行时动态编译所有内核代码。这种方式 无需在安装时进行预编译，极大地简化了部署流程，并提高了灵活性和兼容性。
优异的性能表现: 根据 DeepSeek 官方的测试数据，DeepGEMM 相比基于 CUTLASS 3.6 的优化实现，在普通 GEMM (密集模型) 中最高可提速 2.7 倍。即使在分组 GEMM (MoE 模型) 的连续性布局和掩码布局下，DeepGEMM 的性能提升也 可达到 1.2 倍以上。实测计算性能 最高可达 1358 TFLOPS，内存带宽峰值 高达 2668 GB/s。
两级累加 (Two-level Accumulation): 为了解决 FP8 张量核在累加计算时可能出现的精度问题，DeepGEMM 采用了基于 CUDA Core 的两级累加技术。这种技术能够在保证计算精度的前提下，充分发挥 FP8 的性能优势。
简洁高效的代码: DeepGEMM 的代码库设计 简洁而高效。核心代码仅有 300 行左右，但却实现了专家级优化的内核性能。这不仅降低了维护成本，也方便开发者进行二次开发和定制。

DeepGEMM 的应用场景

DeepGEMM 作为一款高性能的矩阵乘法库，可以广泛应用于以下场景：

深度学习模型训练: 加速各种深度学习模型的训练过程，特别是计算密集型的模型，如大型语言模型、视觉 Transformer 模型等。
深度学习模型推理: 提升模型推理速度，降低延迟，尤其适用于对实时性要求较高的应用场景。
科学计算: 在需要进行大规模矩阵运算的科学计算领域，例如物理模拟、气象预测、生物信息学等，DeepGEMM 也能发挥重要作用。
高性能计算 (HPC): 作为高性能计算的基础组件，DeepGEMM 可以为各种 HPC 应用提供强大的矩阵运算加速能力。

总结与展望

DeepSeek 开源 DeepGEMM，无疑为 AI 社区贡献了一个宝贵的加速工具。它充分利用了 FP8 数据类型的优势，并结合细粒度缩放、JIT 编译等先进技术，实现了极致的矩阵乘法性能。DeepGEMM 的开源，不仅能够帮助研究人员和开发者 更高效地训练和部署 AI 模型，也将 推动 FP8 等低精度计算技术在 AI 领域的普及和应用。

DeepSeek 也表示，DeepGEMM 在某些特定形状的矩阵运算上可能还存在优化空间，并 欢迎开发者社区积极参与，共同优化和完善 DeepGEMM。相信在社区的共同努力下，DeepGEMM 将会变得更加强大，为 AI 技术的进步贡献更大的力量。