DeepGEMM 开源：英伟达 GPU 高性能 GEMM 内核库

DeepGEMM 是一个统一的高性能张量核内核库，汇集了现代大语言模型的关键计算原语——GEMM（FP8、FP4、BF16）、带通信重叠的融合 MoE（Mega MoE）、闪电索引器的 MQA 评分、HyperConnection（HC）等——整合到一个 cohesive 的 CUDA 代码库中。所有内核均通过 DeepJIT 在运行时编译，安装时无需 CUDA 编译。

DeepGEMM 借鉴了 CUTLASS 和 CuTe 的一些概念，但避免了对其模板或代数的重度依赖。该库设计简洁，核心内核函数数量有限，是学习 NVIDIA GPU 内核优化技术的清晰、易懂的资源。

尽管设计轻量，DeepGEMM 在各种矩阵形状下的性能可与或超过专家调优的库。

2026.09.30 DeepGEMM Ascend 来了！查看 DeepGEMM-Ascend 了解更多详情

添加了更多优化，包括 locality domain 特性，查看 #462 了解更多详情

2026.09.10: Sparse Indexer、Mega Gate、Mega mHC、DeepJIT、MoE 和 Indexer 优化等。

请参阅 #432 了解更多详情。

2026.04.16: Mega MoE、FP8xFP4 GEMM、FP4 Indexer、PDL、更快的 JIT 编译等。

请参阅 #304 了解更多详情。

关于 Mega MoE 基准测试，请参阅 #316。

2025.09.28: DeepGEMM 现已支持 DeepSeek v3.2 闪电索引器的评分内核（加权 ReLU MQA logits）。

请参阅 #200 了解更多详情。

2025.07.20: DeepGEMM 现已支持 SM90/SM100，并进行了全面重构，具有低 CPU 开销的 JIT CPP 模块。由于 NVCC 12.9 将自动执行 FFMA 交错，所有后续优化将不再受支持。

请参阅 #112 了解更多详情。

2025.05.14: DeepGEMM 现已为 dense 和 MoE backward 提供权重梯度内核！详见 #95。

2025.04.18: DeepGEMM 现已在 H800 上实现高达 1550 TFLOPS！详见 #74、#78、#81、#86 和 340d988。

环境要求

NVIDIA SM90 或 SM100 架构 GPU

编译器和标准库需支持 C++20 <format>

CUDA Toolkit 12.9 或更高版本

PyTorch 2.3 或更高版本

CUTLASS 4.0 或更高版本（可通过 Git submodule 克隆）

安装

# Submodule must be cloned
git clone --recursive git@github.com:deepseek-ai/DeepGEMM.git
cd DeepGEMM

# Link some essential includes and build the C++ extension
cat develop.sh
./develop.sh

cat install.sh
./install.sh

然后，在你的 Python 项目中 import deep_gemm，即可使用！

GEMM 内核

该库为 NVIDIA GPU 提供优化的 GEMM 内核，命名约定为：D = C + A @ B。输入形状布局为 NT（A 不转置，B 转置）。SM90 实现仅支持 NT 内存布局（行优先、列优先），而 SM100 实现支持所有内存布局（NT、TN、NN、TT）。例如，fp8_gemm_nt 将执行 D = C + A @ B.T。

对于两种架构，LHS 缩放因子需要 TMA 对齐和转置布局。SM90 和 SM100 的缩放因子数据格式不同：

• SM90 要求缩放因子为 FP32 格式。
• SM100 要求缩放因子为 packed UE8M0 格式，每个 torch.int 打包 4 个 UE8M0。

请注意，输入转置或 FP8 转换等操作必须由用户单独处理，请自行实现或将其融合到先前内核中。虽然该库提供了一些简单的 PyTorch 工具函数，但这些可能导致性能下降，我们的主要关注点是优化 GEMM 内核本身。

分组 GEMM

要执行基本的非分组 FP8 GEMM，请调用 fp8_gemm_{nt, nn, tn, tt} 函数。更多详情请参阅函数文档。

与 CUTLASS 中的传统分组 GEMM 不同，DeepGEMM 仅在 M 轴上分组，而 N 和 K 必须保持固定。此设计专为 MoE 模型中专家共享相同形状的场景定制。对于训练前向传播或推理预填充，每个专家可能处理不同数量的 token，我们将这些 token 连接成单个张量，称为"连续"布局。注意，每个专家段必须对齐到 GEMM M 块大小（get_mk_alignment_for_contiguous_layout()）。更多详情请参阅 m_grouped_fp8_gemm_{nt, nn}_contiguous 函数文档。

