内容目录

DeepSeek开源周Day3 发布！DeepGEMM，300行核心代码，加速V3、R1训练推理效率

在DeepSeek的OpenSourceWeek的第三天的开源项目出来了：

介绍内容：

#开源周第三天：DeepGEMM

隆重推出 DeepGEMM —— 一个支持密集矩阵与混合专家（MoE）矩阵乘法的 FP8 GEMM 库，为 V3/R1 的训练与推理提供强劲动力。

⚡ 在 Hopper GPU 上实现高达 1350+ FP8 TFLOPS 的性能
✅ 无繁重依赖，简洁如教程
✅ 完全即时编译（Just-In-Time）
✅ 核心逻辑仅约 300 行代码 —— 却在多数矩阵规模上超越专家调优的内核
✅ 支持密集矩阵布局及两种 MoE 布局

GitHub 地址：https://github.com/deepseek-ai/DeepGEMM

DeepGEMM最大的特点是支持FP8和MOE，是DeepSeek-V3的核心技术，这次DeepSeek开源这么核心的模块，可以说对开源的支持是非常给力的了！！！

详解DeepGEMM与DeepSeek-V3

GEMM（General Matrix Multiplication，通用矩阵乘法）在 DeepSeek-V3 中扮演着至关重要的角色。作为一种基础的矩阵运算，GEMM 在深度学习模型中广泛应用于数据转换、权重应用等任务，尤其是在大型语言模型（LLM）中，其重要性不言而喻。DeepSeek-V3 作为一款基于 Transformer 架构的混合专家（Mixture of Experts, MoE）模型，GEMM 在其多个关键组件中发挥着核心作用。以下是对 GEMM 在 DeepSeek-V3 中作用的详细分析：

MoE 架构中的 GEMM

DeepSeek-V3 的核心架构是 DeepSeekMoE，一种混合专家（MoE）设计。在 MoE 模型中，计算过程分为多个“专家”（Experts），每个专家通常是一个前馈神经网络（Feed-Forward Network, FFN）。GEMM 在这些 FFN 中用于执行线性变换，即输入向量与权重矩阵的乘法，这是每个专家处理输入的主要步骤。

此外，MoE 架构中的门控机制（Gating Mechanism）负责根据输入动态选择激活哪些专家。这一过程通常涉及计算每个专家的得分，可能需要将输入与门控权重进行矩阵乘法，或者对专家输出的结果进行加权组合。这些操作同样依赖于 GEMM，使得门控机制能够高效地分配计算资源。因此，GEMM 在 DeepSeekMoE 的专家网络和门控机制中都是不可或缺的。

FP8 精度优化与 DeepGEMM 库

DeepSeek-V3 在训练和推理过程中使用了 FP8（8 位浮点数）精度，以加速计算并减少内存占用。为此，DeepSeek 团队开发了 DeepGEMM 库，这是一个专门为 FP8 GEMM 操作设计的开源库，特别针对 NVIDIA Hopper GPU 进行了优化。Hopper GPU 配备了 FP8 张量核心（Tensor Cores），能够显著加速低精度矩阵乘法。

DeepGEMM 的设计解决了 FP8 精度下的一些挑战。例如，FP8 张量核心在累积计算时的精度可能不足，而 DeepGEMM 通过 CUDA 核心的两级累积（Two-level Accumulation）策略弥补了这一问题。这种优化使得 DeepSeek-V3 能够在 FP8 精度下高效执行 GEMM 操作，同时保持数值稳定性。通过 DeepGEMM，DeepSeek-V3 的矩阵乘法性能得到了显著提升，尤其是在现代 GPU 上。

这么核心的实现都开源了，看来这次真如DeepSeek所说，一直保持着车库精神！！！

特点分析：DeepGEMM的五大“致命武器”

1️⃣ 极简代码，极致性能

300行核心代码的暴力美学：抛弃臃肿模板，直击MMA指令本质，代码量堪比教学示例，性能却碾压传统库。
JIT编译黑科技：运行时动态优化，让每个GEMM形状都拥有“定制化内核”，小矩阵性能提升高达2.7倍！

2️⃣ FP8精度救星：CUDA核心两级累加

