AI Feeds
开源项目

DeepGEMM

面向 NVIDIA GPU 的高性能 tensor core kernel 库，把 LLM 常用的 FP8/FP4/BF16 GEMM、融合 MoE（Mega MoE）、MQA scoring、HyperConnection 等计算原语收拢到一套 CUDA 代码中，通过 DeepJIT 在运行时编译，安装阶段不需要编译 CUDA。亮点在于代码量克制、核心 kernel 函数少，比 CUTLASS 那套模板体系更好读，同时在大批矩阵形状上性能追平甚至超过专家手工调优的库，适合想研究 GPU kernel 优化或给推理框架换后端的人。MIT 协议。

README

DeepGEMM

DeepGEMM 是一个统一的高性能 tensor core kernel 库，它将现代大语言模型的关键计算原语 —— GEMM（FP8、FP4、BF16）、带通信重叠的融合 MoE（Mega MoE）、用于 lightning indexer 的 MQA scoring、HyperConnection（HC）等 —— 汇集到一个统一、内聚的 CUDA 代码库中。所有 kernel 均在运行时通过 DeepJIT 编译，安装过程中无需进行 CUDA 编译。

DeepGEMM 借鉴了 CUTLASS 和 CuTe 的一些概念，但避免对其模板或代数体系产生重度依赖。该库以简洁为设计目标，只包含数量有限的核心 kernel 函数，因此是学习 NVIDIA GPU kernel 优化技术的一份清晰且易上手的资源。

尽管设计轻量，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 的 benchmark 请参考 #316。
  • 2025.09.28：DeepGEMM 现已支持 DeepSeek v3.2 lightning indexer 的 scoring kernel（weighted ReLU MQA logits）。
    • 详情请查看 #200。
  • 2025.07.20：DeepGEMM 现已同时支持 SM90/SM100，并通过一个低 CPU 开销的 JIT CPP 模块完成了全面重构。
    • 由于 NVCC 12.9 会自动进行 FFMA interleaving，所有后置优化将不再支持。
    • 详情请查看 #112。
  • 2025.05.14：DeepGEMM 现已提供用于 dense 和 MoE 反向传播的权重梯度 kernel！详情请查看 #95。
  • 2025.04.18：DeepGEMM 在 H800 上现已达到最高 1550 TFLOPS！详情请查看 #74、#78、#81、#86 和 340d988。

快速开始

环境要求

  • NVIDIA SM90 或 SM100 架构 GPU
  • Python 3.8 或更高版本
  • 支持 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 即可开始使用！

接口

注意事项

本库为 NVIDIA GPU 提供优化后的 GEMM kernel，命名约定为：D = C + A @ B。输入形状布局为 NT（A 非转置，B 转置）。SM90 实现仅支持 NT 内存布局（row-major、col-major），而 SM100 实现支持所有内存布局（NT、TN、NN、TT）。例如，fp8_gemm_nt 会执行 D = C + A @ B.T。

对于这两种架构，LHS scaling factor 都必须采用 TMA 对齐且转置的布局。并且 SM90 与 SM100 的 scaling factor 数据格式不同：

  • SM90 要求 scaling factor 为 FP32 格式。
  • SM100 要求 scaling factor 为 packed UE8M0 格式，即将 4 个 UE8M0 打包进一个 torch.int。

请注意，诸如输入转置或 FP8 casting 之类的操作必须由用户自行处理，请独立实现它们或将它们融合进前置 kernel 中。虽然本库提供了一些简单的 PyTorch 工具函数，但这些函数可能导致性能变慢，我们的主要关注点始终是 GEMM kernel 本身的优化。

普通稠密 GEMM（非分组）

要执行一个基础的非分组 FP8 GEMM，请调用 fp8_gemm_{nt, nn, tn, tt} 函数。更多细节请参考该函数的文档。

分组 GEMM（连续布局）

与 CUTLASS 中传统的分组 GEMM 不同，DeepGEMM 仅按 M 轴分组，而 N 和 K 必须保持固定。这一设计针对 MoE 模型中各专家共享相同形状的场景。在训练前向传播或推理 prefill 阶段，每个专家可能处理数量不等的 token，我们将这些 token 拼接成单个 tensor，即所谓的“连续”（contiguous）布局。注意每个专家分段必须对齐到 GEMM 的 M block size（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（masked 布局）

在推理 decoding 阶段，当启用 CUDA graph 且 CPU 无法感知每个专家接收的 token 数量时，我们支持 masked 分组 GEMM。通过提供一个 mask tensor，kernel 只计算有效的部分。

