DeepSeek 发布 DeepGEMM 重大更新：Mega MoE 融合算子与 FP4 精度支持

4 月 16 日，DeepSeek 开源了其高性能 CUDA 算子库 DeepGEMM 的重大更新（PR #304），新增 Mega MoE 融合算子、FP8×FP4 GEMM、FP4 Indexer、PDL 支持等，同时在 JIT 编译速度和 MoE GEMM 性能上做了显著优化。

DeepGEMM 目前 6.3k star、861 fork，是专为现代大模型（尤其是 DeepSeek V3/R1 系列的 MoE 架构）设计的底层算子库，支持 NVIDIA SM90（Hopper）和 SM100（Blackwell）架构。

Mega MoE：把整个 MoE 层融成一个核函数

这次更新的核心是 Mega MoE——一个将 MoE 前向传播的多个阶段融合为单一 mega-kernel 的技术。

传统 MoE 推理的计算流程是分阶段的：先做 EP（Expert Parallelism）dispatch 把 token 分发到对应专家，每个专家独立做 linear1 → SwiGLU → linear2，最后 combine 汇总结果。每个阶段之间有同步点，GPU 的计算单元和 NVLink 通信带宽无法充分重叠。

Mega MoE 的做法是把 dispatch、linear1（FP8×FP4）、SwiGLU、linear2（FP8×FP4）、combine 全部融合到一个核函数里，让 NVLink 通信和 Tensor Core 计算在时间上重叠。这类似于 GPU kernel fusion 的经典思路，但应用范围更大——不是融合两三个算子，而是融合整个 MoE 层。

实现上有个技术前提：需要多进程启动 + 对称内存（Symmetric Memory）。多个 GPU 通过 symmetric buffer 共享数据，每个 rank 的 expert 权重预转换到 FP4 布局，运行时直接从 symmetric buffer 读取输入、写入输出，避免显式通信同步。

目前的限制：仅支持 FP8×FP4 MoE（activation FP8、weight FP4），EP ≤ 8 已测试，需要 PyTorch ≥ 2.9。

FP4 精度支持

除了 Mega MoE 之外，DeepGEMM 新增了 FP4 相关的完整支持：

• FP8×FP4 GEMM：activation FP8、weight FP4 的通用矩阵乘法，FP4 权重使用 UE8M0 缩放因子打包格式（4 个 UE8M0 打包进一个 int32）。这意味着 MoE 模型的专家权重可以量化到 4-bit，显存占用减半。
• FP4 Indexer：基于 MQA logits 的专家选择（scoring kernel）也支持 FP4 精度，配合 larger MTP（Multi-Token Prediction）支持。
• 权重预转换工具：transform_weights_for_mega_moe() 接口将 FP4 权重转换为 Mega MoE 所需的内存布局。

从 FP8 到 FP4，意味着在 H800 上单卡能塞下更多专家，或者在同等模型规模下减少 GPU 卡数。对 DeepSeek V3 这种 671B 参数、256 个路由专家的 MoE 架构，4-bit 权重量化的工程价值非常直接。

其他更新

• PDL（Programmatic Dependent Launch）：一种减少 kernel launch 开销的技术，对高频调用的小 kernel 有意义。
• GEMM 启发式重构 + A/B 交换优化：让 MoE GEMM 性能显著提升。
• JIT 编译加速：首次运行时的编译时间缩短。
• DeepEPv2 MoE GEMM 布局：兼容 DeepSeek 的 EP 通信库 DeepEPv2。
• 代码量：PR #304 改动 12,135 行新增 / 3,253 行删除，涉及 109 个文件。

定位与影响

DeepGEMM 的定位很清晰：不是一个通用 BLAS 库，而是专为 MoE 大模型的训练和推理优化的专用算子集合。它的设计理念是"轻量 + JIT"——安装时不需要编译，运行时按需即时编译，根据实际 shape 自动选择最优 kernel 配置。

从 DeepSeek 的工程路线看，他们一直在垂直整合 MoE 推理的每一层：DeepEP 做跨卡通信，DeepGEMM 做单卡计算，Mega MoE 把两者通过 symmetric memory 和 kernel fusion 打通。这是一个完整的 MoE 推理栈。

性能数据方面，README 中标注"性能对比稍后发布"，此前 DeepGEMM 在 H800 上已达到 1550 TFLOPS。Mega MoE 在减少阶段间同步开销、重叠通信与计算之后，实际提升幅度还需要等官方 benchmark。

对社区来说，这个更新的直接价值在于：如果你在做 MoE 模型的推理优化（不管是 DeepSeek 架构还是类似的 expert-parallel 路线），DeepGEMM 提供了一个可以参考甚至直接使用的 kernel 级实现。MIT 协议，无附加限制。

来源：GitHub - deepseek-ai/DeepGEMM · PR #304