DeepSeek 开源第三弹:V3/R1 训练推理关键秘籍,核心代码仅 300 行

开源周的第三天,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 代码如下:

bash 复制代码
# Submodule must be cloned
git clone --recursive [email protected]: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

安装代码如下:

arduino 复制代码
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 入口:
github.com/deepseek-ai...

--- ---

相关推荐
zskj_zhyl7 分钟前
数据驱动的温暖守护:智慧康养平台如何实现 “千人千面” 的精准照护?
人工智能·科技·生活
视觉语言导航15 分钟前
ICRA-2025 | 视觉预测助力机器人自主导航!NavigateDiff:视觉引导的零样本导航助理
人工智能·机器人·具身智能
Seon塞翁23 分钟前
2025年Q1 AI 学习应用总结
人工智能·学习
jndingxin39 分钟前
OpenCV 图形API(30)图像滤波-----腐蚀操作函数erode()
人工智能·opencv
AIGC-Lison42 分钟前
Stable Diffusion ComfyUI 基础教程(一) ComfyUI安装与常用插件
人工智能·stable diffusion·教程·ai绘画·sd·sd教程
__Benco1 小时前
OpenHarmony子系统开发 - 热管理(四)
人工智能·harmonyos
正在走向自律1 小时前
数字人:打破次元壁,从娱乐舞台迈向教育新课堂(4/10)
人工智能·3d建模·语音识别·娱乐·数字人
爱编程的王小美1 小时前
如何玩转Hugging Face
人工智能·大模型
二川bro1 小时前
浅层神经网络:从数学原理到实战应用的全面解析
人工智能·深度学习·神经网络
奇舞精选1 小时前
DeepSeek V3-0324探索-使用prompt生成精美PPT
aigc·deepseek