请使用 m_grouped_fp8_gemm_nt_masked 并查阅相关文档。一个示例用法是将 DeepEP 的低延迟 kernel 输出作为输入。

用于 indexer 的 V3.2 MQA kernel

该 kernel 家族有两个版本：非 paged（用于 prefill）和 paged（用于 decode）。 以非 paged 版本 fp8_fp4_mqa_logits 为例，其主要输入为：

  • q，一个 (q_data, q_sf) 元组；SM100 接受带 packed UE8M0 scale 的 MXFP4/MXFP8 数据
  • kv，一个 (kv_data, kv_sf) 元组，逻辑形状为 [seq_len_kv, head_dim]
  • weights，形状为 [seq_len, num_heads] 的 tensor（SM100 上为 BF16）
  • cu_seq_len_k_start 和 cu_seq_len_k_end，形状为 [seq_len] 的 int tensor
  • max_seqlen_k，任意 query 行的最大有效 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

更多细节以及 paged 版本 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-kernel 中，使 NVLink 通信与 tensor core 计算相互重叠。它需要多进程启动并使用对称内存。用法：

# 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)

包含多进程配置与 benchmark 的完整示例，请参考 tests/test_mega_moe.py。

工具函数

除上述 kernel 外，本库还提供了一些工具函数：

  • deep_gemm.set_num_sms / get_num_sms：设置/获取可使用的最大 SM 数量
  • deep_gemm.set_tc_util / get_tc_util：设置/获取近似的 tensor core 利用率
  • 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：设置/获取连续布局下 group 级别的 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：约束 block size 为给定值的倍数
  • deep_gemm.transform_sf_into_required_layout：将 scaling factor 转换为所需布局
  • deep_gemm.get_tma_aligned_size：获取所需的 TMA 对齐大小
  • deep_gemm.get_mn_major_tma_aligned_tensor：获取一个 MN-major 的 TMA 对齐 tensor
  • deep_gemm.get_mn_major_tma_aligned_packed_ue8m0_tensor：获取一个 MN-major 的 TMA 对齐 tensor（同时将 FP32 打包为 UE8M0）
  • deep_gemm.get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor：K 分组 GEMM 的打包 kernel

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

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

  • 通用
    • DG_JIT_DEBUG：0 或 1，启用 JIT 调试特性，包括编译器命令与 PTXAS 输出、加载耗时报告、line info，以及 PTX/SASS dump；默认 0
    • DG_PRINT_CONFIGS：0 或 1，打印每个 shape 所选的 config，默认 0
  • JIT 缓存
    • DG_JIT_CACHE_DIR：字符串，编译后 kernel 的缓存目录（或以 : 分隔的目录列表）；查找按从前到后的顺序搜索所有路径（命中即返回），未命中时编译结果写入第一个路径，默认 $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，断言编译后的 kernel 无寄存器溢出，默认 0
    • DG_JIT_CHECK_NO_LOCAL_MEMORY：0 或 1，断言编译后的 kernel 未使用 local memory，默认 0
    • DG_JIT_PRINT_LOAD_TIME：0 或 1，打印 kernel 加载耗时，默认 0
  • 调试与 profiling
    • DG_JIT_WITH_LINEINFO：0 或 1，为 profiling 工具嵌入源码行信息，默认 0
    • DG_JIT_DUMP_ASM：0 或 1，同时 dump PTX 与 SASS，默认 0
    • DG_JIT_DUMP_PTX：0 或 1，dump PTX 输出，默认 0
    • DG_JIT_DUMP_SASS：0 或 1，dump SASS 输出，默认 0
    • DG_COMM_KERNEL_DEBUG：0 或 1，为调试目的在每次 Mega MoE 调用前将对称缓冲区清零，默认 0
    • DG_USE_NVIDIA_TOOLS：0 或 1，在外部 NVIDIA 工具下运行时跳过内部 profiling，默认 0
  • 构建选项
    • DG_SKIP_CUDA_BUILD：0 或 1，安装期间跳过 CUDA 扩展构建，默认 0
    • DG_FORCE_BUILD：0 或 1，强制本地构建而非下载预编译 wheel，默认 0

更多示例与细节，请参考测试代码或查阅相应的 Python 文档。

致谢

DeepGEMM 受到 CUTLASS 项目的启发。向开发者们致以感谢与敬意！

许可证

本代码仓库基于 MIT License 发布。

引用

@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}},
}
开源项目deepseek-ai2026-10-06原文

相关内容