DeepGEMM：为大模型训练与推理提供高性能 GPU 矩阵计算内核
一分钟看懂
DeepGEMM 是 DeepSeek 开源的 CUDA 内核库，把现代大模型常用的高性能矩阵乘法与相关计算集中到一个代码库中。它覆盖 FP8、FP4、BF16 GEMM、分组 GEMM、MoE 融合计算和索引器评分等场景，并通过 DeepJIT 在运行时编译内核，因此安装阶段不需要完成 CUDA 内核编译。
你可以用它做什么
- 在支持的 NVIDIA GPU 上运行低精度或 BF16 矩阵乘法。
- 为混合专家模型使用连续或掩码布局的分组 GEMM。
- 把专家并行通信、两层线性计算、SwiGLU 与合并过程融合进 Mega MoE 内核。
- 研究 NVIDIA Tensor Core、TMA、JIT 编译和大模型内核优化方法。
- 在性能工程中测试不同矩阵形状、精度和内存布局的实现。
主要能力
- 支持面向 SM90 与 SM100 架构的 FP8、FP4 和 BF16 计算内核。
- 提供普通稠密 GEMM、面向 MoE 的分组 GEMM、权重梯度以及 MQA 索引评分内核。
- 使用 DeepJIT 按需编译，并提供缓存、调试、汇编输出与编译器选择配置。
- 提供 Python 接口、布局转换与对齐辅助函数，测试目录中包含使用和基准示例。
- 代码结构相对集中，既可用于集成，也可作为 GPU 内核优化的学习材料。
适合哪些人
适合大模型训练或推理框架开发者、CUDA 性能工程师以及研究低精度矩阵计算的人员。它不是面向普通用户的模型运行应用，也不适合没有受支持 NVIDIA GPU 或 CUDA 开发环境的设备。
怎么开始使用
先准备 NVIDIA SM90 或 SM100 GPU、Python 3.8 以上版本、支持 C++20 <format> 的编译器、CUDA Toolkit 12.9 以上版本、PyTorch 2.3 以上版本和 CUTLASS 4.0 以上版本。开发者应递归克隆子模块，再按仓库中的 develop.sh 或 install.sh 构建安装，随后在 Python 中导入 deep_gemm。具体张量布局、缩放因子格式和函数调用应以 README 与测试代码为准。
使用前需要知道
硬件和工具链要求较高，当前官方要求 SM90 或 SM100 架构，不能把结果直接推广到其他 GPU。SM90 与 SM100 支持的内存布局和缩放因子格式不同；输入转置、FP8 转换等操作通常需要用户在前序内核中自行处理。项目公布的性能数据依赖具体硬件、形状和配置，集成前应在目标环境重新验证。代码采用 MIT 许可证。