DeepGEMM

DeepGEMM 是一个专为干净高效的一般矩阵乘法（GEMM）设计的库。它支持 FP8 和 BF16 数据类型，适用于常规和混合专家（MoE）分组场景。该库使用 CUDA 编写，在安装时无需编译内核，而是通过一个轻量级的即时编译（JIT）模块在运行时动态编译所有内核。

DeepGEMM 借鉴了 CUTLASS 和 CuTe 中的一些概念，但避免了对其模板或代数系统的过度依赖。相反，该库以简洁为目标，仅包含少量核心内核函数，使其成为学习 NVIDIA GPU 内核优化技术的一个清晰且易于上手的资源。

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

新闻

• 2025年9月28日：DeepGEMM 现已支持用于 DeepSeek v3.2 光速索引器的打分内核（加权 ReLU MQA logits）。
  • 更多详情请参阅 #200。
• 2025年7月20日：DeepGEMM 现在同时支持 SM90 和 SM100 架构，并进行了全面重构，引入了一个低 CPU 开销的 JIT CPP 模块。
  • NVRTC 和后编译 SASS 优化已被全部禁用。
  • NVRTC 将在后续版本中重新支持。
  • 由于 NVCC 12.9 会自动进行 FFMA 交错处理，所有后优化功能将不再适用。
  • 更多详情请参阅 #112。
• 2025年5月14日：DeepGEMM 现已提供密集模型和 MoE 反向传播的权重梯度内核！详情请见 #95。
• 2025年5月7日：DeepGEMM 现已支持 NVRTC，编译速度最高可提升 10 倍！详情请参阅 #94。请使用 DG_JIT_USE_NVRTC=1 来启用此功能（某些情况下可能会有性能损失）。
• 2025年4月18日：DeepGEMM 在 H800 上已达到高达 1550 TFLOPS 的性能！详情请参阅 #74、#78、#81、#86 以及 340d988。

路线图

• 针对分组连续布局的更多正确性测试
• 输出端共享内存交错
• 兼容 TMA 组播的 MoE 调度器
• 修复不可分割形状下的 TMA 组播兼容性问题
• 跳过 M 维度上的无用计算
• 使用 NVRTC 作为更快的编译器
• 测试用的 Sanitizer 工具
• 密集模型的权重梯度内核
• MoE 模型的权重梯度内核
• 改进 get_best_configs 建模
• 支持 CUDA PDL
• 为部分形状增加更大的 TMA 组播尺寸
• 使用 CUTLASS 重构 MMA 模板
• 移除 N 和 K 维度的形状限制
• 实现 BF16 内核
• 分割/流式优化
• Ampere 架构专用内核
• 完善文档

快速入门

需求

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

开发

# 必须克隆子模块
git clone --recursive git@github.com:deepseek-ai/DeepGEMM.git
cd DeepGEMM

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

# 测试所有 GEMM 实现
python tests/test_layout.py
python tests/test_attention.py
python tests/test_core.py

安装

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 格式，即 4 个 UE8M0 打包成一个 torch.int。

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

普通密集 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 函数文档。

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

分组 GEMM（掩码布局）

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

为此可使用 m_grouped_fp8_gemm_nt_masked，并参考相关文档。一个典型用法是将来自 DeepEP 的低延迟内核输出作为输入。

V3.2 MQA 内核用于索引器

该内核系列有两个版本：非分页版（用于预填充）和分页版（用于解码）。

以非分页版 fp8_mqa_logits 为例，它有 6 个输入：

• q，形状为 [seq_len, num_heads, head_dim] 的 E4M3 张量
• kv，形状为 [seq_len_kv, head_dim] 的 E4M3 张量，以及形状为 [seq_len_kv] 的浮点 SF 张量
• weights，形状为 [seq_len, num_heads] 的浮点张量
• cu_seq_len_k_start 和 cu_seq_len_k_end，形状为 [seq_len] 的整数张量
• 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，并按如下方式计算 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。

工具函数

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

• deep_gemm.set_num_sms: 设置使用的最大 SM 数量
• deep_gemm.get_num_sms: 获取当前设置的最大 SM 数量（若未设置，则返回设备的 SM 总数）
• deep_gemm.set_tc_util: 设置近似的 Tensor Core 利用率
• deep_gemm.get_tc_util: 获取当前的 Tensor Core 利用率
• deep_gemm.transform_sf_into_required_layout: 将缩放因子转换为所需布局
• deep_gemm.get_tma_aligned_size: 获取所需的 TMA 对齐大小
• deep_gemm.get_mk_alignment_for_contiguous_layout: 获取分组连续布局的组级对齐要求
• 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
• JIT 缓存相关
  • DG_JIT_CACHE_DIR: 字符串，用于存储编译后内核的缓存目录，默认为 $HOME/.deep_gemm
• NVCC/NVRTC 选择
  • DG_JIT_USE_NVRTC: 0 或 1，使用 NVRTC 替代 NVCC，编译速度更快，但在某些情况下性能可能较低，默认为 0
  • DG_JIT_NVCC_COMPILER: 字符串，指定 NVCC 编译器路径；默认会在 torch.utils.cpp_extension.CUDA_HOME 中查找
• 编译选项
  • DG_JIT_PTXAS_VERBOSE: 0 或 1，显示详细的 PTXAS 编译器输出，默认为 0
  • DG_JIT_PRINT_COMPILER_COMMAND: 0 或 1，打印 NVCC 编译命令，默认为 0