我们还为 MoE 权重反向传播提供了 K 轴分组 API（M 和 N 必须保持固定），请参阅 k_grouped_fp8_gemm_tn_contiguous。

掩码分组 GEMM

在推理解码阶段，当启用 CUDA graph 且 CPU 不知道每个专家收到的 token 数量时，我们支持掩码分组 GEMM。通过提供掩码张量，内核仅计算有效部分。

使用 m_grouped_fp8_gemm_nt_masked，并参阅相关文档。使用示例是将 DeepEP 的低延迟内核输出作为输入。

MQA 评分

该内核系列有两个版本，非分页（用于预填充）和分页（用于解码）。以非分页版本 fp8_fp4_mqa_logits 为例，其主要输入为：

• q：一个 (q_data, q_sf) 元组；SM100 接受 MXFP4/MXFP8 数据和 packed UE8M0 缩放因子
• kv：一个 (kv_data, kv_sf) 元组，逻辑形状为 [seq_len_kv, head_dim]
• weights：形状为 [seq_len, num_heads] 的张量（SM100 上为 BF16）
• cu_seq_len_k_start 和 cu_seq_len_k_end：形状为 [num_tokens] 的 int 张量
• max_seqlen_k：任何查询行最大的有效 KV 跨度

输出被压缩为 [seq_len, max_seqlen_k]；第 i 行从第零列开始存储其有效 KV 跨度。对于 q 中的每个 token i，它将遍历 [cu_seq_len_k_start[i], cu_seq_len_k_end[i]) 中的所有 token j，并计算对应的压缩 logit：

kv_j = kv[0][j, :] * kv[1][j].unsqueeze(1)  # [head_dim]
out_ij = q[i, :, :] @ kv_j  # [num_heads]
out_ij = out_ij.relu() * weights[i, :]  # [num_heads]
out_ij = out_ij.sum()  # Scalar

关于分页版本 fp8_fp4_paged_mqa_logits 的更多详情，请参阅 tests/test_attention.py。

Mega MoE

Mega MoE 将 EP dispatch、linear 1 和 linear 2（FP8xFP4 或 FP8xFP8）、SwiGLU 以及 EP combine 融合并重叠到单个 mega 内核中，叠加 NVLink 通信和张量核计算。它需要多进程启动和对称内存。用法：

# Allocate symmetric memory buffer
# NOTES: requires PyTorch >= 2.9
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',  # Use 'fp8xfp8' for FP8 routed-expert weights
)

# Transform weights (FP4 or FP8 with UE8M0 SF) into the required layout
transformed_l1, transformed_l2 = deep_gemm.transform_weights_for_mega_moe(l1_weights, l2_weights)

# (Optional) Localize weights into locality domains
transformed_l1 = (deep_gemm.localize(transformed_l1[0]), transformed_l1[1])
transformed_l2 = (deep_gemm.localize(transformed_l2[0]), transformed_l2[1])
deep_gemm.destroy_localizer()

# Copy inputs into the buffer before each call
# You may fuse these into previous kernels
buffer.x[:num_tokens].copy_(x_fp8)
buffer.x_sf[:num_tokens].copy_(x_sf)
buffer.topk_idx[:num_tokens].copy_(topk_idx)
buffer.topk_weights[:num_tokens].copy_(topk_weights)

# Run the fused mega MoE kernel
y = torch.empty((num_tokens, hidden), dtype=torch.bfloat16, device='cuda')
deep_gemm.fp8_fp4_mega_moe(y, transformed_l1, transformed_l2, buffer)

关于多进程设置和基准测试的完整示例，请参阅 tests/test_mega_moe.py。

工具函数

该库除了上述内核外还提供一些工具函数：

