标题：DeepGEMM结合Hopper架构优化矩阵乘法，即时编译助力高性能计算——DeepSeek开源第三弹V3/R1训练推理核心代码仅300行

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

2025-02-26 10：40：49

来源：量子位
作者：金磊 发自 凹非寺

在开源周的第三天，DeepSeek团队正式推出了其最新的训练推理优化工具——DeepGEMM。这款FP8 GEMM（通用矩阵乘法）库不仅支持密集矩阵和混合专家（MoE）矩阵运算，而且核心代码仅有约300行，却能实现卓越的性能。

深入了解GEMM

GEMM（通用矩阵乘法）是线性代数中的基础运算，在科学计算、机器学习及深度学习领域中扮演着重要角色。由于其计算量巨大，GEMM的性能优化至关重要。DeepGEMM正是为此而生，具备以下亮点：

– 高性能：在英伟达Hopper架构GPU上，DeepGEMM可达到超过1350 FP8 TFLOPS的性能。
– 简洁性：核心逻辑精简至300行代码，性能优于传统调优内核。
– 即时编译（JIT）：采用完全即时编译方式，适应不同硬件和矩阵大小。
– 无重依赖：设计轻量级，部署和使用简单。
– 多矩阵布局支持：支持密集矩阵和两种MoE布局，适用于多种应用场景。

DeepGEMM的应用场景

DeepGEMM主要用于加速深度学习中的矩阵运算，尤其在大规模模型训练和推理中表现出色。它能够显著提升计算效率，特别适用于需要高效计算资源的场景。许多网友对这次开源表示赞赏，有人甚至将其比作数学界的超级英雄。

技术细节

DeepGEMM是一款用CUDA编写的库，安装时无需编译，运行时通过轻量级即时编译（JIT）模块动态生成优化代码。目前，该库仅支持英伟达Hopper张量核心，并采用两级累积方法解决FP8张量核心在计算累积时的精度问题。

与CUTLASS和CuTe相比，DeepGEMM的设计更为简洁，只有一个核心内核函数，代码量约为300行。尽管设计轻巧，但其性能可以匹配或超过各种矩阵形状的专家调优库。

性能测试

团队在H800 GPU上使用NVCC 12.8测试了DeepSeek-V3/R1推理中可能使用的所有形状（包括预填充和解码）。结果显示，DeepGEMM的计算性能最高可达1358 TFLOPS，内存带宽最高可达2668 GB/s。与基于CUTLASS 3.6的优化实现相比，加速比最高可达2.7倍。

安装与使用

要使用DeepGEMM，需满足以下依赖项：
– 支持Hopper架构的GPU（sm_90a）
– Python 3.8及以上
– CUDA 12.3及以上（推荐12.8）
– PyTorch 2.1及以上
– CUTLASS 3.6及以上

安装步骤如下：
“`bash
# 克隆仓库并初始化子模块
git clone –recursive git@github.com：deepseek-ai/DeepGEMM.git

# 创建符号链接以包含第三方库（如CUTLASS和CuTe）
python setup.py develop

# 测试即时编译
python tests/test_jit.py

# 测试所有GEMM实现
python tests/test_core.py
“`

安装完成后，可以在Python项目中导入`deep_gemm`，并调用相应的接口函数进行矩阵运算。

优化策略

DeepGEMM还提供了几种优化方式，包括：
– JIT设计：所有内核在运行时编译，支持动态选择最优块大小和流水线阶段。
– 细粒度缩放：通过两层累加解决FP8精度问题，优化SM利用率。
– FFMA SASS交错：修改SASS指令的yield和reuse位，提高性能。

感兴趣的小伙伴可以访问[GitHub链接](https：//github.com/deepseek-ai/DeepGEMM)查看详情。

One More Thing

最近，英伟达股票持续下跌，但其2025财年第四季度业绩报告即将在北京时间27日凌晨发布，值得期待。

本文来源： 量子位【阅读原文】