Deepseek 开源周第三日：DeepGEMM-开源库极限提升 NVIDIA H800 计算性能

Deepseek 开源周第三日的开源项目，DeepGEMM 是一个专为高效 FP8 通用矩阵乘法 (GEMM) 设计的库，支持密集型和混合专家 (MoE) GEMM。它用约 300 行核心代码，将 NVIDIA H800 的计算性能提升至 2.7 倍，帮助训练和推理模型实现更高效率。

核心特点

DeepGEMM 的核心特点包括：

• 干净的实现：核心逻辑仅约 300 行代码，易于理解和学习。
• 高性能：在多种矩阵形状下，性能优于专家调整的库。
• 支持多种布局：支持普通布局和两种 MoE 布局。
• 轻量级 JIT 编译：安装时无需编译，所有内核在运行时即时编译。

性能测试

我们在 NVIDIA H800 上测试了 DeepGEMM 的性能，结果如下：

密集模型的正常 GEMM

M	N	K	计算 (TFLOPS)	内存带宽 (GB/秒)	加速倍数
64	2112	7168	206	1688	2.7x
64	24576	1536	289	2455	1.7x
128	2112	7168	352	1509	2.4x
4096	2112	7168	1058	527	1.1x

用于 MoE 模型的分组 GEMM

DeepGEMM 支持两种 MoE 布局：连续布局和掩码布局。以下是测试结果：

连续布局

#群组	每组 M	N	K	计算 (TFLOPS)	内存带宽 (GB/秒)	加速倍数
4	8192	4096	7168	1297	418	1.2x
8	4096	4096	7168	1288	494	1.2x

掩码布局

#群组	每组 M	N	K	计算 (TFLOPS)	内存带宽 (GB/秒)	加速倍数
1	1024	4096	7168	1233	924	1.2x
2	512	4096	7168	1040	1288	1.2x

快速启动

要使用 DeepGEMM，您需要以下环境：

• Hopper 架构 GPU，sm_90a 必须支持
• Python 3.8 或更高版本
• CUDA 12.3 或更高版本（建议使用 12.8 或更高版本）
• PyTorch 2.1 或更高版本
• CUTLASS 3.6 或更高版本（可以通过 Git 子模块克隆）

有关更多详细信息，请参阅GitHub 项目页面。