打破张量核心精度枷锁：通过CUDA核心的FP32中间累加，解决Hopper FP8张量核心的误差累积问题，兼顾速度与数值稳定性。
细粒度缩放因子：每个矩阵块独立量化，让MoE模型中的专家权重动态适应计算需求。

3️⃣ MoE优化的“双生暗器”

连续布局：预填充阶段暴力拼接token，TMA对齐实现零浪费数据传输；
掩码布局：解码阶段动态激活专家，避免CPU-GPU通信瓶颈，CUDA Graph友好性拉满！

4️⃣ 硬件级“手术刀”优化

TMA异步猛兽：Tensor Memory Accelerator多播+预取，数据搬运效率提升300%；
FFMA指令级调教：修改SASS二进制码，让FFMA指令交错执行，关键场景性能再涨10%；
非对齐块暴走：打破2^N块大小魔咒，SM利用率直冲100%！

5️⃣ 开发者友好革命

零依赖极简部署：pip install即用，告别CUTLASS生态的“依赖地狱”；
透明调试接口：环境变量一键开启PTXAS反汇编、寄存器复用分析，硬核玩家调试利器；
PyTorch无缝衔接：deep_gemm模块即插即用，AI工程师的“瑞士军刀”。

为什么开发者应该尖叫？

当其他库还在比拼“代码重量”时，DeepGEMM用1/100的代码量实现了2倍性能——这不仅是技术突破，更是一场开发范式的革命。它证明：

高性能≠复杂：删繁就简的架构设计，让核心逻辑如水晶般透明；
硬件潜力未耗尽：通过SASS级微操，Hopper GPU的隐藏性能被彻底榨干；
开源社区的逆袭：轻量级代码+MIT协议，个人开发者也能挑战科技巨头的优化霸权！

无论你是苦于MoE模型部署的AI工程师，还是渴望理解FP8极限优化的CUDA极客，DeepGEMM都将为你打开一扇新世界的大门。让我们一起，亲历这场代码的性能核爆！

附录：项目详细介绍

为了方便大家了解项目，我对项目介绍进行了翻译，内容如下：

DeepGEMM

DeepGEMM 是一个专为简洁高效的 FP8 通用矩阵乘法（GEMM）设计的库，支持细粒度缩放，如 DeepSeek-V3 所提出的。它支持普通矩阵乘法和混合专家（MoE）分组矩阵乘法。该库使用 CUDA 编写，安装时无需编译，所有内核在运行时通过轻量级的即时编译（JIT）模块进行编译。

目前，DeepGEMM 仅支持 NVIDIA Hopper 张量核心。为了解决 FP8 张量核心累加不精确的问题，它采用了 CUDA 核心的两级累加（提升）。虽然它借鉴了 CUTLASS 和 CuTe 的一些概念，但避免了对它们模板或代数的重度依赖。相反，该库设计简洁，核心内核函数仅约 300 行代码，使其成为学习 Hopper FP8 矩阵乘法和优化技术的清晰且易于理解的资源。

尽管设计轻量，DeepGEMM 在各种矩阵形状下的性能与专家调优的库相当甚至更优。

性能

我们在 H800 GPU 上测试了 DeepSeek-V3/R1 推理中可能使用的所有形状（包括预填充和解码阶段，但不包括张量并行），使用 NVCC 12.8。所有加速指标均基于我们内部基于 CUTLASS 3.6 精心优化的实现进行比较。

DeepGEMM 在某些形状上表现不佳，欢迎感兴趣的开发者提交优化 PR。

密集模型的普通 GEMM

