@Ryrenz：重磅消息：DeepSeek 开源了 DeepGEMM——他们用于训练模型的 GPU kernel 库，从同样的显卡中挤出更多算力……

X AI KOLs Timeline 工具
gpu-kernels cuda deepseek fp8 moe open-source jit-compilation

摘要

DeepSeek 开源了其训练模型所用的 GPU kernel 库 DeepGEMM，将 FP8/FP4/BF16 GEMM、融合 MoE 通信重叠、MQA 打分等核心算子统一到单一 CUDA 代码库中，通过 DeepJIT 运行时编译，官方称性能可达 H800 上 1550 TFLOPS，并已推出昇腾（Ascend）版本。

🏆 重磅消息：DeepSeek 开源了 DeepGEMM——他们用于训练模型的 GPU kernel 库，从同样的显卡中挤出了更多算力。 在 GitHub 上，它已经获得了 8909 个 star 和 1401 个 fork，而 9 月 30 日，他们又顺手发布了昇腾（Ascend）版本。 凡是做过大模型推理的人都懂那些微调时的煎熬日子：一套 kernel 给 FP8 矩阵乘法用，另一套给 MoE 分发和通信重叠用，还有一套给索引器打分 kernel 用——每一套都要针对不同架构手工编写，更别提等 CUDA 编译环境那半个小时了。 DeepGEMM 把所有这些算子统一打包进了一个 CUDA 代码库：FP8、FP4、BF16 的矩阵乘法；融合了通信重叠的 MoE；甚至连 DeepSeek v3.2 闪电索引器中的打分 kernel 也包含在内。所有 kernel 在运行时由 DeepJIT 即时编译——安装时不需要编译步骤，帮你省下盯着编译器浪费的那截寿命。它借鉴了 CUTLASS 和 CuTe 的思路，但没有把你淹没在模板的海洋里；自己动手改也不会让你绝望。 在性能方面，官方声称它在各种矩阵形状下都能匹敌甚至超越专家调优的库，在 H800 上最高可达 1550 TFLOPS。 一个好的底层库不是堆功能——它只是把你的显卡原封不动地还给你，而且更快。 GitHub：
查看原文
查看缓存全文

缓存时间: 2026/10/11 19:16

🏆 炸裂消息：DeepSeek 开源了 DeepGEMM——他们训练模型时使用的 GPU kernel 库，让同一张卡榨出更多算力。在 GitHub 上已经收获 8909 个 star 和 1401 个 fork，9 月 30 号他们又顺手发布了 Ascend 版本。

做过大模型推理的朋友都懂调优时的煎熬：一套 kernel 给 FP8 矩阵乘，一套给 MoE 分发与通信重叠，还有一套给 indexer 打分 kernel——每一套都要针对不同架构手写，更别提等 CUDA 编译环境那半小时了。DeepGEMM 把这些原语全部打包进一个统一的 CUDA 代码库：FP8、FP4、BF16 矩阵乘，带通信重叠的融合 MoE，甚至 DeepSeek v3.2 闪电 indexer 的打分 kernel。所有 kernel 在运行时由 DeepJIT 即时编译——安装时无需编译步骤，省下你盯着编译器发呆的那段生命。它借鉴了 CUTLASS 和 CuTe 的思路，但没有让你淹没在模板的汪洋大海里，自己改起来也不会绝望。性能方面，官方宣称在各种矩阵形状下匹配甚至超越专家调优的库，在 H800 上最高可达 1550 TFLOPS。一个好的底层库不会堆砌功能——它只是把你的卡还给你，而且更快。

GitHub：

deepseek-ai/DeepGEMM

来源：https://github.com/deepseek-ai/DeepGEMM

DeepGEMM

