DeepGEMM简介
今天是DeepSeek开源周的第三天，继FlashMLA和DeepEP之后，DeepSeek开源了DeepGEMM库 。作为一个专注于FP8精度通用矩阵乘法的高性能库，DeepGEMM在提供极致性能的同时保持了令人惊讶的代码简洁性。
DeepGEMM以其轻量级设计和出色性能成为大模型计算加速领域的亮点。该库不仅支持常规的密集矩阵乘法，还特别优化了混合专家模型(MoE)场景下的分组GEMM计算，成为DeepSeek-V3/R1训练和推理的核心动力。
在技术实现上， DeepGEMM采用了细粒度缩放的FP8计算，在Hopper GPU上可达到惊人的1350+ TFLOPS。更令人印象深刻的是，它的核心逻辑仅约300行代码，却能在大多数矩阵尺寸上超越专家调优的内核性能。
MoE提速效果
DeepGEMM还具有无需预编译、完全即时(JIT)编译、无重度依赖等特点，使其既可作为生产环境中的高性能组件，也是学习现代GPU优化技术的理想教程。
DeepGEMM Github地址： https://github.com/deepseek-ai/DeepGEMM
矩阵乘法GEMM
通用矩阵乘法 (GEMM) 是线性代数、机器学习、统计学和许多其他领域的常见算法。它提供了比上一个教程更有趣的权衡空间，因为有很多方法可以分解计算。这包括使用阻塞、内积、外积和脉动阵列技术。 可以说，矩阵乘法是科学计算和深度学习中最广泛的运算符之一。
设
，和
，如果
，则
，并且
访问计算的元素
如图所示。
计算
在矩阵乘法中的过程。
为什么GEMM 是深度学习的核心
GEMM（General Matrix Multiply，通用矩阵乘法）是深度学习神经网络优化中的一个关键函数。自1979年作为BLAS（Basic Linear Algebra Subprograms，基本线性代数子程序）库的一部分首次提出以来，GEMM就成为了科学计算和高效矩阵运算的核心工具。 尽管许多人在开始深入研究神经网络优化时并不熟悉GEMM，但它却在深度学习的计算效率和能效提升中扮演着至关重要的角色。
GEMM的核心任务是执行矩阵乘法操作
，这在神经网络的训练过程中频繁出现，尤其是在全连接层和卷积层中。
神经网络的计算大多依赖于大量的矩阵乘法，因此优化这一操作就直接关系到网络的运行效率。
在大规模神经网络中，如何提高矩阵乘法的计算速度，减少功耗，已成为研究的重点之一。
为了更直观地理解GEMM的影响，我们可以参考贾扬清在其论文中的一张图表。该图表展示了在不同硬件平台上执行矩阵乘法时，GEMM优化带来的性能提升。通过高效实现GEMM，深度学习的计算资源可以被更充分地利用，从而显著提高训练和推理过程的速度，同时降低功耗。
这打破了使用 Alex Krizhevsky 的 Imagenet 架构进行图像识别的典型深度卷积神经网络的时间分配。所有以 fc（全连接）或 conv（卷积）开头的层都是使用 GEMM 实现的，几乎所有时间（GPU 版本的 95%，CPU 版本的 89%）都花在这些层上。
那么什么是 GEMM？它代表通用矩阵到矩阵乘法，它本质上就是按照字面意思来做，将两个输入矩阵相乘得到一个输出矩阵。它与在 3D 图形世界中习惯的矩阵运算的区别在于， 它处理的矩阵通常非常大 。例如，典型网络中的单个层可能需要将 256 行、1,152 列的矩阵与 1,152 行、192 列的矩阵相乘，以产生 256 行、192 列的结果。简单来说，这需要 5700 万次（256 x 1,152，x 192）浮点运算，现代架构中可能会有数十个这样的层，所以我经常看到网络需要数十亿次 FLOP 来计算单个帧。
全连接层
全连接层是经典的神经网络，已经存在了几十年，最简单的方法可能是先了解一下 GEMM 是如何应用于全连接层的。全连接层的每个输出值都会查看输入层中的每个值，将它们全部乘以该输入索引对应的权重，然后将结果相加以获得其输出。就上图而言，它看起来像这样：
有“k”个输入值，有“n”个神经元，每个神经元对每个输入值都有自己的一组学习权重。有“n”个输出值，每个神经元一个，通过对其权重和输入值进行点积计算得出。
卷积层
使用 GEMM 作为卷积层并不是一个显而易见的选择。卷积层将其输入视为二维图像，每个像素都有多个通道，就像具有宽度、高度和深度的经典图像一样。不过，与过去处理的图像不同，通道数量可以达到数百个，而不仅仅是 RGB 或 RGBA！
卷积运算通过获取多个权重“核”并将它们应用于整个图像来产生输出。输入图像和单个核如下所示：
每个核都是另一个三维数字数组，其深度与输入图像相同，但宽度和高度要小得多，通常为 7×7。为了产生结果，将核应用于输入图像上的点网格。在应用核的每个点上，所有相应的输入值和权重都会相乘，然后相加以产生该点的单个输出值。以下是视觉效果：
MoE 的矩阵运算
Transformer 是一种基于注意力机制的深度学习模型，其核心包括多头自注意力（Multi-Head Self-Attention, MHSA）和前馈神经网络（Feed-Forward Network, FFN）。这些层的计算大量依赖 GEMM，例如：
- 多头自注意力 ：计算过程包括查询（Query）、键（Key）和值（Value）的矩阵乘法。例如，输入矩阵
通过线性变换生成
，然后计算注意力得分
其中
是一个 GEMM 操作。
- 前馈神经网络 ：FFN 通常包括两个线性层和非线性激活函数，如
其中
的矩阵乘法也是 GEMM。
对于 MoE（混合专家模型），其计算涉及多个专家网络的选择和组合。MoE 层通常包括一个门控网络（Gating Network）决定哪个专家处理输入，然后对专家输出进行加权求和。分组 GEMM 在这里尤为重要，例如，对于输入
和专家权重
，计算
是分组 GEMM 的典型场景。DeepGEMM 的优化使得这种计算在 MoE 场景下更加高效。