M	N	K	计算性能	内存带宽	加速比
64	2112	7168	206 TFLOPS	1688 GB/s	2.7x
64	24576	1536	289 TFLOPS	2455 GB/s	1.7x
64	32768	512	219 TFLOPS	2143 GB/s	1.8x
64	7168	16384	336 TFLOPS	2668 GB/s	1.4x
64	4096	7168	287 TFLOPS	2320 GB/s	1.4x
64	7168	2048	295 TFLOPS	2470 GB/s	1.7x
128	2112	7168	352 TFLOPS	1509 GB/s	2.4x
128	24576	1536	535 TFLOPS	2448 GB/s	1.6x
128	32768	512	358 TFLOPS	2103 GB/s	1.5x
128	7168	16384	645 TFLOPS	2604 GB/s	1.4x
128	4096	7168	533 TFLOPS	2221 GB/s	2.0x
128	7168	2048	510 TFLOPS	2277 GB/s	1.7x
4096	2112	7168	1058 TFLOPS	527 GB/s	1.1x
4096	24576	1536	990 TFLOPS	786 GB/s	1.0x
4096	32768	512	590 TFLOPS	1232 GB/s	1.0x
4096	7168	16384	1358 TFLOPS	343 GB/s	1.2x
4096	4096	7168	1304 TFLOPS	500 GB/s	1.1x
4096	7168	2048	1025 TFLOPS	697 GB/s	1.1x

MoE 模型的分组 GEMM（连续布局）

#Groups	每组M	N	K	计算性能	内存带宽	加速比
4	8192	4096	7168	1297 TFLOPS	418 GB/s	1.2x
4	8192	7168	2048	1099 TFLOPS	681 GB/s	1.2x
8	4096	4096	7168	1288 TFLOPS	494 GB/s	1.2x
8	4096	7168	2048	1093 TFLOPS	743 GB/s	1.1x

MoE 模型的分组 GEMM（掩码布局）

#Groups	每组 M	N	K	计算性能	内存带宽	加速比
1	1024	4096	7168	1233 TFLOPS	924 GB/s	1.2x
1	1024	7168	2048	925 TFLOPS	968 GB/s	1.2x
2	512	4096	7168	1040 TFLOPS	1288 GB/s	1.2x
2	512	7168	2048	916 TFLOPS	1405 GB/s	1.2x
4	256	4096	7168	932 TFLOPS	2064 GB/s	1.1x
4	256	7168	2048	815 TFLOPS	2047 GB/s	1.2x

快速开始

要求

Hopper 架构 GPU，必须支持 sm_90a
Python 3.8 或更高版本
CUDA 12.3 或更高版本（推荐 12.8 或更高以获得最佳性能）
PyTorch 2.1 或更高版本
CUTLASS 3.6 或更高版本（可通过 Git 子模块克隆）

开发

# 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，尽情使用吧！

接口

注意事项

该库仅包含 GEMM 内核。它要求 LHS 缩放因子为 TMA 对齐并转置，且仅支持 NT 格式（非转置的 LHS 和转置的 RHS）。对于转置或其他 FP8 转换操作，请独立实现或将它们融合到先前的内核中。虽然该库提供了一些简单的 PyTorch 工具函数，但这些可能会导致性能下降，我们的主要重点是优化 GEMM 内核本身。

普通密集 GEMM（非分组）

要执行基本的非分组 FP8 GEMM，请调用 deep_gemm.gemm_fp8_fp8_bf16_nt 函数。更多详情请参阅函数文档。

分组 GEMM（连续布局）

与 CUTLASS 中的传统分组 GEMM 不同，DeepGEMM 仅对 M 轴进行分组，而 N 和 K 必须保持不变。这种设计适用于 MoE 模型中专家共享相同形状的场景。

在训练前向传递或推理预填充阶段，每个专家可能处理不同数量的 token，我们将这些 token 连接成一个张量，称为“连续”布局。注意，每个专家段必须对齐到 GEMM M 块大小（get_m_alignment_for_contiguous_layout()）。

更多信息请参阅 m_grouped_gemm_fp8_fp8_bf16_nt_contiguous 函数文档。

分组 GEMM（掩码布局）

在推理解码阶段，当启用 CUDA 图且 CPU 不知道每个专家接收的 token 数量时，我们支持掩码分组 GEMM。通过提供掩码张量，内核仅计算有效部分。

使用 m_grouped_gemm_fp8_fp8_bf16_nt_masked 并参阅相关文档。示例用法是将 DeepEP 的低延迟内核输出作为输入。

工具函数

除了上述内核外，该库还提供了一些工具函数：

deep_gemm.set_num_sms：设置要使用的最大 SM 数量
deep_gemm.get_num_sms：获取当前 SM 最大数量
deep_gemm.get_m_alignment_for_contiguous_layout：获取分组连续布局的组级对齐要求
deep_gemm.get_tma_aligned_size：获取所需的 TMA 对齐大小
deep_gemm.get_col_major_tma_aligned_tensor：获取列主序的 TMA 对齐张量

