DeepGEMM质变：从GEMM库到统一LLM计算原语平台

自2025年初首次亮相以来，DeepGEMM一直是DeepSeek团队面向NVIDIA GPU（SM90/SM100）打造的高性能Tensor Core内核库。

然而，本次PR #304标志着DeepGEMM完成了一次根本性的蜕变——从一个”干净高效的GEMM库”升级为”现代大语言模型的统一计算原语平台”。其README中的自我定位已从“a library designed for clean and efficient GEMMs”正式更新为“a unified, high-performance tensor core kernel library that brings together the key computation primitives of modern large language models”。

• [Public release 26/04] Introducing Mega MoE, FP4 Indexer and other features/fixes
• Mega MoE仍处于开发与优化阶段，欢迎关注并提出优化建议！
• 免责声明：本次发布仅涉及DeepGEMM的开发进展，与内部模型发布无关。
• https://github.com/deepseek-ai/DeepGEMM/pull/304

本次发布由9位贡献者协作完成，共涉及109个文件、新增代码超过12000行。核心变化可归纳为三大方向：Mega MoE融合内核、FP4 Indexer（MQA logits），以及全面的GEMM启发式重构与工程优化。下面将逐一进行深入解析。

本文目录

• 一、Mega MoE：通信与计算深度重叠的”超级内核”
• 1.1 问题本质：MoE推理的NVLink通信瓶颈
• 1.2 解决方案：五阶段融合 + 通信/计算重叠
• 1.3 Buffer布局设计
• 二、FP4 Indexer：低精度MQA Logits计算
• 2.1 问题背景
• 2.2 关键实现
• 三、动态Swap A/B：MoE GEMM的关键性能优化
• 3.1 问题：SM100上MoE场景的Multicast失效
• 3.2 解决方案：在SM100上交换A/B操作数
• 四、全面的启发式重构
• 4.1 从硬编码到结构化配置
• 4.2 SM90引入L1/L2带宽建模
• 五、其他重要改进
• 5.1 PDL（Programmatic Dependent Launch）
• 5.2 JIT编译加速与分布式文件系统修复
• 5.3 FP8 × FP4 GEMM
• 5.4 C++20标准升级
• 总结

一、Mega MoE：通信与计算深度重叠的”超级内核”

1.1 问题本质：MoE推理的NVLink通信瓶颈

在Expert Parallelism（EP）模式下，MoE层的执行包含五个串行步骤：dispatch → linear1 (FP8×FP4) → SwiGLU → linear2 (FP8×FP4) → combine。

其中，dispatch和combine步骤涉及跨节点的NVLink通信。传统做法是将通信与计算拆分为独立的kernel串行执行，这导致NVLink带宽利用率低下，同时SM在等待通信完成时处于空转状态。

1.2 解决方案：五阶段融合 + 通信/计算重叠

Mega MoE的核心创新在于将上述五个步骤融合进单个mega-kernel，在kernel内部实现NVLink通信与Tensor Core计算的流水线重叠。这一设计依赖于以下关键技术：

关键技术1：对称内存（Symmetric Memory）架构

在csrc/apis/mega.hpp中可以看到，Mega MoE使用了PyTorch >= 2.9提供的对称内存分配机制，所有rank共享同一地址空间的buffer。SymBuffer结构（见deep_gemm/include/deep_gemm/layout/sym_buffer.cuh）通过offset映射实现跨rank指针翻译：

// deep_gemm/include/deep_gemm/layout/sym_buffer.cuh  
template <uint32_t kNumRanks = kNumMaxRanks>  
struct SymBuffer {  
int64_t base;  
int64_t offsets[kNumMaxRanks];  
uint32_t rank_idx;  

template <typename ptr_t>  
CUTLASS_DEVICE ptr_t map(const ptr_t& ptr, const uint32_t& dst_rank_idx) const {  
int64_t mapped_ptr = offsets[dst_rank_idx] + reinterpret_cast<int64_t>(ptr);  
return *reinterpret_cast<ptr_t*>(&mapped_ptr);  
}  
};  

