DeepGEMM Vibe 项目

DeepGEMM：DeepSeek 开源的 FP8/FP4 高性能 CUDA 矩阵运算核心库（B 线 Vibe）

项目简介

DeepGEMM 是 DeepSeek 团队开源的统一高性能 GPU 张量核心（Tensor Core）算子库，专注于为大语言模型训练与推理提供核心矩阵乘法（GEMM）原语。项目用极简的 C++/CUDA 代码把当前主流 LLM 的关键计算内核整合到同一个代码库，包括 FP8/FP4/BF16 精度的 GEMM、带通信重叠的融合 MoE（Mega MoE）、闪电索引器的 MQA 打分内核、以及 HyperConnection 等新算子。所有算子在运行时通过 DeepJIT 即时编译，无需安装 CUDA 编译工具链。截至 2026-09，最新版本已支持 NVIDIA SM90（Hopper）与 SM100（Blackwell）双架构，并在 H800 上实现最高 1550 TFLOPS 的实测算力。MIT 开源协议下 GitHub 仓库已收获 8600+ Star 与 1300+ Fork，是国产 AI 基础设施类项目中关注度最高的内核库之一。

核心功能

• 统一 GEMM 内核族：支持 FP8、FP4、BF16 三种主流精度的稠密矩阵乘法，提供 NT/TN/NN/TT 全内存布局接口（SM100），可作为 LLM 训练前向、反向与推理 Prefill/Decode 阶段的统一底层算子。
• 融合 MoE 与 Mega MoE：针对混合专家模型提供分组 GEMM（连续布局 + 掩码布局）实现，并进一步把 EP dispatch、Linear1/Linear2、SwiGLU、EP combine 融合进单 mega-kernel，把 NVLink 通信与张量核心计算重叠执行，显著提升专家路由效率。
• 闪电索引器 MQA 打分内核：为 DeepSeek V3.2 的 lightning indexer 提供加权 ReLU MQA logits 内核，分非分页（prefill）与分页（decode）两个版本，输出压缩为 [seq_len, max_seqlen_k] 的索引分数张量。
• Mega mHC / Mega Gate / Sparse Indexer：2026-09 最新更新新增的多头连接、门控与稀疏索引器算子，全部走 FP8 路径，是 DeepSeek-V4 推测路线下的关键基础组件。
• DeepJIT 即时编译：基于 DeepJIT 在运行时按形状生成最优 kernel，无需预编译多种 shape 二进制；支持 PDL（Programmatic Dependent Launch）、确定性算法、SM 占用上限、TC 利用率上限等运行时调参。
• 昇腾 NPU 移植版本：官方同时开源 DeepGEMM-Ascend，将上述算子平移到华为昇腾平台，方便国产算力栈直接复用相同的接口与精度策略。

技术实现

DeepGEMM 在 CUTLASS 与 CuTe 的设计思想上做了大幅简化：仅保留最少量核心 kernel 函数，避免依赖繁重的模板元编程与代数封装，使得内核源码可读、可改、可教学。在算法层面，所有 FP8 GEMM 通过 TMA（Tensor Memory Accelerator）完成 HBM↔SMEM 的异步搬运，并支持 LHS scaling factor 的 UE8M0 packed 布局以降低访存开销。MoE 内核采用 M 轴分组、N/K 轴固定的布局，把同一 shape 的多个 expert token 拼接为 contiguous layout，最大化 tensor core 利用率。Mega MoE 则通过对称显存（symmetric memory）配合多进程启动，把 EP 通信与 GEMM 计算在同一个 mega-kernel 内流水化。运行时通过 DeepJIT 缓存（默认 $HOME/.dj）复用已编译的 PTX/SASS，并提供 DG_JIT_DEBUG / DG_PRINT_CONFIGS / DG_JIT_DUMP_PTX 等丰富的调试与 profiling 环境变量，便于内核调优与性能对比。

快速上手

1. 克隆仓库（含子模块）：git clone --recursive git@github.com:deepseek-ai/DeepGEMM.git，进入目录后执行 cat develop.sh 查看构建脚本说明。
2. 编译 C++ 扩展：运行 ./develop.sh 完成 CUDA 扩展的本地链接；如使用预构建 wheel，则设置 DG_FORCE_BUILD=0 让安装阶段跳过本地编译。
3. 安装运行时：./install.sh，或直接 pip install deep_gemm（已发布至 PyPI 与 Hugging Face Hub）。
4. 在 Python 项目中 import deep_gemm 调用 fp8_gemm_nt(a, b, c, a_sf, b_sf) 等接口即可运行 GEMM；MoE 场景调用 m_grouped_fp8_gemm_nt_contiguous 或 fp8_fp4_mega_moe；MQA 打分调用 fp8_fp4_mqa_logits。
5. 进阶调试时打开 DG_JIT_DEBUG=1 DG_JIT_DUMP_PTX=1 DG_PRINT_CONFIGS=1 即可在 stdout 看到编译命令、PTX 与按形状选择的 config。

适用人群

• 大模型训练/推理基础设施工程师：需要为 FP8/FP4 路线替换 cuBLAS 或自研算子，以榨干 H100/H800/B200 算力的工程团队。
• MoE 模型研究者：在做 DeepSeek-V3 / Qwen-MoE / Mixtral 等 MoE 架构实验，需要高性能分组 GEMM 与融合通信算子的实验室。
• 国产 NPU 适配团队：正在把 LLM 内核移植到昇腾、寒武纪等国产算力平台的工程师，可直接基于 DeepGEMM-Ascend 做二次开发。
• 算子库学习者：希望阅读一份「简洁、可读、不靠模板元编程」的高质量 CUDA kernel 教材，理解 tensor core、TMA、PDL、UE8M0 等前沿概念的开发者。
• AI 系统性能优化方向的硕博生：在论文或课程项目中需要对比 hand-tuned kernel 与编译器生成 kernel 性能差异的研究者。

官方链接

• GitHub 仓库：https://github.com/deepseek-ai/DeepGEMM
• 昇腾移植版：https://github.com/deepseek-ai/DeepGEMM-Ascend
• 引用 BibTeX：见 README Citation 段落（Chenggang Zhao et al., 2025）
• License：MIT