DeepSeek开源FP8计算库DeepGEMM，300行代码性能超越专家优化，R2版本即将发布。
原文标题：DeepSeek-R2曝5月前上线！第三弹DeepGEMM 300行代码暴击专家优化内核
原文作者：数据派THU
冷月清谈：
DeepSeek 最新开源了FP8计算库DeepGEMM，仅用300行代码就实现了对V3/R1训练和推理的强大支持，性能超越专家优化内核。该库支持稠密和MoE模型，具有轻量级依赖、即时编译等优点，核心逻辑简洁易懂。DeepGEMM在Hopper GPU上实现了高达1350+ FP8 TFLOPS的算力，并在多种矩阵规模下超越了专家级优化内核。同时，DeepSeek-R2预计5月前发布，将带来更好的编码和多语言推理能力。DeepSeek的前两个开源项目FlashMLA和DeepEP也已获得大量关注。
怜星夜思：
1、DeepGEMM只用了300行代码就达到了如此高的性能，这背后的核心技术是什么？除了文章提到的JIT编译、TMA特性等，还有什么其他关键技术吗？
2、DeepSeek开源的这几个项目，FlashMLA、DeepEP和DeepGEMM，它们之间有什么联系和区别？在实际应用中应该如何选择？
3、DeepSeek-R2 的发布会对 AI 行业带来哪些影响？它是否真的能打破目前少数巨头的垄断？
2、DeepSeek开源的这几个项目，FlashMLA、DeepEP和DeepGEMM，它们之间有什么联系和区别？在实际应用中应该如何选择？
3、DeepSeek-R2 的发布会对 AI 行业带来哪些影响？它是否真的能打破目前少数巨头的垄断？
原文内容
来源：新智元本文约3000字，建议阅读5分钟
DeepSeek开源第三弹，是支持稠密和MoE模型的FP8计算库——DeepGEMM，支持V3/R1训推。仅凭300行代码，就超过了专家优化的内核。开发者惊叹：DeepSeek有最好的GPU工程师，仿佛拥有某种编译器黑魔法！更令人兴奋的是，DeepSeek-R2有望在5月前提前发布。
近期，DeepSeek发布了DeepGEMM。
这是一个支持稠密和MoE模型的FP8 GEMM（通用矩阵乘法）计算库，可为V3/R1的训练和推理提供强大支持。
仅用300行代码，DeepGEMM开源库就能超越专家精心调优的矩阵计算内核，为AI训练和推理带来史诗级的性能提升！
DeepGEMM库具有以下特征：
-
在Hopper GPU上实现高达1350+ FP8 TFLOPS的算力
-
极轻量级依赖，代码清晰易懂
-
完全即时编译，即用即跑
-
核心逻辑仅约300行代码，却在大多数矩阵规模下超越专家级优化内核
-
同时支持密集布局和两种MoE布局
开发者惊叹道：才300行代码，就能打败专家优化的内核？！
要么是DeepSeek真的破解了GPU运算的天机，要么我们就是见证了有史以来最高级的编译器黑科技。
总之，这个DeepGEMM听起来简直是数学界的超级英雄，比飞快的计算器还要快。
它改变了我们使用FP8 GEMM库的方式，简单、快速、开源。这就是AI计算的未来！
同时，外媒还曝出了另一个重磅消息：原计划在5月初发布的DeepSeek-R2，现在发布时间将再次提前！
在DeepSeek-R2中，将实现更好的编码，还能用英语以外的语言进行推理。
业内人士预测，DeepSeek-R2的发布，将是AI行业的一个关键时刻。目前DeepSeek在创建高成本效益模型上的成功，已经打破了该领域少数主导玩家的垄断。
DeepSeek开源两天，前两个项目爆火程度难以想象。FlashMLA已在GitHub斩获近10k星标，DeepEP的星标已有5k。
DeepGEMM
DeepGEMM是一个专为清晰高效的FP8通用矩阵乘法（General Matrix Multiplications，GEMMs）设计的库，它采用了DeepSeek-V3中提出的细粒度缩放技术。
该库支持常规矩阵乘法和混合专家模型（Mix-of-Experts，MoE）分组矩阵乘法。DeepGEMM使用CUDA编写，无需在安装时进行编译，而是通过轻量级即时编译（Just-In-Time，JIT）模块在运行时编译所有内核。
目前，DeepGEMM仅支持NVIDIA Hopper张量核。为了解决FP8张量核在累加计算时的精度问题，该库采用了基于CUDA核心的两级累加（提升）技术。
虽然DeepGEMM借鉴了CUTLASS和CuTe的一些概念，但避免了过度依赖它们的模板或代数系统。
相反，该库追求设计简洁，仅包含一个核心内核函数，代码量仅约300行。这使其成为学习Hopper FP8矩阵乘法和优化技术的理想入门资源。
尽管采用轻量级设计，DeepGEMM在处理各种矩阵形状时的性能都能够达到甚至超越经专家调优的库。
性能
研究人员在配备NVCC 12.8的H800上测试了DeepSeek-V3/R1推理过程中，可能使用的所有矩阵形状（包括预填充和解码阶段，但不包括张量并行计算）。
所有性能提升指标均与基于CUTLASS 3.6内部精心优化的实现进行对比计算得出。
DeepGEMM在某些矩阵形状下的表现还不够理想，如果你对此感兴趣，可以提交优化相关的Pull Request（拉取请求）。
稠密模型的常规GEMM
下表展示了不同矩阵维度（M、N、K）下DeepGEMM库的性能数据，结果显示在某些配置（如 M=128, N=2112, K=7168）下实现了高达 2.4 倍的加速，反映了DeepGEMM在优化GPU矩阵计算方面的效率和灵活性。
MoE模型的分组GEMM（使用连续存储布局）
MoE模型的分组GEMM（使用掩码存储布局）
快速入门
要求
-
NVIDIA Hopper架构GPU（需支持sm_90a计算能力）
-
Python v3.8或更高版本
-
CUDA v12.3及以上版本（强烈建议使用v12.8或更新版本以获得最佳性能）
-
PyTorch v2.1及以上版本
-
CUTLASS v3.6或更高版本 （可通过Git子模块[submodule]方式克隆获取）
开发
下面代码是DeepGEMM项目的安装和测试指南。
首先，通过命令克隆仓库及其子模块。然后，创建第三方库（CUTLASS和CuTe）的符号链接以便开发。接着，测试JIT编译功能。最后，测试所有GEMM实现。
br
安装
下面代码使用脚本安装Python包，会将包及其依赖项安装到系统中以便在项目中使用。
br
接下来，在你的Python项目中导入deep_gemm，就可以开始使用啦！
优化技术
注意：下面用🐳标记的是，CUTLASS中未包含的技术。
持久化线程束专用化
遵循CUTLASS的设计，DeepGEMM中的内核采用线程束（warp）专用化技术，实现了数据移动、张量核心MMA（矩阵乘累加）指令和CUDA核心提升操作的重叠执行。下图简要说明了这个过程：
TMA线程主要负责数据加载（Data load）和任务分发（TMA issue），用黄色和蓝色表示。数学线程则交替执行WGMA（Wavefront Matrix Multiply-Accumulate）计算（绿色）和数据提升（Promotion，黄色），展示了一种并行计算策略，其中数据加载与矩阵计算和优化操作协同工作，以提高效率和性能。
Hopper TMA特性
张量内存加速器（Tensor Memory Accelerator，TMA）是Hopper架构引入的新硬件特性，用于实现更快速的异步数据移动。具体来说，在以下方面使用TMA：
-
LHS（左矩阵）、LHS缩放因子和RHS（右矩阵）的TMA加载
-
输出矩阵的TMA存储
-
LHS矩阵的TMA多播
-
TMA描述符预取
常见的细节优化
-
使用stmatrixPTX指令
-
针对不同线程束组的寄存器数量精确控制
-
最大化指令重叠，如TMA 存储与非TMA RHS 缩放因子加载的重叠🐳
统一且经过优化的块调度器
-
所有非分组和分组内核使用同一调度器
-
采用光栅化技术提高L2缓存重用率
完全JIT设计 🐳
DeepGEMM采用完全即时编译（JIT）设计，无需在安装时编译。所有内核在运行时通过轻量级JIT实现进行编译。这种方法具有以下优势：
-
GEMM（通用矩阵乘法）形状、块大小和流水线阶段数被视为编译时常量
-
有效节省寄存器空间
-
使编译器能够进行更多优化
-
能够自动选择块大小、线程组数量、最优流水线阶段和TMA（张量内存访问）集群大小
-
即使在不进行自动调优的情况下，也能确定性地选择最优配置
-
完全展开MMA（矩阵乘加）流水线，为编译器提供更多优化机会
-
这一特性对处理小规模矩阵运算尤为重要
-
详细信息请参考kernel文件中的launch_k_iterations部分
总的来说，JIT显著提升了小形状的计算性能，这与Triton编译器采用的方法类似。
非对齐块大小🐳
对于某些形状，采用2的幂次对齐的块大小可能导致SM利用率不足。
例如，当M=256，N=7168时，传统的块大小分配BLOCK_M=128，BLOCK_N=128只能利用 (256/128) * (7168/128) = 112个SM（总共132个）。
为解决这个问题，团队为诸如112这样的非对齐块大小提供了支持，使得 (256/128) * (7168/112) = 128个SM能够充分工作。将这种技术与细粒度缩放结合需要精心优化，但最终能带来显著的性能提升。
FFMA SASS交错优化🐳
团队发现CUTLASS FP8内核在NVCC 12.2和12.3版本之间存在性能差异。
通过比对编译后的SASS代码，可以发现在一系列FADD指令中有一个位按交错模式翻转。
参考开源CUDA汇编器实现后，团队确定这个位控制着让出（yield）操作，可能用于增强线程束级并行性（推测是通过让出当前线程束使其他线程束得以执行）。
为此，团队开发了专门的脚本来修改编译后二进制中的FFMA指令。除了修改让出位，还调整了重用位（当线程束被让出时禁用寄存器重用）。
这种优化通过创造更多MMA指令和提升类FFMA指令重叠的机会，显著提高了细粒度缩放FP8 GEMM的性能（在某些情况下提升超过10%）。
参考资料：
https://x.com/deepseek_ai/status/1894553164235640933
编辑：文婧