这使得GPU线程可以直接通过地址偏移读写远端节点内存，完全无需经过host端协调。

关键技术2：三阶段线程分工

在 csrc/jit_kernels/heuristics/mega_moe.hpp 文件中可以观察到，Mega MoE 内核在其内部将线程划分为三种不同的角色：

• 调度线程组（128 线程）：专门负责 NVLink 通信，实现跨 rank 的 token 分发与收集操作。
• 非尾声线程组（128 线程）：驱动 TMA 数据加载以及 UMMA（Tensor Core）计算的主循环。
• 尾声线程组（256 线程）：执行 SwiGLU 激活函数、FP8 量化以及 TMA 输出的写回任务。

const int num_dispatch_threads = 128;
const int num_non_epilogue_threads = 128;
const int num_epilogue_threads = 256;

关键技术 3：Wave 级专家调度机制

为达成负载均衡，Mega MoE 引入了 num_experts_per_wave 这一概念。该机制的核心在于每个“wave”处理一批专家，并通过启发式计算（考量不均衡因子 kImbalanceFactor = 2）来选定合适的 wave 大小，从而确保所有流多处理器（SM）都有足够的 block 可供执行：

// csrc/jit_kernels/heuristics/mega_moe.hpp#L83-L110
static int get_num_experts_per_wave_for_mega_moe(...) {
constexpr int kImbalanceFactor = 2;
// ...
int num_experts_per_wave = num_l1_blocks_per_expert > 0
? ceil_div(kImbalanceFactor * num_sms, num_l1_blocks_per_expert) : 1;
// Round up to nearest divisor of num_experts_per_rank
while (num_experts_per_wave < max_num_experts_per_wave
and num_experts_per_rank % num_experts_per_wave != 0)
        ++ num_experts_per_wave;
return num_experts_per_wave;
}

关键技术 4：NVLink Barrier 实现

在 deep_gemm/include/deep_gemm/comm/barrier.cuh 文件中，实现了一套完整的 NVLink 跨 rank 同步原语。这套原语涵盖了 grid 内同步、基于 ptx::red_add_rel_sys 的跨节点原子信号通知，并且内置了 30 秒的超时检测机制：

// deep_gemm/include/deep_gemm/comm/barrier.cuh
template <uint32_t kNumRanks, uint32_t kNumSMs, ...>
CUTLASS_DEVICE void nvlink_barrier(...) {
// Grid sync before NVLink signaling
grid_sync<kNumSMs, kGridSyncIndex>(workspace, sm_idx, thread_idx, sync_scope);

if (sm_idx == 0) {
// Send signals to remote ranks
if (thread_idx < kNumRanks)
ptx::red_add_rel_sys(sym_buffer.map(signal_ptr, thread_idx), ...);
// Wait arrival with 30s timeout
// ...
}
// Grid sync after NVLink completion
grid_sync<kNumSMs, kGridSyncIndex>(workspace, sm_idx, thread_idx, sync_scope);
}

1.3 Buffer 布局设计

在 csrc/apis/mega.hpp 文件中，get_symm_buffer_size_for_mega_moe 函数精心设计了对称内存中各区域的排列顺序：Workspace → Input(x, x_sf, topk_idx, topk_weights) → L1 pool(acts, sf, weights) → L2 pool(acts, sf) → Combine buffer(BF16)。所有这些 buffer 在地址空间中是连续排列的，并且 SF buffer 按照 UTCCP 128 元素的粒度进行了对齐。

二、FP4 Indexer：低精度 MQA Logits 计算

2.1 问题背景

DeepSeek v3.2 的 lightning indexer 采用 weighted ReLU MQA logits 来进行评分计算。此前，该方案仅支持 FP8 精度。本次发布新增了 FP4 精度支持（仅限 SM100/Blackwell 架构），并扩展了对更大规模 MTP（Multi-Token Prediction）的支持。

2.2 关键实现

从 csrc/apis/attention.hpp 文件可以看出，新推出的 fp8_fp4_mqa_logits 和 fp8_fp4_paged_mqa_logits API 统一了 FP8 与 FP4 两条处理路径。FP4 路径的核心差异体现在以下几个方面：

