DeepGEMM
SHA-256DeepSeek 开源的高性能 GPU 矩阵计算算子库,为大模型 FP8/FP4 矩阵乘和 MoE 加速而生,安装免编译,性能媲美 CUTLASS。
让大模型矩阵运算在英伟达新显卡上跑到极限速度的开发库。
核心功能
- 支持 FP8、FP4、BF16 矩阵乘法,H800 上峰值可达 1550 TFLOPS
- Mega MoE 融合内核:把专家分发、两层线性层、SwiGLU 激活和结果合并压进一个内核,通信计算重叠
- DeepJIT 运行时编译,安装时无需任何 CUDA 编译,开箱即用
- 代码刻意保持精简,核心内核数量少,是学习 GPU 内核优化的优质教材
- 提供连续/掩码两种分组 GEMM 和 MQA 打分内核,覆盖 MoE 训练与推理解码全流程
避坑指南
- •只支持 NVIDIA SM90(H100/H800)和 SM100(B200)架构,RTX 系列消费显卡和 A100 等旧卡无法运行;它是面向开发者的底层算子库,没有图形界面,不是装完就能用的应用;输入转置和 FP8 量化转换需要用户在自己的前置算子里完成,库本身不代劳;依赖 Linux 环境、CUDA 12.9+ 和支持 C++20 <format> 的编译器,环境配置门槛不低;Mega MoE 还需要多进程启动和对称显存,调试复杂度较高。
适用场景
- 训练或部署 DeepSeek V3 这类大语言模型时,加速 FP8/FP4 低精度矩阵运算
- MoE 模型推理服务中,用 Mega MoE 把专家并行计算与 NVLink 通信重叠以降低延迟
- 作为参考代码,学习英伟达 Tensor Core 内核优化与 TMA 编程技巧
详细介绍
DeepGEMM 是 DeepSeek 开源的统一高性能 GPU 张量核心算子库,把现代大模型的关键计算原语整合进一个干净的 CUDA 代码库:FP8/FP4/BF16 矩阵乘法、通信与计算重叠的融合 MoE(Mega MoE)、lightning indexer 的 MQA 打分内核、HyperConnection 等。所有内核通过 DeepJIT 在运行时编译,安装时完全不需要编译 CUDA。它借鉴了 CUTLASS/CuTe 的思路,但刻意避开了繁重的模板体系,只保留少量核心内核——因此它既是生产级库(H800 上可达 1550 TFLOPS),也是学习英伟达 GPU 内核优化最易读的教材之一。支持 SM90(H100/H800)和 SM100(B200)架构,提供连续/掩码两种布局的分组 GEMM,覆盖 MoE 训练前向与推理解码。需要 CUDA 12.9+、PyTorch 2.3+ 和支持 C++20 的编译器,MIT 协议开源。
常见问题与排障 (2)
故障排除为什么 DeepGEMM 的 test_mqa_logits 在 SM120 Blackwell GPU 上会因按位不匹配而失败?FP8 MQA logits 核函数的输出是否可以安全使用?
该故障在 SM120 上有两个相互独立的原因(例如 RTX PRO 6000 Blackwell、deep_gemm 2.6.1 nv_dev):(1) FP8 自一致性断言失败,原因在于非确定性的*窗口外*写入——在压缩模式下,fp8_fp4_mqa_logits kernel 会将垃圾数据写入超出每行有效 [ks, ke) 窗口之外的输出单元,导致测试的全缓冲区逐位比较不匹配。更深入的隔离测试(30 次重复启动,形状最大至 4096x8192,H=64 D=128)表明,有效区域的输出逐位一致且不含 NaN,因此遵守窗口约束的消费方不受影响:只读取每行 [ks_i, ke_i) 范围内的 logits,并将该测试失败视为测试伪影,而非有效数据上的数据竞争。(2) H=64、D=128 的 MXFP4 路径存在真实的精度失败(diff ~0.62,而容差为 0.02),目前尚无修复——在按架构的同步修复落地之前,应避免在 SM120 上使用 FP4 indexer 路径,这与此前的 SM90(#229)和 SM100(#285,commit 96fbced)MQA 同步修复类似。
故障排除运行 mega_moe 内核时,如何修复 "AttributeError: module 'deep_ep' has no attribute 'ElasticBuffer'" 错误?
此错误的发生是因为 ElasticBuffer 是 DeepEP v2 中引入的新 API,它在 DeepEP v1 中并不存在。要解决此问题,请将您的 DeepEP 安装升级到 v2(参见 https://github.com/deepseek-ai/DeepEP/pull/605)。更新 DeepEP 后,deep_ep.ElasticBuffer 属性将可用,mega_moe 内核应该能够正常运行,不再出现此错误。
标签
快速上手
安装软件
双击下载的安装程序,按提示完成安装
确认硬件为 SM90 或 SM100 架构的 NVIDIA GPU(如 H100、H800、B200),并安装好 CUDA 12.9+ 和 PyTorch 2.3+
执行 git clone --recursive 克隆仓库,注意必须带上子模块
进入项目目录,运行 ./install.sh 完成安装
- 确认硬件为 SM90 或 SM100 架构的 NVIDIA GPU(如 H100、H800、B200),并安装好 CUDA 12.9+ 和 PyTorch 2.3+
- 执行 git clone --recursive 克隆仓库,注意必须带上子模块
- 进入项目目录,运行 ./install.sh 完成安装
已提供 SHA-256 校验码,下载后可自行核对文件完整性
该校验码提取自 GitHub 官方 Release 页面
SHA256 校验码
5b34ed544d56312b257320343fc98ea9bc704f512d38cfcca35f649ef209ded5该校验码提取自 GitHub Release 页面,下载后请自行核对文件完整性
本平台所有 SHA-256 校验码均提取自项目在 GitHub 官方 Release 页面发布的文件,未做任何修改。你可以通过 GitHub Releases 页面自行验证。
开源透明
查看 GitHub 源码卸载说明
先运行 pip uninstall deep_gemm 卸载 Python 包,然后删除克隆下来的源码目录即可;如需彻底清理,可一并删除运行时 JIT 生成的编译缓存目录。
无额外依赖
下载后即可直接使用,无需安装其他运行环境
遇到问题?查看下方 FAQ
2 FAQs