DeepGEMM：统一高性能张量核心 GEMM 计算库编译与使用指南

DeepGEMM 是 DeepSeek 开源的统一高性能张量核心 kernel 库，把现代大模型的关键计算原语——FP8/FP4/BF16 GEMM、融合通信的 Mega MoE、MQA 打分、HyperConnection 等——整合进同一套 CUDA 代码库。所有 kernel 在运行时通过 DeepJIT 即时编译，安装时无需 CUDA 编译。

一、项目官方信息

二、官方图片

DeepGEMM 仓库社交预览图：

三、环境要求

• NVIDIA SM90（Hopper）或 SM100 架构 GPU；
• Python 3.8 及以上；
• 支持 C++20 <format> 的编译器与标准库；
• CUDA Toolkit 12.9 及以上；
• PyTorch 2.3 及以上；
• CUTLASS 4.0 及以上（可由 Git submodule 拉取）。

四、部署步骤

1. 克隆（含子模块）

git clone --recursive https://github.com/deepseek-ai/DeepGEMM
cd DeepGEMM

2. 开发模式构建

./develop.sh

3. 安装为 wheel

./install.sh

之后即可在 Python 中 import deep_gemm。

五、使用方法

基础非分组 FP8 GEMM

库的命名约定为 D = C + A @ B，输入布局为 NT（A 不转置、B 转置）：

import deep_gemm
# 非分组 FP8 GEMM
out = deep_gemm.fp8_gemm_nt(a_fp8, a_sf, b_fp8, b_sf, c)

SM90 仅支持 NT 布局；SM100 支持 NT/TN/NN/TT 全部布局。缩放因子（scale factor）SM90 用 FP32，SM100 用打包的 UE8M0。

分组 GEMM（MoE）

DeepGEMM 的分组只沿 M 轴分组，N/K 固定，适合 MoE 中各专家形状一致的场景：

• 训练前向 / prefill（连续布局）：m_grouped_fp8_gemm_{nt, nn}_contiguous；
• 权重反向（K 轴分组）：k_grouped_fp8_gemm_tn_contiguous；
• 推理 decoding（masked 布局）：m_grouped_fp8_gemm_nt_masked，输入可用 DeepEP 的低延迟 kernel 输出。

Mega MoE 融合 kernel

Mega MoE 把 EP dispatch、linear1/linear2、SwiGLU、EP combine 融合成一个 mega-kernel，重叠 NVLink 通信与张量核心计算：

buffer = deep_gemm.get_symm_buffer_for_mega_moe(
    group, num_experts, num_max_tokens_per_rank, num_topk, hidden, intermediate_hidden,
    mma_type='fp8xfp4',
)
transformed_l1, transformed_l2 = deep_gemm.transform_weights_for_mega_moe(l1_weights, l2_weights)
buffer.x[:num_tokens].copy_(x_fp8)
y = torch.empty((num_tokens, hidden), dtype=torch.bfloat16, device='cuda')
deep_gemm.fp8_fp4_mega_moe(y, transformed_l1, transformed_l2, buffer)

工具函数

• set_num_sms / get_num_sms：限制/查询 SM 数；
• set_tc_util：设置近似张量核心利用率；
• set_pdl：开关 Programmatic Dependent Launch；
• get_mk_alignment_for_contiguous_layout：查询连续布局的对齐。

常用环境变量

• DG_JIT_CACHE_DIR：JIT 缓存目录（默认 $HOME/.dj）；
• DG_JIT_DEBUG=1：打印编译器命令、PTXAS 输出、PTX/SASS dump；
• DG_JIT_PRINT_COMPILER_COMMAND=1：打印编译命令；
• DG_SKIP_CUDA_BUILD=1：安装时跳过 CUDA 扩展构建。

六、常见问题 FAQ

Q1：DeepGEMM 是什么？ 统一高性能张量核心 GEMM 库，覆盖 FP8/FP4/BF16 与 MoE、MQA 等大模型核心算子。

Q2：支持哪些 GPU？ NVIDIA SM90（Hopper）与 SM100 架构。

Q3：安装要 CUDA 编译吗？ 不需要。kernel 运行时由 DeepJIT 即时编译，安装脚本只构建 C++ 扩展。

Q4：怎么跑一次 FP8 GEMM？ deep_gemm.fp8_gemm_nt(a, a_sf, b, b_sf, c)，布局 NT。

Q5：MoE 的分组 GEMM 怎么调？ 连续布局用 m_grouped_fp8_gemm_nt_contiguous，decoding 用 m_grouped_fp8_gemm_nt_masked。

Q6：和 CUTLASS 什么关系？ 借鉴了 CUTLASS/CuTe 思路，但避免重度依赖其模板代数，核心逻辑精简，便于学习 GPU kernel 优化。

七、跳转链接

• GitHub 仓库：https://github.com/deepseek-ai/DeepGEMM
• DeepJIT：https://github.com/deepseek-ai/DeepJIT
• 配套 DeepEP：https://github.com/deepseek-ai/DeepEP