DeepSeek 开源周（三）：DeepGEMM 性能实测及 NV 往事 - 文章 - 开发者社区

点击下方

卡片

，关注“

慢慢学AIGC

”

picture.image

DeepGEMM 简介

继 FlashMLA（DeepSeek 开源周（一）：FlashMLA 在 H100 上的性能实测）和 DeepEP（DeepSeek 开源周（二）：DeepSeek MoE 架构回顾和 DeepEP 性能实测）后，DeepSeek 开源周第三弹喜迎 DeepGEMM，一个 FP8 GEMM 库，支持稠密模型和 MoE 的 GEMM 计算，对 V3/R1 模型训练和推理均有帮助。

在 Hopper GPU 上高达 1350+ FP8 TFLOPS（硬件峰值 1979 TFLOPS 稠密算力，3959 TFLOPS 稀疏算力）；
无繁重的依赖，尽可能精简的代码结构；
完全即时（Just-In-Time，简称 JIT）编译；
核心逻辑仅 300 行代码，在大多数矩阵尺寸上的性能超过专家优化的内核；
支持稠密矩阵布局和两种 MoE 布局；

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

DeepGEMM 代码结构非常清晰：

picture.image

打开 deep_gemm/include/deep_gemm/mma_utils.cuh，PTX 内联汇编代码映入眼帘，为之前外网疯传的“DeepSeek 用 PTX 而不是 CUDA 编程，英伟达护城河不存在了”言论送上锤子（英伟达笑而不语）。

picture.image

PTX 所有指令集可以参考 Nvidia 官方文档：

https://docs.nvidia.com/cuda/parallel-thread-execution/

DeepGEMM 只包含 GEMM 核心，它要求 LHS 缩放因子为 TMA 对齐并且转置，且仅支持 NT 格式（非转置 LHS 和转置 RHS）。其他情况需要基于现有代码做二次开发。

普通密集 GEMM（非分组）

调用 deep_gemm.gemm_fp8_fp8_bf16_nt 函数执行基本的非分组 FP8 GEMM。

分组 GEMM（连续布局）

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

对于训练前向传递或推理预填充，其中每个专家可能处理不同数量的 tokens，我们将这些 tokens 连接成一个单一的张量，称为"连续"布局。请注意，每个专家段必须与 GEMM M 块大小对齐(get_m_alignment_for_contiguous_layout())。

分组 GEMM（掩码布局）

在启用 CUDA Graph 的推理解码阶段，当 CPU 不知道每个专家接收的 tokens 数量时，我们支持掩码分组 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 编译器路径；默认从from 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，打印更多调试信息

DeepGEMM 使用的一些优化技巧

DeepGEMM 借鉴了 Nvidia CUTLASS，站在巨人肩膀上做了进一步优化。下文用 🐳 标记了目前 CUTLASS 中未包含的技术。

持久性 warp-specialization

DeepGEMM 遵循 CUTLASS 设计，DeepGEMM 中的核心是 warp-specialization 的，使数据移动、张量核 MMA 指令和 CUDA 核心 Promotion 能够重叠。（笔者注：warp-specialization 没有找到很好的翻译，保留英文。通俗讲是让专业团队干专业的事，不要分心。对 warp，TMA，WGMMA 比较陌生的读者，可以阅读 Nvidia CUDA 编程手册查阅相关概念）。

下面是展示此过程的简化图：

picture.image

Hopper TMA 特性

张量内存加速器（Tensor Memory Accelerator，TMA）是 Hopper 架构引入的一项新硬件特性，专为更快速和异步的数据移动而设计。具体来说，DeepGEMM 利用 TMA 实现：

对 LHS、LHS 缩放因子和 RHS 矩阵进行 TMA 加载
对输出矩阵进行 TMA 存储
TMA 多播（仅限于 LHS 矩阵）
TMA 描述符预取

常见的细节优化

利用stmatrix PTX 指令
针对不同 warpgroup 的寄存器数量控制
尽可能多地重叠操作，例如重叠 TMA 存储和非 TMA RHS 缩放因子加载 🐳

