DeepSeek开源周 Day03：从DeepGEMM看大模型算力提速的矩阵乘法

在deepseek开源周的第三天，deepgemm库被隆重推出。这是一个专注于fp8精度的通用矩阵乘法的高性能库，它不仅以极致性能著称，还保持了令人惊讶的代码简洁性。

DeepGEMM因其轻量级设计和出色性能在加速大模型计算方面脱颖而出。该库不仅支持常规的密集矩阵乘法，还特别针对混合专家模型(MoE)场景下的分组GEMM计算进行了优化，成为DeepSeek-V3/R1训练和推理的核心动力。

在技术实现上，DeepGEMM采用了细粒度缩放的FP8计算，在Hopper GPU上可达到惊人的1350+ TFLOPS。更值得注意的是，它的核心逻辑仅约300行代码，却能在大多数矩阵尺寸上超越专家调优的内核性能。

DeepGEMM还具有无需预编译、完全即时(JIT)编译、无重度依赖等特点，使其既可作为生产环境中的高性能组件，也是学习现代GPU优化技术的理想教程。

矩阵乘法GEMM通用矩阵乘法 (GEMM) 是线性代数、机器学习、统计学和许多其他领域的常见算法。它提供了比上一个教程更有趣的权衡空间，因为有很多方法可以分解计算。这包括使用阻塞、内积、外积和脉动阵列技术。可以说，矩阵乘法是科学计算和深度学习中最广泛的运算符之一。

设

A in mathbb{R}^{n imes l}，和

B in mathbb{R}^{l imes m}，如果

C = AB，则

C in mathbb{R}^{n imes m}，并且

C_{i,j} = sum_{k=1}^{l} A_{i,k} B_{k,j}访问计算的元素

C_{i,j}如图所示。

C_{x,y}在矩阵乘法中的过程。

为什么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）的矩阵乘法。例如，输入矩阵 X in mathbb{R}^{n imes d} 通过线性变换生成

Q, K, V in mathbb{R}^{n imes d_k}，然后计算注意力得分

Attention(Q, K, V) = ext{softmax}left(frac{QK^T}{sqrt{d_k}} ight)V其中

QK^T 是一个 GEMM 操作。

前馈神经网络：FFN 通常包括两个线性层和非线性激活函数，如FFN(x) = max(0, xW_1 + b_1)W_2 + b_2其中

W_1, W_2 的矩阵乘法也是 GEMM。

对于 MoE（混合专家模型），其计算涉及多个专家网络的选择和组合。MoE 层通常包括一个门控网络（Gating Network）决定哪个专家处理输入，然后对专家输出进行加权求和。分组 GEMM 在这里尤为重要，例如，对于输入

X in mathbb{R}^{n imes d} 和专家权重

W_e in mathbb{R}^{d imes d_e}，计算

Y_e = XW_e是分组 GEMM 的典型场景。DeepGEMM 的优化使得这种计算在 MoE 场景下更加高效。