• Q 和 KV 使用 kPackedFP4 数据类型（两个 FP4 值被打包在一个字节中）。
• Scale factor 采用 torch::kInt32（UE8M0 编码）。
• head_dim 固定为 128，以确保 64B swizzle 对齐。
• TMA descriptor 使用 fp4_unpacked_smem = false，即对应 CU_TENSOR_MAP_DATA_TYPE_16U4_ALIGN8B。

static torch::Tensor fp8_fp4_mqa_logits(
const std::tuple<torch::Tensor, std::optional<torch::Tensor>>& q,
const std::tuple<torch::Tensor, torch::Tensor>& kv, ...) {
const auto [q_fp, q_sf] = q;
const bool is_fp4 = q_sf.has_value();
// FP4 path: head_dim *= 2 (packed), scalar type = kPackedFP4
// FP8 path: scalar type = kFloat8_e4m3fn
}

同时保留了旧的 `fp8_mqa_logits` API 作为兼容性封装，其内部会调用新接口，并将 `std::nullopt` 作为 SF 参数传入。

## 三、动态 Swap A/B：MoE GEMM 的关键性能优化

### 3.1 问题：SM100 上 MoE 场景的 Multicast 失效

在 MoE 的 m-grouped contiguous GEMM 中，矩阵 A（激活值）按 token 进行连续存储，而矩阵 B（权重）则按 expert 分组存放。SM100 上的 TMA multicast 虽然能在 cluster 内广播数据，但由于 A 的 M 维度（token 数量）通常较小且缺乏规律性，导致无法对 A 实施有效的 multicast 操作。

### 3.2 解决方案：在 SM100 上交换 A/B 操作数

本次重构在 `GemmConfig` 中新增了 `Layout::swap_ab` 字段。当 `swap_ab = 1` 时，实际计算转变为 `B^T @ A^T`，这使得原本的 N 维度（权重矩阵的 hidden size，通常较大且规整）变成了 UMMA 的 M 维度，从而能够实现 cluster multicast。

// Always enable swap A/B for m-grouped GEMMs
if (desc.gemm_type == GemmType::MGroupedContiguous || …) {
const bool swap_ab = true;
const auto block_n = 128;
const auto block_m = heuristics_runtime->get_mk_alignment_for_contiguous_layout();
const auto cluster_n = ceil_div(desc.n, block_n) % 2 == 0
and desc.num_sms % 2 == 0 ? 2 : 1;
// …
}

在 SM100 的候选枚举过程中，`swap_ab` 和非 `swap_ab` 两个方向都会被纳入考量（`for (int swap_ab = 0; swap_ab < 2; ++swap_ab)`），并由统一的比较器从中选出最优配置。这一优化对 MoE 场景的性能提升尤为显著。

## 四、全面的启发式重构

### 4.1 从硬编码到结构化配置

整个启发式系统经历了彻底的重构。旧代码中的 `GemmConfig` 是一个内容庞杂的结构体，而新代码将其拆分成了四个清晰的层次：

// csrc/jit_kernels/heuristics/config.hpp#L134-L140
struct GemmConfig {
Layout layout; // block_m/n/k, swap_ab, cluster_m/n
StorageConfig storage_config; // load/store block sizes, swizzle modes
PipelineConfig pipeline_config; // num_stages, smem_size
LaunchConfig launch_config; // num_sms, num_threads
};

### 4.2 SM90 引入 L1/L2 带宽建模

SM90 的配置选择已从“波数最小化”升级为基于 **L1/L2 带宽建模的 cycle 估计**：

// csrc/jit_kernels/heuristics/sm90.hpp#L159-L176
const int l2_bandwidth_per_cycle = std::min(64. * desc.num_sms, 8e6 / 1.3e3);
const int l1_bandwidth_per_cycle = 128 * desc.num_sms;
// …
int64_t num_l2_cycles = (num_bytes_l2_ab + num_bytes_l1_l2_cd) * num_blocks / l2_bandwidth_per_cycle;
int64_t num_l1_cycles = (num_bytes_l1_ab + num_bytes_l1_tc + …) * num_blocks / l1_bandwidth_per_cycle;
float wave_efficiency = static_cast(num_blocks) / (num_waves * desc.num_sms);
int64_t num_cycles = std::max(num_l1_cycles, num_l2_cycles) / wave_efficiency;