统一且优化的块调度器

一个适用于所有非分组和分组内核的调度器
光栅化以增强 L2 缓存重用

完全 JIT 设计 🐳

DeepGEMM 采用完全即时编译（Just-In-Time，JIT）设计，安装时无需编译。所有内核在运行时使用轻量级 JIT 实现进行编译。这种方法提供了几个优势：

GEMM 形状、块大小和流水线阶段数被视为编译时常量
- 节省寄存器
- 编译器可以进行更多优化
自动选择块大小、warpgroup 数量、最佳流水线阶段和 TMA 集群大小
- 但无需自动调优，最佳选择是确定性的
MMA 流水线的完全展开，为编译器提供更多优化机会
- 对小形状非常重要
- 详情请参考内核文件中的launch\_k\_iterations

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

非对齐块大小 🐳

对于某些形状，与 2 的幂对齐的块大小可能导致 SM 利用率不足。例如，使用M=256, N=7168，典型的块大小分配BLOCK\_M=128, BLOCK\_N=128仅使用了 132 个SM中的(256 / 128) * (7168 / 128) = 112个。为解决这个问题，我们支持非对齐的块大小，如 112，使得在这种情况下可以有(256 / 128) * (7168 / 112) = 128个SM工作。将此技术与细粒度缩放一起实现需要谨慎优化，但最终会带来性能提升。

FFMA SASS 交错 🐳

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

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

H100 性能实测

DeepGEMM 非常轻量，安装和运行过程相对简单。


          
git clone https://github.com/deepseek-ai/DeepGEMM --recursive
          
cd DeepGEMM/
          
python setup.py develop
          
python tests/test_jit.py
          
python tests/test_core.py
          
# 安装
          
python setup.py install

笔者在 H100 上运行结果如下：

picture.image

和官方结果对比：

picture.image

部分结果低于官方结果，主要是因为笔者使用的软件环境版本较低（驱动 535.161.08，CUDA 12.6），而官方建议使用 CUDA 12.8 获得最好性能。

picture.image

NV 往事

看到 DeepGEMM FFMA SASS 交错部分文档时，勾起了一段对 NV 和开源社区之间的一段往事。

早在 2015 年，也就是 10 年前，有一位技术极客在 NV 开发者论坛发了一篇帖子，原贴找不到了，目前只能在他的开源项目 MaxAs Wiki 页面保留了完整内容。

链接：

https://github.com/NervanaSystems/maxas/wiki/Introduction

将翻译后的内容直接贴在这里：

MaxAs 项目简介

Scott Gray 于 2015 年 1 月 25 日编辑了此页面 · 4 次修订

我为 Nvidia Maxwell 架构编写了一个功能相当完善的汇编器。这个工具让你可以直接使用与 cuobjdump 输出相同的"sass"语言进行编码。这一切都始于我研究不同的 sgemm 实现，并尝试将这些技术融入到我正在开发的深度学习代码中。我基本上得出结论，用 Nvidia 提供的工具无法充分利用我购买的硬件。不幸的是，Nvidia 不相信"吃自己的狗粮"，他们手动汇编他们的库例程，而不是像我们其他人那样使用 ptxas。ptxas 在寄存器使用方面管理得很糟糕（尤其是当你使用向量内存操作时），在使用交错计算隐藏内存延迟方面做得很差（特别是如果你尝试对内存加载进行双缓冲），并且对某些谓词内存操作处理得很糟糕（即使是在 warp 均匀的情况下），以及其他问题。

无论如何，我越研究我代码的 sass 输出，我越开始意识到应该能够弄清楚我正在使用的所有指令的操作码和控制码，然后直接汇编我自己的代码。经过一个月左右的努力，我实现了这个目标，甚至更多。我现在发现直接用汇编器编码并与硬件交流要比用 cuda c 或 ptx 编码容易得多，不那么令人沮丧。

以下是我集成的主要功能（还有更多计划中的功能）：