• deep_gemm.set_num_sms / get_num_sms：设置/获取使用的最大 SM 数量
• deep_gemm.set_tc_util / get_tc_util：设置/获取近似的张量核利用率
• deep_gemm.set_pdl / get_pdl：启用/禁用 Programmatic Dependent Launch（PDL）
• deep_gemm.use_deterministic_algorithms：启用/禁用确定性算法
• deep_gemm.set_mk_alignment_for_contiguous_layout / get_mk_alignment_for_contiguous_layout：设置/获取连续布局的分组级 M/K 对齐
• deep_gemm.get_theoretical_mk_alignment_for_contiguous_layout：获取理论最小 M/K 对齐
• deep_gemm.set_ignore_compile_dims：配置 JIT 编译期间忽略的维度
• deep_gemm.set_block_size_multiple_of：约束块大小为给定值的倍数
• deep_gemm.transform_sf_into_required_layout：将缩放因子转换为所需布局
• deep_gemm.get_tma_aligned_size：获取所需的 TMA 对齐大小
• deep_gemm.get_mn_major_tma_aligned_tensor：获取 MN-major TMA 对齐张量
• deep_gemm.get_mn_major_tma_aligned_packed_ue8m0_tensor：获取 MN-major TMA 对齐张量（将 FP32 打包为 UE8M0）
• deep_gemm.get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor：K 分组 GEMM 打包内核

环境变量

该库还提供一些可能有用的环境变量：

每个 DG_JIT_* 变量在未设置时会回退到对应的全局 DJ_JIT_* 变量。

通用

• DG_JIT_DEBUG: 0 或 1，启用 JIT 调试功能，包括编译器命令和 PTXAS 输出、加载时报告、行信息和 PTX/SASS 转储；默认为 0
• DG_PRINT_CONFIGS: 0 或 1，打印每个 shape 选择的配置；默认为 0

JIT 缓存

• DG_JIT_CACHE_DIR: 字符串，编译内核的缓存目录（可以是冒号分隔的目录列表）；查找按从头到尾的顺序搜索所有路径（首次命中优先），缓存未命中时编译到第一个路径；默认为 $HOME/.dj

编译器选择

• DG_JIT_NVCC_COMPILER: 字符串，NVCC 编译器路径；否则通过 CUDA_HOME、CUDA_PATH、which nvcc 然后 /usr/local/cuda 查找 CUDA
• DG_JIT_CPP_STANDARD: 整数，C++ 标准版本，默认为 20

编译器输出

• DG_JIT_PRINT_COMPILER_COMMAND: 0 或 1，打印编译命令；默认为 0
• DG_JIT_PTXAS_VERBOSE: 0 或 1，显示详细 PTXAS 输出；默认为 0
• DG_JIT_CHECK_NO_SPILLS: 0 或 1，断言编译内核无寄存器溢出；默认为 0
• DG_JIT_CHECK_NO_LOCAL_MEMORY: 0 或 1，断言编译内核无本地内存使用；默认为 0
• DG_JIT_PRINT_LOAD_TIME: 0 或 1，打印内核加载时间；默认为 0

调试和性能分析

• DG_JIT_WITH_LINEINFO: 0 或 1，为性能分析工具嵌入源行信息；默认为 0
• DG_JIT_DUMP_ASM: 0 或 1，转储 PTX 和 SASS；默认为 0
• DG_JIT_DUMP_PTX: 0 或 1，转储 PTX 输出；默认为 0
• DG_JIT_DUMP_SASS: 0 或 1，转储 SASS 输出；默认为 0
• DG_COMM_KERNEL_DEBUG: 0 或 1，每次 Mega MoE 调用前将对称缓冲区清零用于调试；默认为 0
• DG_USE_NVIDIA_TOOLS: 0 或 1，在外部 NVIDIA 工具下运行时跳过内部性能分析；默认为 0

构建选项

• DG_SKIP_CUDA_BUILD: 0 或 1，安装期间跳过 CUDA 扩展构建；默认为 0
• DG_FORCE_BUILD: 0 或 1，强制本地构建而非下载预编译 wheels；默认为 0

关于更多示例和详情，请参阅测试代码或查阅相应的 Python 文档。

致谢

DeepGEMM 灵感来自 CUTLASS 项目。感谢所有开发者的贡献！

本代码仓库基于 MIT 许可证发布。

@misc{deepgemm2025,
      title={DeepGEMM: clean and efficient BLAS kernel library on GPU},
      author={Chenggang Zhao and Zhean Xu and Liang Zhao and Jiashi Li and Chenhao Xu and Anyi Xu and Shengyu Liu and Kexing Zhou and Kuai Yu},
      year={2025},
      publisher = {GitHub},
      howpublished = {\url{https://github.com/deepseek-ai/DeepGEMM}},
}