DeepSeek 开源进度 3/5：深度学习利器 DeepGEMM

人阅读 2025-02-26 11:58:19

感谢网友刺客、昵称已起烦、HH_KK 的线索投递！

2 月 26 日消息，DeepSeek“开源周”的进度今日来到 3/5：支持稠密和混合专家模型 (MoE) 的 FP8 矩阵乘法 (GEMM) 库，用以驱动 V3 / R1 模型的训练和推理。

在 Hopper GPU 上可实现高达 1350 FP8 TFLOPS 性能
无复杂依赖，代码简洁如教程
完全采用即时编译技术（Just-In-Time）
核心代码仅约 300 行 —— 在大多数矩阵尺寸下超越了专家优化的内核
支持稠密布局和两种 MoE 布局

附开源链接：https://github.com/deepseek-ai/DeepGEMM

官方介绍大意如下：

DeepGEMM 是一个专为高效且清晰的 FP8 通用矩阵乘法（GEMM）设计的库，具备 DeepSeek-V3 所提出的精细化缩放能力。它支持普通的 GEMM 以及 Mix-of-Experts (MoE) 分组 GEMM。

该库基于 CUDA 编写，在安装时无需预编译，而是通过轻量级的即时编译（JIT）模块，在运行时动态编译所有内核。

目前，DeepGEMM 仅支持 NVIDIA Hopper 张量核心。为了应对 FP8 张量核心累加不精确的问题，它使用了 CUDA 核心的两级累加（提升）方法。虽然它借鉴了部分 CUTLASS 和 CuTe 的理念，但并未过度依赖它们的模板或代数结构。

DeepGEMM 的设计简洁，核心内核函数只有大约 300 行代码，方便学习 Hopper FP8 矩阵乘法和优化技术。

尽管采用轻量设计，DeepGEMM 在多种矩阵形状下的性能表现与专家优化的库相当，甚至更好。

我们在 H800 上，使用 NVCC 12.8 测试了 DeepSeek-V3 / R1 推理中可能用到的各种矩阵形状（包括预填充和解码，但不涉及张量并行）。所有加速指标都是相对于我们内部精心优化的 CUTLASS 3.6 实现计算的。

DeepGEMM 在某些矩阵形状下的表现不尽如人意，欢迎有兴趣的朋友提交优化 PR。

广告声明：文内含有的对外跳转链接（包括不限于超链接、二维码、口令等形式），用于传递更多信息，节省甄选时间，结果仅供参考，所有文章均包含本声明。