该库还提供了一些环境变量，可能有用：

DG_CACHE_DIR：字符串，存储编译内核的缓存目录，默认为 $HOME/.deep_gemm
DG_NVCC_COMPILER：字符串，指定的 NVCC 编译器路径；默认从 torch.utils.cpp_extension.CUDA_HOME 中查找
DG_DISABLE_FFMA_INTERLEAVE：0 或 1，禁用 FFMA 交错优化
DG_PTXAS_VERBOSE：0 或 1，显示详细的 PTXAS 编译器输出
DG_PRINT_REG_REUSE：0 或 1，打印 FFMA 交错细节
DG_JIT_PRINT_NVCC_COMMAND：0 或 1，打印 NVCC 编译命令
DG_JIT_DEBUG：0 或 1，打印更多调试信息

更多示例和详细信息，请参阅测试代码或查看相应的 Python 文档。

优化

我们标记了 CUTLASS 中未使用的技术为。

持久化 Warp 专用化

遵循 CUTLASS 设计，DeepGEMM 中的内核是 warp 专用化的，能够重叠数据移动、张量核心 MMA 指令和 CUDA 核心提升。下图展示了这一过程的简化示意图：

Hopper TMA 特性

张量内存加速器（TMA）是 Hopper 架构引入的新硬件特性，旨在实现更快且异步的数据移动。具体来说，我们利用 TMA 进行：

LHS、LHS 缩放因子和 RHS 矩阵的 TMA 加载
输出矩阵的 TMA 存储
LHS 矩阵的 TMA 多播
TMA 描述符预取

常见细节优化

使用 stmatrix PTX 指令
针对不同 warp 组调整寄存器数量
尽可能重叠操作，例如重叠 TMA 存储和非 TMA RHS 缩放因子加载
统一且优化的块调度器
所有非分组和分组内核使用一个调度器
光栅化以增强 L2 缓存重用

完全 JIT 设计

DeepGEMM 采用完全即时编译（JIT）设计，安装时无需编译。所有内核在运行时通过轻量级 JIT 实现进行编译。这种方法有以下优势：

GEMM 形状、块大小和流水线阶段数被视为编译时常量
节省寄存器
编译器可能进行更多优化
自动选择块大小、warp 组数量、最佳流水线阶段和 TMA 集群大小
但不进行自动调优，而是确定性地选择最佳方案
完全展开 MMA 流水线，为编译器提供更多优化机会
对小形状非常重要
详情请参阅内核文件中的 launch_k_iterations

总体而言，JIT 显著提升了小形状的性能，类似于 Triton 编译器的方法。

非对齐块大小

对于某些形状，对齐到 2 的幂的块大小可能导致 SM 利用率不足。例如，对于 M=256、N=7168，典型的块大小分配为 BLOCK_M=128、BLOCK_N=128，结果只有 (256 / 128) (7168 / 128) = 112 个 SM 被利用。为了解决这个问题，我们支持非对齐块大小，如 112，使得 (256 / 128) (7168 / 112) = 128 个 SM 在这种情况下工作。结合细粒度缩放，实现这一技术需要仔细优化，但最终带来了性能提升。

FFMA SASS 交错

我们观察到 CUTLASS FP8 内核在 NVCC 12.2 和 12.3 之间的性能提升。通过比较编译后的 SASS，我们发现一系列 FADD 指令中的一个位以交错模式翻转。参考一些开源的 CUDA 汇编器实现后，我们确定该位控制 yield，这可能会增强 warp 级并行性（仅为猜测，yield 当前 warp 并让其他 warp 工作）。

为了利用这一点，我们开发了一个类似的脚本来修改编译后的二进制文件中的 FFMA 指令。除了简单地修改 yield 位外，我们还翻转了重用位（如果 warp 被 yield，寄存器无法重用）。这一调整通过创建更多机会来重叠 MMA 指令和提升 FFMA 指令，提高了细粒度缩放 FP8 GEMM 的性能（在某些情况下提升 10% 以上）。