寄存器分配：你可以在文件顶部通过变量名到寄存器号的映射来完成这项工作。这样你就可以编写易于理解的代码，而不会被所有寄存器号所掩盖。但主要是，这给予你对分配哪些寄存器的绝对控制权，并且零寄存器溢出。对于性能代码，这很重要，因为在硬件级别，寄存器是分组的，某些组合比其他组合提供更高的吞吐量（我说的是数百 GFlops）。为了帮助这一点，该工具可以自动分配寄存器以避免冲突，并通知你未避免的冲突，以便你可以手动调整映射。该工具还在汇编期间为你优化管理 .reuse 标志（Maxwell 和 cuda 6.5 的新功能）。这些标志进一步减少了 bank conflict 以及寄存器组带宽和整体芯片功耗。

调度块：对于大部分代码，你不想花太多时间优化指令的顺序和停顿。所以我编写了一个基本的调度器来为你完成这项工作。这样你就可以专注于编写清晰、易于维护的代码。但对于性能块代码，你可以选择不自动调度它们，而是非常小心地放置你的指令以最大化吞吐量。
元编程和宏：我在 Perl 中实现了这个汇编器，并将解释器嵌入为元编程工具。这允许你保持代码整洁地卷起，而不必维护无数指令。这使得感觉更像是在脚本语言中开发，而不是汇编。我还为像 XMAD 这样需要展开成多条指令的操作添加了汇编宏。
控制码：放置在调度块中的任何指令都会自动管理所需的停顿计数，以满足相关指令的管道深度。但是控制符号的其他方面我刻意不为你管理。这些主要是内存操作使用的依赖性屏障，用于指示数据何时可用。自动管理这些是一个困难的问题，我觉得最好留给开发者来解决。设置这些代码实际上为 gpu 编程增加了一个有趣的维度，而 cuda c 或 ptx 并没有暴露这一点。
反汇编：有时你只想稍微调整一下编译好的程序。这个工具使这变得非常容易。它可以以易于编辑的格式转储 cubin 代码，你可以直接将其插回去。事实上，该程序并不是设计为从头开始工作的。我没有花时间完全解析 cubin 格式，只是足够能够在适当的位置编辑内核。你至少需要从定义全局变量、共享内存和参数的内核框架开始。该工具会转储编译好的 cubin，然后你从那里接手。还有很多其他小功能可以讨论，你可以在 wiki 中探索它们。我用 Perl 编写了它，但我可能会在某个时候将其转换为 Python（这似乎是最终学习该语言的完美项目）。就目前而言，我发现现在编写性能在理论吞吐量 2% 以内的代码相当容易，对于 GM204 来说，理论吞吐量是 4.9 TFlops（默认时钟）。我从反复折腾 ptxas 中获得的最好结果是一个非常勉强的 75%。

目前的操作码覆盖率约为 80%。我可以将 cublas_device.lib 全部反汇编然后重新汇编，零错误。但仍有更多工作要做：更多操作码（主要是表面、纹理和视频指令）以及更多微基准测试来微调调度器。欢迎任何有兴趣贡献的人。

包含的是一个样例 sgemm 实现，运行速度达到 Maxwell 硬件理论吞吐量的 98%，比 Nvidia 手工汇编的 cublas 实现快达 4.8%。还包括一个简单的框架，用于编写微基准测试，目的是调整调度器。这里有一个简单的页面可以帮助你开始。更多文档和功能即将推出...

--Scott Gray

上面文章作者 Scott Gray 凭借 MaxAs 项目和针对小卷积核的 Winograd 快速算法（https://arxiv.org/pdf/1509.09308）在 CUDA 和深度学习社区具有深远影响力。

picture.image

大佬工作履历同样光鲜，2014 年加入 AI 芯片创业公司 Nervana 并担任主程序员，该公司于 2016 年被 Intel 以 3.5 亿美元收购，随后大佬加入 OpenAI 一直到现在。

picture.image

Scott 大佬在 OpenAI 的工作并没有太多曝光，笔者在 GPT-4 Tech Report 作者列表中看到大佬的身影。

picture.image

在论文最后看到大佬在做 Inference Research：

picture.image

