金磊 发自 凹非寺

量子位 | 公众号 QbitAI

开源周的第三天， DeepSeek把训练推理V3/R1背后的“动力”给亮出来了——

DeepGEMM：一个FP8 GEMM（通用矩阵乘法）库，支持密集（dense）和混合专家（MoE）矩阵乘法运算。

我们先来简单了解一下GEMM。

GEMM，即通用矩阵乘法，是线性代数中的基本运算，是科学计算、机器学习、深度学习等领域中“常客”，也是许多高性能计算任务的核心。

但由于它的计算量往往都比较大，所以GEMM的性能优化是至关重要的一点。

而DeepSeek这次开源的DeepGEMM，依旧是保持了 “高性能+低成本”的特性，亮点如下：

• 高性能：在Hopper架构的GPU上，DeepGEMM能够实现高达1350+FP8 TFLOPS的性能。

• 简洁性：核心逻辑仅约 300 行代码，但性能却优于专家调优的内核。

• 即时编译（JIT）：采用完全即时编译的方式，这意味着它可以在运行时动态生成优化的代码，从而适应不同的硬件和矩阵大小。

• 无重依赖：这个库设计得非常轻量级，没有复杂的依赖关系，可以让部署和使用变得简单。

• 支持多种矩阵布局：支持密集矩阵布局和两种 MoE 布局，这使得它能够适应不同的应用场景，包括但不限于深度学习中的混合专家模型。

简单来说，DeepGEMM主要用于加速深度学习中的矩阵运算，特别是在大规模模型训练和推理中，它特别适用于需要高效计算资源的场景，能够显著提升计算效率。

很多网友们对这次的开源都比较“买单”，有人将DeepGEMM比作数学界的超级英雄，认为它比飞快的计算器还要快，比多项式方程还要强大。

也有人将DeepGEMM的发布比喻为量子态稳定到一个新的现实，称赞其即时编译的干净利落。

当然……也有人开始担心起自己手上的英伟达股票了……

深入了解DeepGEMM

DeepGEMM是一个专门为实现简洁高效的FP8通用矩阵乘法（GEMMs）而打造的库，它还具备细粒度缩放功能，这一设计源于DeepSeek V3。

它既能处理普通的通用矩阵乘法，也能支持MoE分组的通用矩阵乘法。

这个库是用CUDA编写的，安装的时候不需要编译，因为它会在运行时通过一个轻量级的即时编译（JIT）模块来编译所有的内核程序。

目前，DeepGEMM只支持英伟达的Hopper张量核心。

为了解决FP8张量核心在计算累积时不够精确的问题，它采用了CUDA核心的两级累积（提升）方法。

虽然DeepGEMM借鉴了CUTLASS和CuTe里的一些理念，但并没有过度依赖它们的模板或代数运算。

相反，这个库设计得很简洁，只有一个核心内核函数，代码量大概300行左右。

这使得它成为一个简洁易懂的资源，方便大家学习Hopper架构下的FP8矩阵乘法和优化技术。

尽管其设计轻巧，但DeepGEMM的性能可以匹配或超过各种矩阵形状的专家调优库。

那么具体性能如何呢？

团队在H800上使用NVCC 12.8测试了DeepSeek-V3/R1推理中可能使用的所有形状（包括预填充和解码，但没有张量并行）。

下面这张图展示的是用于密集模型的普通DeepGEMM的性能：

从测试结果来看，DeepGEMM 计算性能最高可达1358 TFLOPS， 内存宽带最高可达2668 GB/s。

加速比方面，与基于CUTLASS 3.6的优化实现相比，最高可达2.7倍。

再来看下DeepGEMM支持MoE模型的 连续布局（contiguous layout）的性能：

以及支持MoE模型 掩码布局（masked layout）的性能是这样的：

如何使用？

要想使用DeepGEMM，需先注意一下几个 依赖项，包括：

• 必须支持Hopper架构的GPU，sm_90a。

• Python 3.8及以上。

• CUDA 12.3及以上（推荐12.8）。

• PyTorch 2.1及以上。

• CUTLASS 3.6及以上

Development代码如下：