deepseek-ai/DeepGEMM

DeepGEMM

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

DeepGEMM 借鉴了CUTLASS（https://github.com/nvidia/cutlass）和CuTe（https://github.com/NVIDIA/cutlass/tree/main/include/cute）的部分概念，但避免过度依赖它们的模板或代数系统。该库设计简洁，核心内核函数数量有限，是学习NVIDIA GPU内核优化技术的清晰且易懂的资源。

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

新闻

• 2026.04.16: Mega MoE、FP8xFP8 GEMM、FP4索引器、PDL、更快的JIT编译等。
  • 性能对比将稍后发布。
  • 详见 #304 (https://github.com/deepseek-ai/DeepGEMM/pull/304)。
• 2025.09.28: DeepGEMM现在支持闪电索引器的评分内核（加权ReLU MQA logits），用于DeepSeek v3.2。
  • 详见 #200 (https://github.com/deepseek-ai/DeepGEMM/pull/200)。
• 2025.07.20: DeepGEMM现在支持SM90/SM100，并进行了全面重构，实现了低CPU开销的JIT CPP模块。
  • NVRTC和后编译SASS优化均已禁用。
  • NVRTC将在后续支持。
  • 由于NVCC 12.9会自动进行FFMA交织，所有后编译优化将不再支持。
  • 详见 #112 (https://github.com/deepseek-ai/DeepGEMM/pull/112)。
• 2025.05.14: DeepGEMM现在为dense和MoE反向传播提供权重梯度内核！详见 #95 (https://github.com/deepseek-ai/DeepGEMM/pull/95)。
• 2025.05.07: DeepGEMM现在支持NVRTC，编译速度提升高达10倍！详见 #94 (https://github.com/deepseek-ai/DeepGEMM/pull/94)。请使用 DG_JIT_USE_NVRTC=1 启用（某些情况下可能有性能损失）。
• 2025.04.18: DeepGEMM现在在H800上达到最高 1550 TFLOPS！详见 #74 (https://github.com/deepseek-ai/DeepGEMM/pull/74)、#78 (https://github.com/deepseek-ai/DeepGEMM/pull/78)、#81 (https://github.com/deepseek-ai/DeepGEMM/pull/81)、#86 (https://github.com/deepseek-ai/DeepGEMM/pull/86) 和 340d988 (https://github.com/deepseek-ai/DeepGEMM/commit/340d9880f4a418d943d34260d20a79f41f4c0526)。

快速开始

环境要求

• NVIDIA SM90或SM100架构GPU
• Python 3.8或更高版本
• 支持C++20的编译器
• CUDA Toolkit：
  • SM90需要CUDA 12.3或更高版本
    • 我们强烈建议使用12.9或更高版本以获得最佳性能
  • SM100需要CUDA 12.9或更高版本
• PyTorch 2.1或更高版本
• CUTLASS 4.0或更高版本（可通过Git子模块克隆）
• {fmt} 库（可通过Git子模块克隆）

开发

# 必须克隆子模块
git clone --recursive [email protected]:deepseek-ai/DeepGEMM.git
cd DeepGEMM

# 链接一些必要的头文件并构建CPP JIT模块
cat develop.sh
./develop.sh

安装

cat install.sh
./install.sh

然后在你的Python项目中导入 deep_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

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

• SM90需要FP32格式的缩放因子。
• SM100需要打包的UE8M0格式（https://docs.nvidia.com/cuda/parallel-thread-execution/#alternate-floating-point-data-formats）的缩放因子，即将4个UE8M0打包为单个 torch.int。

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

普通dense GEMM（非分组）

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

分组GEMM（连续布局）

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

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

分组GEMM（掩码布局）

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

请使用 m_grouped_fp8_gemm_nt_masked，并参阅相关文档。使用示例是将DeepEP（https://github.com/deepseek-ai/DeepEP）的低延迟内核输出作为输入。

V3.2 MQA内核用于索引器

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

• q，E4M3张量，形状为 [seq_len, num_heads, head_dim]
• kv，E4M3张量（形状为 [seq_len_kv, head_dim]），带有浮点SF（形状为 [seq_len_kv]）
• weights，浮点张量，形状为 [seq_len, num_heads]
• cu_seq_len_k_start 和 cu_seq_len_k_end，形状为 [seq_len] 的int张量
• clean_logits，是否将未填充的logits清理为 -inf

输出张量形状为 [seq_len, seq_len_kv]，表示token到token的logits。 对于 q 中的每个token i，它将遍历 [cu_seq_len_k_start[i], cu_seq_len_k_end[i]) 中的所有token j， 并将logit out[i, j] 计算为：

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()  # 标量

更多详情和分页版本 fp8_paged_mqa_logits，请参阅 tests/test_attention.py。

Mega MoE

Mega MoE将EP调度、线性1（FP8xFP4）、SwiGLU、线性2（FP8xFP4）和EP合并融合并重叠到单个 mega 内核中，交叠NVLink通信和张量核心计算。它需要多进程启动和对称内存。使用方法：

# 分配对称内存缓冲区
# 注意：需要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
)

# 将权重（FP4和UE8M0 SF）转换为所需布局
transformed_l1, transformed_l2 = deep_gemm.transform_weights_for_mega_moe(l1_weights, l2_weights)

# 每次调用前将输入复制到缓冲区
# 你可以将这些融合到前面的内核中
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)

# 运行融合的mega MoE内核
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：启用/禁用程序化依赖启动（PDL）
• 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主序TMA对齐张量
• deep_gemm.get_mn_major_tma_aligned_packed_ue8m0_tensor：获取MN主序TMA对齐张量（将FP32打包为UE8M0）
• deep_gemm.get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor：K分组GEMM打包内核

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

• 通用
  • DG_JIT_DEBUG：0 或 1，打印JIT调试信息，默认为 0
  • DG_PRINT_CONFIGS：0 或 1，打印每个形状选择的配置，默认为 0
• JIT缓存
  • DG_JIT_CACHE_DIR：字符串，编译内核的缓存目录，默认为 $HOME/.deep_gemm
• 编译器选择
  • DG_JIT_USE_NVRTC：0 或 1，使用NVRTC代替NVCC（编译更快，某些情况下性能可能较低），默认为 0
  • DG_JIT_NVCC_COMPILER：字符串，NVCC编译器路径；默认为 torch.utils.cpp_extension.CUDA_HOME
  • DG_JIT_CPP_STANDARD：整数，C++标准版本，默认为 20
• 编译器输出
  • DG_JIT_PRINT_COMPILER_COMMAND：0 或 1，打印编译命令，默认为 0
  • DG_JIT_PTXAS_VERBOSE：0 或 1，显示详细的PTXAS输出，默认为 0
  • DG_JIT_PTXAS_CHECK：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，强制本地构建而不是下载预构建的wheel，默认为 0
  • DG_JIT_USE_RUNTIME_API：0 或 1，使用CUDA Runtime API加载内核（需要CUDA runtime >= 12.8），默认为 0

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

致谢

DeepGEMM的灵感来自CUTLASS（https://github.com/nvidia/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}},
}