DeepGEMM

人工智能训练和推理

DeepGEMM 是用于人工智能训练和推理的 FP8 GEMM 库，专为密集和专家混合（MoE）矩阵操作而设计，为 DeepSeek V3 和 R1 模型的训练和推理提供强大支持。

DeepGEMM

什么是 DeepGEMM？

DeepSeek 相继开源了 FlashMLA 和 深度教育今天，它将介绍 DeepGEMM是专为 Hopper 架构 GPU 优化的矩阵乘法库。该库支持标准矩阵计算和专家混合（MoE）计算，为 DeepSeek-V3/R1 的训练和推理提供了强大的支持，在 Hopper GPU 上实现了 1350+ FP8 TFLOPS 的高性能。

DeepGEMM deepGEMM的设计简单高效，核心代码仅约300行，在大多数矩阵尺寸下都优于现有解决方案。该库支持三种数据排列：一种标准排列和两种为混合专家模型设计的特殊排列（顺序排列和屏蔽排列）。deepGEMM 采用即时编译，无需在安装时进行编译，代码结构清晰易懂，非常适合学习 GPU 优化技术。

 

DeepGEMM 性能

DeepGEMM 在各种计算场景中都表现出色。与基于 CUTLASS 3.6 的优化实现相比，标准矩阵乘法的速度提高了 1.0 到 2.7 倍。在计算小批量数据（M=64 或 128）时，速度提升最为明显，最高可达 2.7 倍。对于混合专家模型的计算，DeepGEMM 提供的两种特殊数据排列方式也具有显著优势。顺序排列适用于训练和批量推理阶段，速度提高了约 1.1 到 1.2 倍，而遮蔽排列专为实时推理设计，支持使用 CUDA 图形技术，速度也提高了 1.1 到 1.2 倍。

密集模型的正态 GEMM

M N K 计算 内存带宽 提速
64 2112 7168 206 TFLOPS 1688 GB/s 2.7x
64 24576 1536 289 TFLOPS 2455 GB/s 1.7x
64 32768 512 219 TFLOPS 2143 GB/s 1.8x
64 7168 16384 336 TFLOPS 2668 GB/s 1.4x
64 4096 7168 287 TFLOPS 2320 GB/s 1.4x
64 7168 2048 295 TFLOPS 2470 GB/s 1.7x
128 2112 7168 352 TFLOPS 1509 GB/s 2.4x
128 24576 1536 535 TFLOPS 2448 GB/s 1.6x
128 32768 512 358 TFLOPS 2103 GB/s 1.5x
128 7168 16384 645 TFLOPS 2604 GB/s 1.4x
128 4096 7168 533 TFLOPS 2221 GB/s 2.0x
128 7168 2048 510 TFLOPS 2277 GB/s 1.7x
4096 2112 7168 1058 TFLOPS 527 GB/s 1.1x
4096 24576 1536 990 TFLOPS 786 GB/s 1.0x
4096 32768 512 590 TFLOPS 1232 GB/s 1.0x
4096 7168 16384 1358 TFLOPS 343 GB/s 1.2x
4096 4096 7168 1304 TFLOPS 500 GB/s 1.1x
4096 7168 2048 1025 TFLOPS 697 GB/s 1.1x

分组的全球环境监测模型（连续布局）

# 小组 每组 M N K 计算 内存带宽 提速
4 8192 4096 7168 1297 TFLOPS 418 GB/s 1.2x
4 8192 7168 2048 1099 TFLOPS 681 GB/s 1.2x
8 4096 4096 7168 1288 TFLOPS 494 GB/s 1.2x
8 4096 7168 2048 1093 TFLOPS 743 GB/s 1.1x

用于模拟实体模型的分组 GEMM（屏蔽布局）

# 小组 每组 M N K 计算 内存带宽 提速
1 1024 4096 7168 1233 TFLOPS 924 GB/s 1.2x
1 1024 7168 2048 925 TFLOPS 968 GB/s 1.2x
2 512 4096 7168 1040 TFLOPS 1288 GB/s 1.2x
2 512 7168 2048 916 TFLOPS 1405 GB/s 1.2x
4 256 4096 7168 932 TFLOPS 2064 GB/s 1.1x
4 256 7168 2048 815 TFLOPS 2047 GB/s 1.2x

如何使用 DeepGEMM？

要使用 DeepGEMM，需要支持 sm_90a 的 Hopper 架构 GPU、Python 3.8 或更高版本、CUDA 12.3 或更高版本（建议使用 12.8 或更高版本以获得最佳性能）、PyTorch 2.1 或更高版本以及 CUTLASS 3.6 或更高版本。

发展

必须克隆 # 子模块
git clone --recursive [email protected]:deepseek-ai/DeepGEMM.git

# 为第三方（CUTLASS 和 CuTe）包含目录创建符号链接
python setup.py develop

# 测试 JIT 编译
python tests/test_jit.py

# 测试所有 GEMM 实现（正常、连续分组和屏蔽分组）
python tests/test_core.py

安装

python setup.py install

最后，导入 deep_gemm 你就大功告成了！