DeepSeek开源DeepGEMM矩阵乘法加速库，最快加速2.7倍

DeepSeek 开源活动来到了第三天，新项目如约而至。

此次开源库名为 DeepGEMM，GitHub 地址在文末。

DeepSeek 表示，这是一个支持密集和混合专家（MoE，Mixture of Experts）通用矩阵乘法（GEMM，General Matrix Multiplication）的 FP8 GEMM 库，为 V3/R1 的训练和推理提供支持。

在 Hopper GPU 上最高可达 1350+FP8 TFLOPS。其他优点包括：

没有过多的依赖，像教程一样简洁

完全即时编译

核心逻辑约为 300 行，但在大多数矩阵大小上均优于专家调优的内核

支持密集布局和两种 MoE 布局

据 DeepSeek 介绍，DeepGEMM 是一个专门设计的、干净且高效的工具库，它的核心任务是进行一种叫做 GEMM 的数学运算，这是 AI 模型训练和运行中非常常见的一种计算。

DeepGEMM 的特别之处在于，它使用了一种叫做 FP8 的超高效、低精度计算方式，能让运算速度更快，同时占用更少的内存。这种方式在 DeepSeek-V3 中被提出，并且支持精细的缩放调整（fine-grained scaling），让计算更灵活。

除了普通的矩阵乘法，DeepGEMM 还能处理混合专家矩阵乘法。

目前，DeepGEMM 只支持英伟达 Hopper 架构的张量核心。张量核心是 GPU 里的一种特殊硬件，专门用来加速矩阵运算。不过，Hopper 的张量核心在做 FP8 计算时，会出现累加（accumulation，就是把结果一点点加起来的过程）不够精确的问题。

为了解决这个问题，DeepGEMM 采用了两级累加（two-level accumulation）的办法。它借助 CUDA 核心来做更精确的累加，确保结果不会因为硬件限制而出错。

DeepGEMM 借鉴了一些来自 CUTLASS 和 CuTe 的概念。这两个也是 NVIDIA 的高性能矩阵运算库。不过，DeepGEMM 没有完全依赖它们的复杂模板或数学系统。

相反，该库的设计非常简单，只有一个核心内核函数，包含大约 300 行代码。这使其成为学习 Hopper FP8 矩阵乘法和优化技术的、干净且易于访问的资源。

尽管 DeepGEMM 设计轻量，但它的性能足以媲美那些由专家精心调优的库，甚至在某些矩阵形状（matrix shapes，指矩阵的大小和结构）上表现得更好。

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

从性能对比表格上可见，DeepGEMM 最高能加速 2.7 倍。

所有加速指标都是与 DeepSeek 基于 CUTLASS 3.6 的内部精心优化的实现进行比较计算的。

鉴于 DeepGEMM 在某些矩阵形状上的表现并不是很好，DeepSeek 也邀请各路大神帮助优化这个开源库。

在优化方面，DeepSeek 实现了一些 CUTLASS 设计之外的优化。这些也是 DeepGEMM 最具创新的地方。

首先是完全 JIT 设计，安装时无需编译。所有内核均使用轻量级 JIT 实现在运行时进行编译。这种方法的优点包括：

• GEMM 形状、块大小和管道阶段数被视为编译时常量

• 自动选择块大小、warpgroups 数量、最佳管道阶段和 TMA 集群大小

• 全面展开 MMA 管道，为编译器提供更多优化机会

其次是支持非对齐的块大小。在做矩阵乘法时，矩阵会被分成一个个小块（block）来计算。块的大小通常是固定的，比如 128×128（因为 128 是 2 的幂，计算机喜欢这种数字）。但有时候，矩阵的尺寸跟这些块大小不完全匹配，就会浪费计算资源。

举个例子：

假设矩阵的行数（M）是 256，列数（N）是 7168。如果用常见的块大小 BLOCK_M=128（行）和 BLOCK_N=128（列），那么行方向可以分成 256÷128=2 个块，列方向可以分成 7168÷128=56 个块。总共用到的 SM（GPU 的计算单元）是 2×56=112 个。

GPU 通常有更多 SM（比如 128 个），但这里只用了 112 个，剩下的就闲着了，没充分利用。

DeepGEMM 的办法是支持非对齐的块大小，也就是不一定非要是 2 的幂（比如 128），可以根据矩阵尺寸灵活调整。

还是上面的例子，如果把 BLOCK_N 改成 112（而不是 128，不是 2 的幂），行方向还是 256÷128=2 个块，列方向变成了 7168÷112=64 个块。总共用到的 SM 变成 2×64=128 个。

这下正好把 128 个 SM 全用上，没有浪费。

最后一点是，DeepSeek 通过观察和修改底层代码（SASS 汇编指令）实现了性能优化。

简单来说，DeepSeek 发现新版编译器（NVCC 12.3）里有些底层代码变了，性能变得更好。深入研究后，他们觉得这是因为线程调度（yield）的方式变了，能让 GPU 同时干更多工作。

于是，他们模仿这个变化，写了个脚本修改自己的代码（调整 FFMA 指令的某些位），让 MMA 指令（矩阵乘累加）和提升指令更高效地“重叠在一起执行”。

结果是 DeepGEMM 在某些场景下快了 10% 以上，特别适合那种需要灵活调整的 FP8 矩阵乘法。

最后，DeepSeek 致谢了 CUTLASS 项目，称其是 DeepGEMM 的灵感来源。

参考资料：

https://x.com/deepseek_ai/status/1894553164235640933

https://github.com/deepseek-ai/DeepGEMM

排版：刘雅坤