这一改进使得 SM90 在处理小矩阵场景时，能够更精准地选择 block size 和 multicast 策略。

## 五、其他重要改进

### 5.1 PDL（Programmatic Dependent Launch）

新增了 `set_pdl` / `get_pdl` API，在 kernel launch 时设置 `cudaLaunchAttributeProgrammaticStreamSerialization` 属性。这允许 CUDA 运行时在前一个 kernel 尚未完全结束时就开始启动下一个，从而有效减少 kernel 间的执行间隙：

// csrc/jit/handle.hpp#L92-L98
if (enable_pdl) {
auto& attr = attrs[config.numAttrs ++];
attr.id = cudaLaunchAttributeProgrammaticStreamSerialization;
attr.val.programmaticStreamSerializationAllowed = 1;
}
“`

5.2 JIT 编译加速与分布式文件系统修复

• 增量哈希：新增了 IncludeParser（位于 csrc/jit/include_parser.hpp），它仅对 #include <deep_gemm/*> 依赖树进行递归哈希计算，取代了原先对所有头文件内容的完整哈希操作，大幅降低了缓存 key 的计算开销。
• 原子目录重命名：编译产物会先写入一个临时目录，最后通过 std::filesystem::rename 执行原子替换，从根本上解决了分布式文件系统（如 NFS/Lustre）上多进程竞争导致的 JIT 缓存损坏问题。同时，在写入完成后会调用 fsync 确保数据持久化。
• 损坏检测：加载 CUBIN 文件时，会通过 cuLibraryGetKernelCount 验证 kernel 数量必须为 1，若不符合则提示用户清除缓存目录。

5.3 FP8 × FP4 GEMM

新增了 sm100_fp8_fp4_gemm_1d1d.hpp，支持 A 为 FP8、B 为 FP4（或反过来）的混合精度 GEMM 运算，这构成了 Mega MoE 的基础计算单元。

5.4 C++20 标准升级

项目已从 C++17 升级至 C++20（在 CMakeLists.txt 中通过 CMAKE_CXX_STANDARD 20 启用），以支持 std::variant、designated initializers 等现代语言特性。

总结

PR #304 标志着 DeepGEMM 迄今为止规模最大的公开版本发布。其核心价值体现在以下几个方面：

1. Mega MoE 是业界首个将 EP dispatch、linear1、SwiGLU、linear2 和 combine 这五个步骤融合为一个单一 kernel 的开源实现。它利用对称内存、三角色线程分工以及 NVLink barrier 技术，实现了通信与计算的深度重叠。
2. 动态 Swap A/B 有效解决了 SM100 上 MoE GEMM multicast 受限的痛点，虽然原理看似简单，但实际工程优化效果显著。
3. 启发式系统的结构化重构 为后续支持更多硬件架构和更精细的自动调优奠定了坚实基础。
4. JIT 编译的分布式文件系统修复 解决了大规模训练与推理集群中实际部署时面临的痛点。

正如 PR 描述中所说：“Mega MoE is still under development and optimizations, stay tuned”——这仅仅是一个开端。

性能数据将在后续陆续公布，而 DeepEP V2（提供 ElasticBuffer）也即将随之发布，届时 Mega MoE 的完整多进程测试才能顺利跑通。此次发布充分展示了 DeepSeek 团队在 GPU 微架构级优化上的深厚功底，同时也为开源社区提供了学习现代 LLM 推理内核设计的宝贵资源。

相关推荐

• 解锁 GPU GEMM 极致性能！PyTorch 集成 CuteDSL 后端，覆盖 BF16/FP8/FP4
• 如何用 Triton 写出比 SGLang 更快的 MoE Align Block Size
• 微秒级决策，近最优性能！AMD 提出 tritonBLAS：告别耗时参数搜索，用“数学”算出最快矩阵乘法