相信 Scott 大佬随着 OpenAI 市值节节攀升早已实现财富自由，刚好可以做一些 research 相关的事情。

MaxAs 主要针对英伟达 Maxwell 架构，这个是 2015 年的架构，在新的架构上不通用。

比 MaxAs 更早的项目有 asfermi，主要针对 Nvidia Fermi 架构（CUDA Capability 2.0 和 2.1）实现了类似汇编器功能。

Github：https://github.com/hyqneuron/asfermi

asfermi 是一个用于 NVIDIA Fermi 指令集架构(ISA)的汇编器。它可以将汇编后的内核直接输出到 cubin 文件中（一种 NVIDIA 用来存储其机器码的 ELF 格式文件），或者用于修改现有的 cubin 内核。目前支持约 70 条指令。asfermi 可在 32 位和 64 位平台上运行。

Fermi 是 NVIDIA 当前一代图形处理单元(GPU)架构（截至 2011 年）。NVIDIA 并未披露其当前架构的许多细节。虽然 NVIDIA 为 CUDA 提供了一种类似汇编的语言 PTX，但 PTX 实际上被视为一种高级语言，由 ptxas 进行编译和优化。因此，无法使用真实操作码直接访问硬件。借助这个汇编器，可以从指令一对一地直接生成真实操作码。在这种直接交互方式的帮助下，有望揭示 GF1xx 架构的所有重要细节，以便这些知识可用于对 Fermi（及后代）GPU 进行更准确的优化。

国内也有一帮技术极客做类似 MaxAs 和 asfermi 的事情，有些闭源，有些开源。

2014 年，阿里云飞天专用计算组启动了 AsKepler 项目，并在 2015 年 GTC 上公布了性能表现，2015 年云栖大会阿里云 HPC 产品上线，提供 Nvidia Kepler 架构 GPU K40 的同时也提供了 AsKepler 在线汇编器工具方便用户对核心进行调优。笔者是第一批种子用户，因此印象深刻。互联网的记忆有时会消失，目前包括阿里云和 GTC 涉及 AsKepler 的相关内容都已下线。

2020 年，香港浸会大学褚晓文教授（目前就职于香港科技大学）带领团队开发并开源了 TuringAs 汇编器，支持 Volta、Turing、Ampere 架构 GPU。

Github：https://github.com/daadaada/turingas

商汤 HPC 团队在 2022 年开源了 CuAssembler，对多架构的支持更完整（当前主要有 SM60/SM61/SM70/SM75/SM80/SM86），指令编码支持和扩展也更容易。同时，CuAssembler 能更好地嵌入现有 NVIDIA 官方工具链，可以实现大部分代码的功能继承自原有编译流程，从而支持更多 CUDA feature。

Github：https://github.com/OpenPPL/CuAssembler

这个领域必读的论文和报告：

https://inria.hal.science/file/index/docid/789958/filename/112\_Lai.pdf

https://arxiv.org/pdf/1804.06826

https://arxiv.org/pdf/1903.07486

https://arxiv.org/pdf/2402.13499

最后，将话题回到 DeepGEMM。

在 GPU 性能优化方面，大多数人（包括我在内）都认为“够用就行”，如果一张卡性能不够，就两张，通过规模化优势达到目的。这是一条 scale-out 的路线，包括 OpenAI，X.AI，Meta，Google 和 Nvidia 都是沿这条确定性高的路线狂奔。

而 DeepSeek 和前面一些技术极客却在想的是，GPU 我花钱买了，凭啥不给我解锁 100% 的算力？Nvidia 你安的什么心？于是，采用了各种深入底层的汇编级优化，去探索硬件的极限性能。这个投入产出比可能不高，但值得，在国内算力受限，数量短缺的情况下，尤其珍贵。

DeepGEMM 将代码开源，给老外一点小小的震撼，原来还真有人在敲骨吸髓式用 GPU。

不知远在大洋彼岸的 Scott Gray 看到 DeepGEMM，是否曾想到十年前意气风发的屠龙少年，和 GTX 980 为伴的那个下午。

扫描下方