DeepGEMM 是一个统一的高性能 Tensor Core kernel 库，将现代大语言模型的核心计算原语——GEMM（FP8、FP4、BF16）、带通信重叠的融合 MoE（Mega MoE）、闪电 indexer 的 MQA 打分、HyperConnection（HC）等——整合到一个统一的 CUDA 代码库中。所有 kernel 在运行时通过 DeepJIT 编译，安装时无需 CUDA 编译。DeepGEMM 借鉴了 CUTLASS (https://github.com/nvidia/cutlass) 和 CuTe (https://github.com/NVIDIA/cutlass/tree/main/include/cute) 的部分概念，但避免了对其模板或代数体系的重度依赖。该库设计追求简洁，核心 kernel 函数数量有限，是学习 NVIDIA GPU kernel 优化技术的干净且易于上手的资源。尽管设计轻量，DeepGEMM 在各种矩阵形状下的性能仍能匹配甚至超越专家调优的库。

新闻

  • 2026.09.30 - DeepGEMM Ascend 已上线！详情请查看 DeepGEMM-Ascend (https://github.com/deepseek-ai/DeepGEMM-Ascend/)
    • 新增更多优化，包括 locality domain 特性，详见 #462 (https://github.com/deepseek-ai/DeepGEMM/pull/462)
  • 2026.09.10：Sparse Indexer、Mega Gate、Mega mHC、DeepJIT、MoE 与 Indexer 优化等。
    • 详见 #432 (https://github.com/deepseek-ai/DeepGEMM/pull/432)。
  • 2026.04.16：Mega MoE、FP8xFP4 GEMM、FP4 Indexer、PDL、更快的 JIT 编译等。
    • 详见 #304 (https://github.com/deepseek-ai/DeepGEMM/pull/304)。
    • Mega MoE 基准测试请参考 #316 (https://github.com/deepseek-ai/DeepGEMM/pull/316)。
  • 2025.09.28：DeepGEMM 现已支持 DeepSeek v3.2 闪电 indexer 的打分 kernel（加权 ReLU MQA logits）。
    • 详见 #200 (https://github.com/deepseek-ai/DeepGEMM/pull/200)。
  • 2025.07.20：DeepGEMM 现已支持 SM90/SM100，并进行了全面重构，引入低 CPU 开销的 JIT CPP 模块。
    • 由于 NVCC 12.9 将自动进行 FFMA 交错优化，后续优化不再需要支持。
    • 详见 #112 (https://github.com/deepseek-ai/DeepGEMM/pull/112)。
  • 2025.05.14：DeepGEMM 现已提供 dense 和 MoE 反向传播的权重梯度 kernel！详见 #95 (https://github.com/deepseek-ai/DeepGEMM/pull/95)。
  • 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 12.9 或更高版本
  • PyTorch 2.3 或更高版本
  • CUTLASS 4.0 或更高版本（可通过 Git submodule 克隆）

开发安装

# 必须克隆 submodule
git clone --recursive [email protected]:deepseek-ai/DeepGEMM.git
cd DeepGEMM
# 链接必要的头文件并构建 C++ 扩展
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 内存布局（行主序、列主序），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 类型转换等操作需要用户单独处理，请自行实现或融合到前序 kernel 中。虽然本库提供了一些简单的 PyTorch 工具函数，但这些可能导致性能下降，我们的核心优化目标是 GEMM kernel 本身。

普通 dense GEMM（非分组）

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

分组 GEMM（连续布局）

与 CUTLASS 中的传统分组 GEMM 不同，DeepGEMM 仅对 M 轴进行分组，N 和 K 必须保持固定。这种设计专为 MoE 模型中各专家形状相同的场景而定制。在训练前向传播或推理 prefilling 阶段，每个专家处理的 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（masked 布局）

在推理解码阶段，启用 CUDA Graph 且 CPU 不知道每个专家接收多少 token 时，我们支持 masked 分组 GEMM。通过提供 mask 张量，kernel 仅计算有效部分。使用 m_grouped_fp8_gemm_nt_masked 实现此功能，相关详情请参阅文档。一个典型用法是使用 DeepEP (https://github.com/deepseek-ai/DeepEP) 低延迟 kernel 的输出作为输入。

V3.2 Indexer MQA Kernel

该 kernel 系列有两个版本：非分页版本（用于 prefilling）和分页版本（用于 decoding）。以非分页版本 fp8_fp4_mqa_logits 为例，其主要输入为：

  • q，一个 (q_data, q_sf) 元组；SM100 接受 MXFP4/MXFP8 数据及打包的 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，形状为 [seq_len] 的整型张量
  • 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()                       # 标量

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

Mega MoE

Mega MoE 将 EP dispatch、linear 1 和 linear 2（FP8xFP4 或 FP8xFP8）、SwiGLU 和 EP combine 融合并重叠为单个超大 kernel，实现 NVLink 通信与 Tensor Core 计算的重叠。需要多进程启动并使用对称内存。

用法：

# 分配对称内存缓冲区
# 注意：需要 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',  # FP8 路由专家权重请使用 'fp8xfp8'
)

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

# （可选）将权重局部化到 locality domain 中
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()

# 每次调用前将输入复制到缓冲区中
# 可以将这些操作融合到前序 kernel 中
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 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。

工具函数

除了上述 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：启用/禁用程序化依赖启动（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 主序 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 打包 kernel

本库还提供一些环境变量，可能对调试有用：

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

  • 通用
    • DG_JIT_DEBUG：0 或 1，启用 JIT 调试功能，包括编译器命令和 PTXAS 输出、加载时间报告、行号信息和 PTX/SASS 转储；默认 0
    • DG_PRINT_CONFIGS：0 或 1，打印每种形状选择的配置，默认 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 无本地内存使用，默认 0
    • DG_JIT_PRINT_LOAD_TIME：0 或 1，打印 kernel 加载时间，默认 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

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

致谢

DeepGEMM 受到 CUTLASS (https://github.com/nvidia/cutlass) 项目的启发。感谢并致敬 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-ai/DeepGEMM

GitHub Trending (daily)

# deepseek-ai/DeepGEMM 来源：[https://github.com/deepseek-ai/DeepGEMM](https://github.com/deepseek-ai/DeepGEMM) # DeepGEMM DeepGEMM 是一个统一的高性能张量核心内核库，整合了现代大语言模型的关键计算原语——GEMM（FP8、FP4、BF16）、带通信重叠的融合 MoE（Mega MoE）、闪电索引器的 MQA 评分、超连接（HC）等——形成一个统一的 CUDA 代码库。所有内核均在运行时编译

@YRSM_Simon: 太厉害了

X AI KOLs Timeline

shi3z分享在不支持FP4的A100 GPU上运行DeepSeek v4.1 Flash的优化经验，将推理速度从33 tok/s提升至673 tok/s，超过官方API速度。