DeepGEMM 指南：FP8 Kernel、Mega MoE 以及 Hopper/Blackwell

如果你对 DeepGEMM 只能记住一件事，那就是：它是一个小巧的 CUDA 库，专注于优化几种精心挑选的 FP8 GEMM 和 MoE 算子形状，在 CuTe 和 CUTLASS 原语的基础上进行了手工调优，并通过 JIT 编译的 Python API 进行封装。它不是一个通用的 BLAS 库。它是 DeepSeek 在发布 V3、R1、V3.2 以及现在的 FP8 版 V4 时，梦寐以求的 kernel 层实现。

1. 一句话定义

DeepGEMM 是一个统一的、JIT-compiled CUDA tensor core 库，适用于 FP8/FP4 GEMM、BF16 GEMM、MoE grouped GEMM、用于 V3.2 lightning indexer 的 MQA scoring kernels、HyperConnection，以及一个将 NVLink 通信与 tensor core 计算重叠（overlap）的 fused Mega MoE kernel。 它针对 SM90 (Hopper) 和 SM100 (Blackwell) 架构，采用 MIT 许可协议，并为 DeepSeek 自有的 V3/R1 训练和推理提供动力。

把它想象成 DeepSeek's 推理栈的底层。 DeepEP 负责 expert-parallel 通信，模型代码负责路由 tokens，而 DeepGEMM 则是 tensor cores 实际执行计算消耗周期的地方。

2. 为什么它会存在

在 DeepGEMM 出现之前，训练巨型 FP8 MoE 的开源团队有三个糟糕的选择：使用 cuBLAS 并接受它提供的任何 scaling 和累加方式；为每种形状编写自定义 CUTLASS kernel 并承担编译时间的代价；或者退而求其次使用 BF16，从而耗尽 HBM 带宽和 tensor core 的机时。

DeepSeek's 的约束非常明确：在 FP8 下训练和提供超大规模 MoE 模型服务，具备细粒度 per-tile scaling 的精度，并使用与 Mixture-of-Experts 模型中专家实际调用方式相匹配的 grouped GEMM 形状（一个 M-axis，固定的 N 和 K）。CUTLASS 可以实现这一点，但它也会让你的构建过程淹没在 templates 中。

DeepGEMM's 的方案非常精简且有针对性。选择重要的形状。将每个 kernel 手写为可读的 CUDA。将 CUTLASS's 的 templates 隐藏在轻量级 runtime 之后。当 call site 出现时，JIT-compile 对应的 specialization。README 明确指出：该库是 “为简单而设计，仅包含有限数量的核心 kernel 函数，使其成为学习 NVIDIA GPU kernel 优化技术的简洁且易于获取的资源。”

3. 心智模型：核心组件

DeepGEMM 足够小，你可以完全掌握它的全貌。它由五个核心部分组成。

图表 (ASCII):

Python call site (PyTorch tensors, FP8 + scale factors)
        |
        v
+----------------------------+
|   deep_gemm.fp8_gemm_nt   |   <- thin Python wrapper
|   m_grouped_fp8_gemm_*    |
|   fp8_paged_mqa_logits    |
|   fp8_fp4_mega_moe        |
+----------------------------+
        |
        | shape + dtype + arch -> JIT key
        v
+----------------------------+
|  JIT module (CPP / NVRTC) |   <- compiles + caches kernel
|  $HOME/.deep_gemm cache   |
+----------------------------+
        |
        v
+----------------------------+
|  CUDA kernel              |   <- WGMMA / TMA / cluster scheduling
|  ~few hundred lines core  |       Hopper SM90 or Blackwell SM100
|  CuTe / CUTLASS primitives|
+----------------------------+
        |
        v
   tensor cores: FP8 / FP4 mma -> FP32 acc -> bf16/fp32 out

一个有趣的设计选择是边界。DeepGEMM 完成了你对 GEMM 算子的所有预期。而其周边的操作 — 转置、FP8 类型转换、缩放因子准备 — README 明确告知你需要自行处理或融合到前一个算子中。

