金磊发自凹非寺
量子位公众号 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 代码如下：
# Submodule must be cloned git clone --recursive git@github.com:deepseek-ai/DeepGEMM.git # Make symbolic links for third-party (CUTLASS and CuTe) include directories python setup.py develop # Test JIT compilation python tests/test_jit.py # Test all GEMM implements (normal, contiguous-grouped and masked-grouped) python tests/test_core.py
安装代码如下：
python setup.py install
在上述步骤之后，您的 Python 项目中导入 deep_gemm 即可。
接口方面，对于普通的 DeepGEMM，可调用 deep_gemm.gemm_fp8_fp8_bf16_nt 函数，支持 NT 格式（非转置 LHS 和转置 RHS）。
对于分组的 DeepGEMM，连续布局情况下是m_grouped_gemm_fp8_fp8_bf16_nt_contiguous；掩码布局情况下是m_grouped_gemm_fp8_fp8_bf16_nt_masked。
DeepGEMM 还提供设置最大 SM 数量、获取 TMA 对齐大小等工具函数；支持环境变量，如 DG_NVCC_COMPILER、DG_JIT_DEBUG 等。
除此之外，DeepSeek 团队还提供了几种优化的方式，包括：
- JIT 设计：所有内核在运行时编译，无需安装时编译；支持动态选择最优块大小和流水线阶段。
- 细粒度缩放：通过 CUDA 核心两层累加解决 FP8 精度问题；支持非 2 的幂次方块大小，优化 SM 利用率。
- FFMA SASS 交错：通过修改 SASS 指令的 yield 和 reuse 位，提高性能。
感兴趣的小伙伴可以戳文末 GitHub 链接查看详情哦~
One More Thing
英伟达这几天的股票……嗯……一直再跌：
不过在北京时间 27 日凌晨，英伟达 2025 财年第四季度业绩报告也即将出炉，我们可以期待一下它的表现~
DeepGEMM GitHub 入口：