• 启发式选择
  • DG_PRINT_CONFIGS: 0 或 1，打印针对每种形状所选的配置，默认为 0

更多示例和详细信息，请参阅 测试代码 或相应的 Python 文档。

致谢

DeepGEMM 灵感来源于 CUTLASS 项目。感谢并致敬其开发者！

许可证

本代码仓库根据 MIT 许可证 发布。

DeepGEMM 快速上手指南

DeepGEMM 是一个专为 NVIDIA GPU 设计的高效通用矩阵乘法（GEMM）库，支持 FP8 和 BF16 精度，特别针对混合专家模型（MoE）的分组场景进行了优化。它采用轻量级 JIT（即时编译）机制，无需预编译内核，代码简洁且性能卓越。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 硬件架构：NVIDIA SM90 (Hopper, 如 H800/H100) 或 SM100 (Blackwell) GPU。
• 操作系统：Linux。
• Python：3.8 或更高版本。
• 编译器：支持 C++20 标准的编译器。
• CUDA Toolkit：
  • SM90：CUDA 12.3+（强烈推荐 12.9+ 以获得最佳性能）。
  • SM100：CUDA 12.9+。
• 深度学习框架：PyTorch 2.1 或更高版本。
• 依赖库：
  • CUTLASS 4.0+
  • {fmt} 库
  • (注：这两个库可通过 Git 子模块自动获取)

安装步骤

DeepGEMM 通过源码安装，利用 Git 子模块管理依赖。

1. 克隆仓库（包含子模块）

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

   提示：如果网络连接 GitHub 较慢，可配置国内 Git 镜像或使用代理加速克隆过程。

2. 构建开发环境与 JIT 模块 运行提供的脚本以链接必要的头文件并构建 C++ JIT 模块：

   cat develop.sh
   ./develop.sh
   

3. 正式安装 执行安装脚本完成最终配置：

   cat install.sh
   ./install.sh
   

4. 验证安装（可选） 运行测试脚本确保 GEMM 实现正常工作：

   python tests/test_layout.py
   python tests/test_attention.py
   python tests/test_core.py
   

基本使用

安装完成后，即可在 Python 项目中导入 deep_gemm 使用。DeepGEMM 的核心接口遵循 $D = C + A \times B$ 的命名规范，默认输入布局为 NT（A 非转置，B 转置）。

最简单的使用示例：普通稠密矩阵乘法 (FP8)

以下示例展示如何调用基础的 FP8 GEMM 接口。请注意，缩放因子（Scaling Factors）需要根据架构要求进行特定的内存布局转换。

import torch
import deep_gemm

# 假设已准备好符合要求的 FP8 张量 A, B 和累加器 C
# 注意：实际使用中需确保张量维度对齐，并正确处理缩放因子 (scale)
# 此处仅为接口调用示意

# 定义输入形状 (示例)
M, N, K = 1024, 1024, 1024

# 创建模拟数据 (实际使用时应为 E4M3/E5M2 等 FP8 格式)
# A: [M, K], B: [N, K] (因为接口期望 B 是转置传入的，即原始 B 为 [K, N])
A = torch.randn(M, K, dtype=torch.float8_e4m3fn, device='cuda')
B = torch.randn(N, K, dtype=torch.float8_e4m3fn, device='cuda') 
C = torch.zeros(M, N, dtype=torch.bfloat16, device='cuda') # 输出通常为高精度

# 准备缩放因子 (SM90 需要 FP32 格式且需 TMA 对齐，具体请参考 transform_sf_into_required_layout)
# 这里假设 scale_a 和 scale_b 已经过 deep_gemm 工具函数处理
scale_a = torch.ones((M,), dtype=torch.float32, device='cuda') 
scale_b = torch.ones((N,), dtype=torch.float32, device='cuda')

# 执行 NT 布局的矩阵乘法: D = C + A @ B.T
# 对应函数：fp8_gemm_nt
deep_gemm.fp8_gemm_nt(
    a=A,
    b=B,
    c=C,
    scale_a=scale_a,
    scale_b=scale_b,
    # 其他可选参数如 alpha, beta 等视具体需求而定
)

print("GEMM 计算完成")

关键注意事项

1. 内存布局：

   • SM90：仅支持 NT 布局（Row-major A, Col-major B）。
   • SM100：支持所有布局（NT, TN, NN, TT）。
   • 如果需要其他布局（如 NN），请在调用前自行对输入张量进行转置或使用对应的 _nn, _tn 后缀接口（若架构支持）。

2. 缩放因子 (Scaling Factors)：

   • 库不负责输入数据的 FP8 转换或转置，这些操作需在调用前由用户完成。
   • SM90：缩放因子必须为 FP32 格式。
   • SM100：缩放因子必须为打包的 UE8M0 格式（每 4 个值打包进一个 torch.int）。
   • 推荐使用库提供的工具函数 deep_gemm.transform_sf_into_required_layout 来处理缩放因子布局。

3. MoE 分组支持：

   • 对于 MoE 场景，库提供了 m_grouped_fp8_gemm_*_contiguous（连续布局）和 m_grouped_fp8_gemm_*_masked（掩码布局）接口，专门用于处理专家共享形状但 Token 数量动态变化的场景。

更多高级用法（如权重梯度计算、MQA 索引器等）请参考 tests/ 目录下的测试代码或 Python 文档字符串。