4. 最小端到端示例

获取代码仓库，构建 JIT 模块，并运行单个 FP8 GEMM。两条命令完成安装，第三条是 Python 冒烟测试，用于端到端测试 JIT 路径。

# 1. Clone with submodules (CUTLASS, fmt)
git clone --recursive [email protected]:deepseek-ai/DeepGEMM.git
cd DeepGEMM

# 2. Build the CPP JIT module
./develop.sh

# 3. Run the FP8 dense test
python tests/test_fp8.py

一旦包导入成功，调用处就是一个普通的 PyTorch op。注意布局约定：fp8_gemm_nt 表示 A 不转置，B 转置（因此它计算 D = C + A @ B.T）。

import torch
import deep_gemm

# Allocate FP8 (E4M3) tensors with separate FP32 scale factors
m, n, k = 4096, 4096, 4096
a = torch.randn(m, k, device='cuda').to(torch.float8_e4m3fn)
b = torch.randn(n, k, device='cuda').to(torch.float8_e4m3fn)
c = torch.zeros(m, n, device='cuda', dtype=torch.bfloat16)

# Scale factors: FP32 for SM90; SM100 wants packed UE8M0 in torch.int
a_sf = torch.ones(m, k // 128, device='cuda')   # per-tile scaling
b_sf = torch.ones(n, k // 128, device='cuda')

# Wrappers expect tuples of (tensor, scale_factor)
deep_gemm.fp8_gemm_nt((a, a_sf), (b, b_sf), c)

# c now holds D = C + A @ B.T in BF16

这就是一个基础稠密 GEMM 的全部接口。第一次调用会触发 JIT 编译，并将算子缓存至 $HOME/.deep_gemm。从第二次调用开始，就是普通的 CUDA 启动。

5. 深度解析：FP8 与细粒度缩放

FP8 在 Hopper 架构上有两种格式：E4M3（精度更高，范围更小）和 E5M2（范围更大，精度更低）。DeepGEMM's 的 GEMM 使用 E4M3 作为输入，并在 FP32 中进行累加，然后在存储时转换回 BF16 或 FP32。这是标准的 FP8 方案。有趣的地方在于缩放。

全张量缩放（每个矩阵一个 FP32）成本很低，但当激活值存在离群点时会损失精度。细粒度缩放为每个 tile 分配一个缩放因子 — 例如，每 128 列 K 分配一个缩放因子。这就是 k // 128 在上述示例中的来源。在 WGMMA 累加之前，每个 tile 都会进行独立缩放。

DeepGEMM 在内核中内置了两个特定于架构的特性：

• SM90 需要 FP32 缩放因子，且 LHS 缩放因子采用 TMA 对齐的转置布局。
• SM100 需要打包的 UE8M0 缩放因子：将四个 UE8M0 浮点数打包进单个 torch.int，同样需要 TMA 对齐。

UE8M0 是一种 8 位无符号仅指数浮点数（无符号位，无尾数）。它是 CUDA 12.9+ 用于 Blackwell's 原生块缩放路径的格式。DeepGEMM 暴露了get_mn_major_tma_aligned_packed_ue8m0_tensor 作为帮助函数，引导你进入正确的布局。

# Pack FP32 scale factors into the SM100 layout
import deep_gemm
sf_packed = deep_gemm.get_mn_major_tma_aligned_packed_ue8m0_tensor(a_sf)
# sf_packed is a torch.int tensor with 4 UE8M0 values per element

还有一个值得关注的待解决 bug。Issue #309 / PR “fix: correct operator precedence in pack_ue8m0_to_int assertion” 指出，在 deep_gemm/utils/math.py was always passing because == binds tighter than & in Python. The fix is one line of parens. The lesson is the more interesting part: even kernel libraries have boring Python bugs hiding under the FP8 math.

6. Deep Dive: The JIT Compile Path

DeepGEMM doesn't precompile every (M, N, K, dtype, layout) combination. It compiles on first call, caches by JIT key, and reuses. The 2025.07.20 refactor (PR #112) replaced the older path with a low-CPU-overhead C++ JIT module.

Two compiler backends are wired in:

• NVCC (default) — slower compile, reliably best perf. NVCC 12.9+ does FFMA interleaving automatically, so DeepGEMM dropped its old post-compilation SASS pass.
• NVRTC — up to 10× faster compile (per PR #94), but may have lower performance on some shapes. Opt in with DG_JIT_USE_NVRTC=1.

The cache lives at $HOME/.deep_gemm 默认。使用 DG_JIT_CACHE_DIR。 DG_JIT_PRINT_COMPILER_COMMAND=1。导出 PTX/SASS 以供检查 DG_JIT_DUMP_SASS=1。确认没有发生本地内存溢出，使用 DG_JIT_PTXAS_CHECK=1— 溢出会严重破坏 FP8 性能。

# Trace what the JIT is doing on a fresh kernel
DG_JIT_DEBUG=1 \
DG_PRINT_CONFIGS=1 \
DG_JIT_PRINT_COMPILER_COMMAND=1 \
DG_JIT_PTXAS_CHECK=1 \
python -c "import deep_gemm; deep_gemm.fp8_gemm_nt(a_pair, b_pair, c)"

注意多进程竞态。Issue #301 记录了当多个 rank 尝试在共享文件系统上并发编译并写入同一个 kernel 时，会出现 JIT 缓存竞态条件。目前有一个正在处理的 PR 修复此问题。在合并之前，安全的模式是从 rank 0 预热缓存，执行 barrier，然后 fan out。

7. 深度解析：分组 GEMM（MoE 的故事）

混合专家模型（Mixture-of-Experts）将每个 token 路由到专家的子集。在 DeepSeek-V3-Pro 中，大约有 384 个专家，top-k=6。路由之后，你会得到一个不规则的 batch：专家 0 得到了一些 token，专家 17 得到了很多 token，专家 200 则一个也没得到。每个专家应用自己的权重矩阵。

朴素方案：每个专家执行一个 GEMM。这对 Tensor Core 来说是灾难性的。更好的方案：分组 GEMM。DeepGEMM 对 CUTLASS 分组 GEMM 的改进在于针对 MoE 进行了优化。

• M 轴分组，连续（training / prefill）：沿 M 轴拼接所有专家的 token。N 和 K 保持固定（因为专家共享形状）。每个专家分段必须对齐到 get_mk_alignment_for_contiguous_layout()。
• M 轴分组，掩码（decode）: 在 CUDA graphs 下，CPU 无法获知每个专家（expert）分配到了多少个 token。通过传递一个 mask tensor；kernel 仅计算有效块。该设计旨在处理以下组件的输出： DeepEP's 的低延迟调度（dispatch）。
• K-分组，连续（MoE 反向传播）: M 和 N 固定，K 分组。用于权重梯度传递（PR #95）。

# Grouped FP8 GEMM in masked-decode mode
import deep_gemm

# x_fp8 shape: [num_tokens_max, hidden]; mask shape: [num_tokens_max]
deep_gemm.m_grouped_fp8_gemm_nt_masked(
    (x_fp8, x_sf),
    (w_fp8, w_sf),
    out=y,
    expert_offsets=expert_offsets,   # int32 prefix sum from DeepEP
    mask=mask,                       # which tokens are real this step
)

掩码路径（masked path）允许你在专家分配随步数变化时，依然在推理解码（inference decode）过程中保持 CUDA graphs 开启。仅此一点就极具价值：否则，每个解码步骤都必须承担重新捕获 graph 的开销。

8. 深度解析：Mega MoE 与 Lightning Indexer

2026 年 4 月发布的版本（PR #304）标志着 DeepGEMM 不再仅仅是一个 kernel 库，而开始演变为一个运行时（runtime）。其中最重要的两个部分是： Mega MoE 以及 FP4 indexer。

Mega MoE 将整个 MoE 块 — 包括专家并行调度（expert-parallel dispatch）、FP8×FP4 linear 1、SwiGLU 激活、FP8×FP4 linear 2 以及 EP combine — 融合进一个单一的 mega-kernel 中。它通过将调度和结合（combine）视为同一调度任务的一部分，在对称内存缓冲区（symmetric memory buffers）上运行，从而实现了 NVLink 通信与 tensor core 计算的重叠。symm-mem 路径要求 PyTorch ≥ 2.9。初始版本仅支持 FP8×FP4 MoE。

# Mega MoE skeleton (multi-process, symmetric memory)
buffer = deep_gemm.get_symm_buffer_for_mega_moe(
    group, num_experts, num_max_tokens_per_rank, num_topk, hidden, intermediate_hidden
)
transformed_l1, transformed_l2 = deep_gemm.transform_weights_for_mega_moe(l1, l2)

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)

y = torch.empty((num_tokens, hidden), dtype=torch.bfloat16, device='cuda')
deep_gemm.fp8_fp4_mega_moe(y, transformed_l1, transformed_l2, buffer)

在索引器端，V3.2 引入了一个 “lightning indexer”，它在运行完整注意力机制之前，通过加权 ReLU MQA logits 对 token 进行评分。DeepGEMM 发布了 fp8_mqa_logits (prefill, non-paged) 以及fp8_paged_mqa_logits (decode, paged)，带有 cu_seq_len_k_start/end 每个 query token 的范围。四月发布的版本增加了一个 FP4 indexer 路径，支持更大的 MTP，因此 KV cache 本身可以是 FP4 格式，且节省的 KV 字节数呈线性增长。此外还有一个融合的kv_cache + kv_cache_sf 布局（交错值 + 每行缩放因子），该布局引入了一个 stride bug，已在公开 issue 中指出并提交了修复 PR。

维护团队在 PR #316中发布了 EP8 下 V4-Flash 和 V4-Pro 的 Mega MoE 数据。我们不打算逐行重复这些数据，因为它们与特定的 commit 绑定，并会在下个版本发生变化。结果的规律是一致的：小 batch 受带宽限制（bandwidth-bound），Mega MoE 的优势在于延迟（latency）；大 batch 受计算限制（compute-bound），优势在于 TFLOPS。

9. 我们做错的地方

当我们第一次将 DeepGEMM 接入一个侧边项目的 FP8 训练器时，遇到了三个坑。

第一：我们假设 FP8 类型转换（casting）是 DeepGEMM 的工作。 事实并非如此。README 明确指出：“诸如输入转置或 FP8 类型转换之类的操作必须由用户单独处理，请独立实现它们或将其融合到之前的 kernel 中。” 我们当时有一个 .to(torch.float8_e4m3fn) 在我们的前向传播中，这抵消了 FP8 GEMM 带来的所有收益。修复方法是将类型转换（cast）融合到前一个 LayerNorm 中。在下一次运行中，我们的吞吐量大幅提升，这并非因为 DeepGEMM，而是因为我们不再误读 DeepGEMM 所负责的范围。

第二：我们盲目地在不同 rank 之间共享 JIT 缓存。 分布式启动触发了 Issue #301's 的竞态问题。两个 rank 竞相编译同一个 kernel；其中一个写入了部分缓存文件，而另一个读取了它。症状看起来像是 WGMMA 深层随机出现的 IMA 错误。有效的模式是：rank 0 预编译，所有 rank 在缓存目录上设置同步屏障（barrier），然后恢复运行。

第三：我们搞混了 SM100 的缩放因子（scale-factor）打包方式。 将 FP32 缩放因子推送到一个预期 packed UE8M0 的 kernel 中，会静默产生错误的输出。kernel 并未崩溃；loss 曲线看起来只是偏离了 0.5%。教训是：在 Blackwell 上，务必通过 get_mn_major_tma_aligned_packed_ue8m0_tensor。这个打包辅助函数（pack helper）的存在是有原因的。

10. 真实世界的工作流

MoE 推理引擎实际上是如何接入 DeepGEMM 的：

11. 常见错误（源自 Issues 和 Reddit）

• 尝试在 SM120（RTX 5090, RTX Pro 6000 Blackwell）上运行。 Issue #317 以及开启的 SM120 PR表明 tf32_hc_prenorm_gemm 并且 paged MQA logits kernel 在 SM120 上会断言 “Unsupported architecture”。 根本原因：在以下路径中没有 SM120 的实现： impls/；仅存在 SM90 和 SM100 路径。权宜之计是通过社区 PR 逐个 kernel 添加 SM120 特化实现。
• SM100 上的缩放因子（scaling）不匹配。 当 SM100 需要打包的 UE8M0 时，传入 FP32 缩放因子会静默产生错误结果。 根本原因：kernel 为每个torch.int读取四个 UE8M0 值；FP32 张量会被解释为 32 位的垃圾数据。请务必通过 helper 处理。
• 分布式任务中的 JIT 缓存竞态。 症状：随机的 IMA、kernel 缺失、共享对象文件不完整。 根本原因：在 $HOME/.deep_gemm 所在的网络文件系统上存在并发写入（Issue #301）。请从单个 rank 进行预热。
• fused KV cache 中的步长（Stride）错误。 该 fp8_fp4_paged_mqa_logits 路径分解了一个融合的 [values | SF] 使用扁平布局步长（flat-layout strides）而非交错布局的 KV cache。根本原因：混淆了两种有效的布局；测试辅助函数使用了错误的一种，从而掩盖了小块（small blocks）下的 bug。
• 将 fp8_gemm_nt 视为通用的 matmul。 N/T 表示非转置/转置；SM90 仅支持 NT，SM100 支持全部四种。 根本原因：从 torch matmul 复制粘贴而来，而该 matmul 并没有此约束。

12. 性能、扩展性与成本说明

DeepSeek's 的核心宣传点是 “在 Hopper 上高达 1350+ FP8 TFLOPS”。这个数字是真实的，维护者在发布推文中也提到了这一点，它也是在 H800 上针对特定 commit 和特定 shape 测得的峰值。请将其作为参考基准，而非性能保证。

实际运行中真正重要的指标：

• FP8 与 BF16 的吞吐量对比：当 kernel 融合良好时，在 Hopper tensor cores 上大约有 2× 的提升。DeepGEMM 旨在将常数保持在接近 1 的水平，而不是像粗糙的 FP8 路径那样回落到约 1.4×。
• HBM 压力：FP8 激活使内存流量减半。对于长上下文解码（受带宽限制的场景），这是更大的优势。
• JIT 编译开销: 首次调用时，NVCC 编译是耗时最长的部分。NVRTC 能显著降低这一开销。在缓存命中的情况下，两者的开销几乎可以忽略不计。
• Mega MoE: PR #316 中发布的单批次数据随 batch size 增加而上升，并在接近单卡 Tensor Core 饱和点时趋于平缓。小批次主要受 fused dispatch 延迟的影响。
• 成本考量: DeepGEMM 本身是免费的 (MIT)。其成本主要在于将其正确集成到运行时环境中的工程时间，以及维护与上游兼容的 scale-factor 流水线。

该库本身并不会直接带来功耗优势。它提升的是 Tensor Core 利用率和内存带宽，这才是大规模推理服务中影响账单的关键因素。

13. 适用人群（以及不适用人群）

在以下情况下使用 DeepGEMM：

• 你在 H100/H200/H800/B200/GB200 上部署或训练 DeepSeek 风格的 FP8 MoE。
• 你需要用于长上下文解码的 V3.2 lightning indexer kernels。
• 你正在构建研究代码库，并希望有一个精简、可读的 Hopper FP8 kernel 设计参考。
• 你已经在使用 DeepEP，并希望获得 Mega MoE 的 overlapped dispatch/combine 特性。

在以下情况下请跳过：

• 你需要通用的 BLAS。请使用 cuBLAS 或完整的 CUTLASS。
• 你只关注 BF16/FP16 稠密 LLM。FP8 布局处理对你来说是多余的开销。
• 你目前在 SM120 工作站级 Blackwell 上运行。请等待 SM120 PR 合并，或使用 CUTLASS 路径。
• 你处于迭代早期，需要支持各种 shape 且能 Just Work。DeepGEMM 刻意保持了较窄的适用范围。

14. 社区动态

DeepSeek 团队在 2025 年 2 月的 “开源周” 期间发布了 DeepGEMM，其原始推文对该库定位的描述比任何第三方总结都更加精准。

Reddit's r/LocalLLaMA 帖子关于最初发布的讨论是了解社区脉搏的最佳途径。主流反应是技术上的尊重：一个拥有严谨 Tensor Core 调度的小型库，在 DeepSeek 关注的矩阵形状（shapes）上，性能超越了内部调优的 CUTLASS。几位评论者指出，“约 300 行”的说法是指核心 kernel 逻辑，而非整个仓库，且称其为“教程”有些言过其实。这两点其实并不矛盾。

Hacker News 的 讨论挖掘得更深。高赞评论集中在 SASS 级别的 FFMA 交错（自 NVCC 12.9 自动调优后已弃用）、选择公开而非隐藏 UE8M0 打包（packing），以及一个反复出现的反对观点：DeepGEMM 令人印象深刻，但它也 局限。几位工程师指出，“优于专家调优的 kernel”主要适用于 DeepSeek-V3 使用的特定形状分布，一旦进行相同的手动调优，CUTLASS 在更奇特的形状上仍然胜出。

反对的声音很重要。最有价值的质疑解读是：DeepGEMM 最好被理解为针对 DeepSeek 负载的生产级 FP8 参考实现。它不是通用的 BLAS，也不是 CUTLASS 的终结。团队自己在 README 中向 CUTLASS 致谢时也承认了这一点。

15. 结论：值得使用吗？

DeepGEMM 赢得了“源码即文档”的勋章。这些算子足够精简，一下午就能读完；同时又足够硬核，每次阅读都能让你有所收获。单凭这一点就非常罕见。

16. 宏观视角

DeepGEMM 是 DeepSeek 一直在公开构建的“FP8 与 MoE 优先”技术栈中的一环。 DeepEP 负责通信。 DeepGEMM 负责计算。上层的模型代码同时依赖这两者。这种分离代表了 MoE 推理的未来：算子层需要感知缩放（scaling），通信层需要感知专家分配（expert assignment），而运行时（runtime）则将它们缝合在一起。

本指南背后更深层的背景是：FP8 已成为前沿 MoE 模型激活精度的首选，而 FP4 也正朝着 KV 缓存和权重的方向演进。DeepGEMM 在 2026 年 4 月新增的 FP8×FP4 GEMM 和 FP4 索引器，是 Blackwell 时代推理服务发展方向的风向标。如果你想了解 DeepSeek 本身的发展脉络，我们的 DeepSeek-V3.2 指南涵盖了将极速索引器算子引入此仓库的模型架构。若想从更高维度了解研究型工程代码库如何随 LLM 一同演进，我们的 Karpathy LLM 知识库文章是一个很好的补充。如果你想通过现有的客户端在本地运行 DeepSeek 模型，这份 免费的 Claude Code 指南解释了如何将 Claude Code 指向任何开源权重提供商，包括 DeepSeek。

17. 常见问题解答

18. 术语表

19. 所有来源 & 链接

来源归属

主要来源

• deepseek-ai/DeepGEMM (GitHub)
• DeepGEMM README
• DeepGEMM 合并请求 (完整动态)
• DeepGEMM Issues
• deepseek-ai/DeepEP (通信侧配套组件)

社区

• X: DeepSeek 发布 DeepGEMM 的公告
• Reddit: r/LocalLLaMA 发布讨论帖
• Hacker News 讨论帖

文档

• NVIDIA CUTLASS
• CuTe (CUTLASS 布局/分块代数)
• PTX 替代浮点格式 (UE8M0 等)

内部链接

• DeepSeek-V3.2 指南
• 免费 Claude Code：支持在任何 LLM 提供商上运行
• Karpathy 的 LLM 知识库