deepseek-ai/DeepGEMM

# DeepGEMM DeepGEMM 是一个统一的、高性能的张量核心内核库，将现代大型语言模型的关键计算原语——GEMM（FP8、FP4、BF16）、融合 MoE（重叠通信的 Mega MoE）、用于闪电索引器的 MQA 评分、HyperConnection（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、FP8xFP4 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 现在提供稠密和 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 工具包： - 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 ``` ### 安装 ``` cat install.sh ./install.sh ``` 然后，在 Python 项目中导入 `deep_gemm` 并开始使用！ ## 接口 #### 注意事项 本库为 NVIDIA GPU 提供了基于命名约定 `D = C + A @ B` 的优化 GEMM 内核。输入布局为 NT（非转置 A，转置 B）。虽然 SM90 实现仅支持 NT 内存布局（行主序，列主序），但 SM100 实现支持所有内存布局（NT、TN、NN、TT）。例如，`fp8_gemm_nt` 将执行 `D = C + A @ B.T`。 对于两种架构，LHS 缩放因子需要具有 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 内核本身。 #### 普通稠密 GEMM（非分组） 要执行基本的非分组 FP8 GEMM，请调用 `fp8_gemm_{nt, nn, tn, tt}` 函数。更多细节请参考函数文档。 #### 分组 GEMM（连续布局） 与 CUTLASS 中的传统分组 GEMM 不同，DeepGEMM 仅对 M 轴进行分组，而 N 和 K 必须保持固定。这种设计专为 MoE 模型中专家共享相同形状的场景而定制。对于训练前向传播或推理预填充，其中每个专家可能处理数量可变的令牌，我们将这些令牌连接成一个张量，称为“连续”布局。请注意，每个专家段必须对齐到 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 不知道每个专家将接收多少令牌时，我们支持掩码分组 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]` - `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`， 并计算： ``` 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 ``` 更多细节以及分页版本 `fp8_paged_mqa_logits` 请参考 `tests/test_attention.py`。 #### Mega MoE Mega MoE 融合并重叠了 EP 分派、线性 1（FP8xFP4）、SwiGLU、线性 2（FP8xFP4）和 EP 合并，形成一个巨大的内核，同时重叠 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` 更多示例和细节请参考 [测试代码](tests/test_core.py) 或查阅相应的 Python 文档。 ## 致谢 DeepGEMM 灵感来源于 [CUTLASS](https://github.com/nvidia/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}}, } ```

标签：BF16, CUDA, DLL 劫持, DNS解析, FP4, FP8, GEMM, GPU 加速, HC, HyperConnection, JIT, Mega MoE, MoE, MQA, NVRTC, SASS优化, SM100, SM90, Vectored Exception Handling, 低CPU开销, 内核优化, 凭据扫描, 反向传播, 多查询注意力, 大语言模型, 开源项目, 张量核心, 性能比较, 权重梯度, 深度学习, 矩阵乘法, 索引器, 融合专家, 运行时编译, 逆向工具, 高性能计算