编程 DeepGEMM 深度解析：DeepSeek 开源的 FP8 GEMM 内核如何重塑 AI 推理性能边界

2026-04-21 05:16:09 +0800 CST views 9

DeepGEMM 深度解析：DeepSeek 开源的 FP8 GEMM 内核如何重塑 AI 推理性能边界

引言：当 1550 TFLOPS 成为新的起点

2025 年 4 月，DeepSeek 团队开源了 DeepGEMM——一个专为现代大语言模型设计的高性能 FP8 GEMM（General Matrix Multiply）内核库。这个看似技术细节的开源项目，实则是 AI 基础设施领域的一次重要突破。

在 H800 GPU 上，DeepGEMM 实现了高达 1550 TFLOPS 的计算性能，这一数字意味着什么？它接近 NVIDIA H800 Tensor Core 的理论峰值算力，代表着 FP8 矩阵运算在实际工作负载中所能达到的效率极限。

但 DeepGEMM 的意义远不止于性能数字。它代表了深度学习计算范式的三个关键转变：

从 FP16/BF16 到 FP8 的精度迁移——在保持模型质量的同时，将计算吞吐提升 2 倍
从静态编译到 JIT（Just-In-Time）运行时编译——突破传统 CUDA 库的安装和部署限制
从通用库到领域专用优化——针对 MoE（Mixture of Experts）架构的深度定制

本文将从底层硬件架构出发，深入剖析 DeepGEMM 的设计哲学、核心算法与工程实现，并探讨它对 AI 推理基础设施的深远影响。

第一章：FP8——AI 计算的精度革命

1.1 为什么需要 FP8？

在深度学习的发展历程中，模型精度的选择经历了一场持续的"降维"革命：

精度格式	位宽	动态范围	典型应用场景
FP32	32-bit	~83 dB	训练主路径、梯度累积
FP16	16-bit	~12 dB	混合精度训练、推理
BF16	16-bit	~78 dB	大模型训练（Google 主推）
FP8 (E4M3)	8-bit	~18 dB	推理、部分训练场景
FP8 (E5M2)	8-bit	~48 dB	梯度计算、对范围敏感场景

FP8 的出现并非偶然。随着 Transformer 模型规模指数级增长，计算和内存带宽成为核心瓶颈。NVIDIA Hopper 架构（SM90）首次原生支持 FP8 Tensor Core 运算，而 Blackwell 架构（SM100）进一步扩展了 FP8/FP4 的支持能力。

FP8 的核心优势体现在三个维度：

计算吞吐翻倍：在相同的 Tensor Core 单元上，FP8 运算的吞吐量是 FP16 的 2 倍。这意味着在内存带宽充足的情况下，理论性能可以直接翻倍。

内存占用减半：模型权重和激活值以 FP8 存储，显存占用减少 50%，这使得更大 batch size 的推理成为可能，直接提升吞吐量。

通信开销降低：在分布式训练中，梯度以 FP8 传输可以显著减少 NVLink/InfiniBand 的通信压力。

1.2 FP8 的精度挑战与解决方案

然而，FP8 的 8-bit 位宽带来了严峻的精度挑战。E4M3 格式只有 4 位指数和 3 位尾数，动态范围仅约 18 dB，远低于 FP16 的 12 位指数。

DeepGEMM 采用了 细粒度缩放（Fine-grained Scaling） 技术来解决这一问题：

# FP8 GEMM 的基本计算模式
# D = C + A @ B
# 其中 A, B 是 FP8 张量，需要配合缩放因子

import torch
import deep_gemm

# A: [M, K] FP8 (E4M3)
# B: [N, K] FP8 (E4M3)，注意是转置布局
# C: [M, N] BF16/FP16
# scale_A: [M] FP32 或 packed UE8M0
# scale_B: [N] FP32 或 packed UE8M0

D = deep_gemm.fp8_gemm_nt(
    A_fp8,      # LHS operand
    scale_A,    # LHS scaling factor
    B_fp8,      # RHS operand (transposed)
    scale_B,    # RHS scaling factor
    C_bf16      # accumulator
)

细粒度缩放的核心思想是：对矩阵的每一行（或每个块）维护独立的缩放因子，在计算前进行反量化。这种按行/按块的缩放策略，相比全局缩放能更好地保留数值精度。

对于 Hopper 架构（SM90），DeepGEMM 要求缩放因子以 FP32 格式存储；而对于 Blackwell 架构（SM100），则支持更紧凑的 packed UE8M0 格式——将 4 个 UE8M0 值打包到一个 32-bit 整数中，进一步降低内存开销。

第二章：DeepGEMM 架构设计——极简主义的胜利

2.1 设计哲学：做减法而非加法

DeepGEMM 的设计深受 CUTLASS 和 CuTe 的影响，但采取了截然不同的路径：

"The library is designed for simplicity, with only a limited number of core kernel functions, making it a clean and accessible resource for learning NVIDIA GPU kernel optimization techniques."

传统的 CUDA 矩阵运算库（如 CUTLASS）采用高度模板化的设计，通过 C++ 模板元编程实现各种配置的组合爆炸。这种设计虽然灵活，但带来了编译时间长、代码复杂度高、学习曲线陡峭等问题。

DeepGEMM 的核心设计决策包括：

运行时 JIT 编译：所有内核在运行时通过轻量级 JIT 模块编译，安装时无需 CUDA 编译
有限的核心函数：只提供少量精心优化的核心 kernel，而非覆盖所有可能的配置
显式控制流：避免复杂的模板代数，使用直接的 CUDA C++ 编写核心逻辑
SM90/SM100 双架构支持：同时支持 Hopper 和 Blackwell 架构，但保持代码简洁

2.2 JIT 编译架构

DeepGEMM 的 JIT 编译系统是其最独特的设计之一。传统 CUDA 库需要在安装时编译大量 kernel variant，导致：

安装时间长（可能需要数十分钟）
二进制体积庞大
无法针对具体工作负载优化

DeepGEMM 的解决方案是：

// JIT 编译流程
// 1. 根据运行时参数（M, N, K, 数据类型等）生成专门的 kernel 代码
// 2. 使用 NVRTC 或 CPP JIT 模块进行即时编译
// 3. 缓存编译结果供后续调用

// 启用 NVRTC 加速编译（可选）
// export DG_JIT_USE_NVRTC=1
// 可获得最高 10 倍的编译速度提升

这种设计的优势在于：

快速安装：无需在安装时编译，pip install 后即可使用
针对性优化：根据实际的矩阵形状生成最优 kernel
低 CPU 开销：通过精心设计的 CPP JIT 模块，将编译开销降到最低

2.3 核心 Kernel 家族

DeepGEMM 提供了以下核心 kernel 家族：

2.3.1 基础 FP8 GEMM

# 非分组 FP8 GEMM
# D = C + A @ B.T
deep_gemm.fp8_gemm_nt(A, scale_A, B, scale_B, C)

# 支持多种内存布局
deep_gemm.fp8_gemm_nn(A, scale_A, B, scale_B, C)  # A @ B
deep_gemm.fp8_gemm_tn(A, scale_A, B, scale_B, C)  # A.T @ B
deep_gemm.fp8_gemm_tt(A, scale_A, B, scale_B, C)  # A.T @ B.T

2.3.2 分组 GEMM（Grouped GEMM）

针对 MoE 模型的特殊需求，DeepGEMM 提供了 M 轴分组 GEMM：

# 连续布局的分组 GEMM
# 适用于训练前向或推理 prefilling 阶段
# 不同 expert 处理的 token 数不同，但共享相同的 N, K
deep_gemm.m_grouped_fp8_gemm_nt_contiguous(
    A,          # 拼接后的输入 [total_M, K]
    scale_A,    # 缩放因子
    B,          # 权重 [num_experts, N, K]
    scale_B,    # 权重缩放因子
    C,          # 输出
    group_sizes # 每个 expert 的 M 维度大小
)

2.3.3 掩码分组 GEMM（Masked Grouped GEMM）

针对推理 decoding 阶段，当 CUDA Graph 启用且 CPU 无法预知每个 expert 接收的 token 数时：

# 使用掩码 tensor 指定有效计算区域
# 可与 DeepEP 的低延迟 kernel 输出配合使用
deep_gemm.m_grouped_fp8_gemm_nt_masked(
    A, scale_A, B, scale_B, C,
    mask_tensor  # 指定有效 token 位置
)

2.3.4 MQA Scoring Kernel

为 DeepSeek V3.2 的 lightning indexer 设计的专用 kernel：

# 计算 weighted ReLU MQA logits
# 用于高效的注意力分数计算
deep_gemm.fp8_mqa_logits(
    q,                  # query: [seq_len, num_heads, head_dim]
    kv,                 # key/value: [seq_len_kv, head_dim]
    kv_scale,           # KV 缩放因子
    weights,            # 权重: [seq_len, num_heads]
    cu_seq_len_k_start, # 序列起始位置
    cu_seq_len_k_end,   # 序列结束位置
    clean_logits        # 是否清理未填充的 logits
)

2.3.5 Mega MoE——融合核的极致

DeepGEMM 的巅峰之作是 Mega MoE kernel，它将以下操作融合为单个 mega-kernel：

EP（Expert Parallelism）dispatch
Linear 1（FP8xFP4）
SwiGLU 激活
Linear 2（FP8xFP4）
EP combine

# 分配对称内存缓冲区（需要 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 kernel
y = torch.empty((num_tokens, hidden), dtype=torch.bfloat16, device='cuda')
deep_gemm.fp8_fp4_mega_moe(y, transformed_l1, transformed_l2, buffer)

Mega MoE 的关键优化在于：

NVLink 通信与 Tensor Core 计算重叠：通过精细的流水线调度，隐藏通信延迟
FP8xFP4 混合精度：权重以 FP4 存储，进一步减少内存占用
对称内存（Symmetric Memory）：多进程共享内存，避免数据拷贝

第三章：性能优化技术深度剖析

3.1 Tensor Core 利用率最大化

NVIDIA Tensor Core 是执行矩阵乘法的专用单元。DeepGEMM 通过以下策略最大化 Tensor Core 利用率：

3.1.1 Warp-level 并行

// 在 Hopper 架构上，使用 wgmma（Warp Group Matrix Multiply Accumulate）指令
// 每个 warp group 可以独立执行矩阵乘法

// 典型的 warp 分组策略：
// - 将 M 维度划分为多个 tile，每个 tile 由一个 warp group 处理
// - N 维度同样划分，形成 2D 并行网格
// - K 维度通过流水线迭代处理

3.1.2 共享内存优化

// 使用共享内存进行数据复用
// - 将 A 和 B 矩阵的 tile 加载到共享内存
// - 多个 warp 复用同一份数据
// - 通过异步拷贝（cp.async）隐藏内存延迟

// DeepGEMM 的共享内存布局经过精心设计
// 以最大化 bank conflict-free 访问

3.1.3 寄存器压力管理

FP8 GEMM kernel 需要在寄存器中保存中间结果。DeepGEMM 通过精细的寄存器分配策略，避免寄存器溢出（spill）到本地内存。

3.2 内存布局与 TMA

Hopper 架构引入了 Tensor Memory Accelerator（TMA），可以硬件级地处理张量数据的加载和存储。

DeepGEMM 对内存布局的要求：

# LHS (A) 缩放因子布局要求：
# - TMA 对齐
# - 转置布局

# 对于 SM90 (Hopper):
# - 缩放因子使用 FP32 格式

# 对于 SM100 (Blackwell):
# - 缩放因子使用 packed UE8M0 格式
# - 4 个 UE8M0 打包到一个 torch.int

3.3 流水线与指令级并行

DeepGEMM 采用软件流水线技术，将数据加载、计算、存储重叠执行：

// 典型的双缓冲流水线：
// - Buffer 0 进行计算时，Buffer 1 加载下一阶段数据
// - 阶段完成后交换 buffer
// - 这种设计隐藏了内存访问延迟

// 在 Hopper 上，结合 cp.async 和 wgmma 指令
// 可以实现更深的流水线层级

3.4 PDL（Programmatic Dependent Launch）

DeepGEMM 支持 NVIDIA 的 PDL 技术，允许 kernel 之间建立依赖关系，实现更精细的执行控制：

# 启用/禁用 PDL
deep_gemm.set_pdl(True)

# PDL 可以在某些场景下减少 kernel launch 开销
# 并提高多 kernel 流水线的效率

第四章：与 DeepSeek 生态的深度集成

4.1 DeepGEMM 在 DeepSeek V3 中的角色

DeepGEMM 并非孤立存在，它是 DeepSeek 开源生态的关键组件：

DeepSeek 开源生态
├── DeepSeek-V3/V3.2: 大语言模型
├── DeepSeek-R1: 推理模型
├── DeepEP: 专家并行通信库
├── DeepGEMM: FP8 GEMM 内核库 ← 本文主角
└── ...

在 DeepSeek V3 的训练和推理中，DeepGEMM 负责：

前向传播：FP8 矩阵乘法，包括 dense 和 MoE 层
反向传播：权重梯度计算（K-axis grouped GEMM）
推理优化：Masked grouped GEMM 支持动态 batching
注意力计算：MQA scoring kernel 加速 indexer

4.2 与 DeepEP 的协同

DeepEP 是 DeepSeek 开源的专家并行通信库，DeepGEMM 与其紧密配合：

# 典型的工作流程：
# 1. DeepEP 执行 all-to-all 通信，分发 token 到各 expert
# 2. DeepEP 输出 token 分配信息（可作为 mask）
# 3. DeepGEMM 使用 masked grouped GEMM 执行 expert 计算
# 4. DeepEP 执行 all-to-all 聚合结果

# 在 Mega MoE 场景中，DeepEP 的通信与 DeepGEMM 的计算完全重叠

4.3 FP4 与混合精度

DeepGEMM 的 Mega MoE kernel 支持 FP8xFP4 混合精度，这是目前最激进的量化策略之一：

# FP4 格式：1 位符号 + 2 位指数 + 1 位尾数
# 动态范围仅约 6 dB，但存储密度是 FP8 的 2 倍

# 在 Mega MoE 中：
# - 激活值使用 FP8 (E4M3)
# - 权重使用 FP4
# - 通过精细的缩放因子保持精度

这种混合精度策略使得在相同显存容量下，可以部署更大的模型或支持更大的 batch size。

第五章：实战——使用 DeepGEMM 优化你的推理系统

5.1 环境准备

# 硬件要求：
# - NVIDIA SM90 (Hopper, H100/H800) 或 SM100 (Blackwell)
# - 推荐使用 CUDA 12.9+ 以获得最佳性能

# 安装依赖：
# - Python 3.8+
# - PyTorch 2.1+
# - C++20 编译器

# 克隆并安装 DeepGEMM
git clone --recursive git@github.com:deepseek-ai/DeepGEMM.git
cd DeepGEMM

# 查看安装脚本
cat develop.sh  # 开发环境设置
cat install.sh  # 安装脚本

# 执行安装
./develop.sh
./install.sh

5.2 基础使用示例

import torch
import deep_gemm

# 设置设备
assert torch.cuda.is_available()
device = torch.device('cuda')

# 创建测试数据
M, N, K = 1024, 2048, 4096

# FP8 张量（E4M3 格式）
A_fp8 = torch.randn(M, K, dtype=torch.float8_e4m3fn, device=device)
B_fp8 = torch.randn(N, K, dtype=torch.float8_e4m3fn, device=device)

# 缩放因子（FP32）
scale_A = torch.ones(M, dtype=torch.float32, device=device)
scale_B = torch.ones(N, dtype=torch.float32, device=device)

# 累加器（BF16）
C_bf16 = torch.zeros(M, N, dtype=torch.bfloat16, device=device)

# 执行 FP8 GEMM
D = deep_gemm.fp8_gemm_nt(A_fp8, scale_A, B_fp8, scale_B, C_bf16)

print(f"Output shape: {D.shape}")
print(f"Output dtype: {D.dtype}")

5.3 MoE 模型推理优化

import torch
import deep_gemm

# MoE 模型配置
num_experts = 8
top_k = 2
hidden_dim = 4096
intermediate_dim = 11008
num_tokens = 1024

# 模拟 MoE 前向传播
def moe_forward(x, expert_weights, topk_indices, topk_weights):
    """
    x: [num_tokens, hidden_dim] BF16
    expert_weights: [num_experts, intermediate_dim, hidden_dim] FP8
    topk_indices: [num_tokens, top_k]
    topk_weights: [num_tokens, top_k]
    """
    batch_size, hidden = x.shape
    output = torch.zeros_like(x)
    
    # 对每个 expert 执行 grouped GEMM
    for expert_id in range(num_experts):
        # 获取该 expert 处理的 token
        mask = (topk_indices == expert_id).any(dim=-1)
        if not mask.any():
            continue
            
        tokens_for_expert = x[mask]  # [num_tokens_for_expert, hidden]
        
        # 转换为 FP8
        tokens_fp8 = tokens_for_expert.to(torch.float8_e4m3fn)
        
        # 计算缩放因子（简化示例，实际应使用更精细的缩放策略）
        scale_x = torch.ones(tokens_fp8.size(0), device=x.device, dtype=torch.float32)
        scale_w = torch.ones(intermediate_dim, device=x.device, dtype=torch.float32)
        
        # 执行 FP8 GEMM: up-projection
        up_proj = torch.zeros(tokens_fp8.size(0), intermediate_dim, 
                             device=x.device, dtype=torch.bfloat16)
        up_proj = deep_gemm.fp8_gemm_nt(
            tokens_fp8, scale_x,
            expert_weights[expert_id], scale_w,
            up_proj
        )
        
        # SwiGLU 激活（简化）
        gate, up = up_proj.chunk(2, dim=-1)
        activated = torch.nn.functional.silu(gate) * up
        
        # 累加到输出（考虑 topk 权重）
        expert_mask = (topk_indices == expert_id)
        for i, (token_idx, weight) in enumerate(zip(
            torch.where(expert_mask)[0],
            topk_weights[expert_mask]
        )):
            output[token_idx] += activated[i] * weight
    
    return output

5.4 性能调优技巧

import deep_gemm

# 1. 设置使用的 SM 数量
# 在某些场景下，限制 SM 使用可以提高多任务并行效率
deep_gemm.set_num_sms(100)  # 使用 100 个 SM

# 2. 设置 Tensor Core 利用率目标
# 这是一个近似值，用于指导 kernel 选择
deep_gemm.set_tc_util(0.9)  # 目标 90% Tensor Core 利用率

# 3. 启用 PDL（如果适用）
deep_gemm.set_pdl(True)

# 4. 对于连续布局的 grouped GEMM，设置对齐要求
# 这影响 kernel 如何选择 tile 大小
deep_gemm.set_mk_alignment_for_contiguous_layout(128)

# 5. 启用 NVRTC 加速编译（开发/调试阶段）
# export DG_JIT_USE_NVRTC=1
# 注意：某些场景下可能有性能损失

第六章：性能对比与基准测试

6.1 与 CUTLASS 的性能对比

DeepGEMM 官方声称其性能"matches or exceeds expert-tuned libraries across various matrix shapes"。这主要得益于：

针对特定形状的 JIT 优化：相比通用库的静态模板，JIT 可以为特定 (M, N, K) 生成最优代码
简化的代码路径：减少模板元编程带来的开销
针对 DeepSeek 工作负载的定制：优化了 MoE 场景下的典型矩阵形状

6.2 H800 上的实测数据

根据官方数据，DeepGEMM 在 H800 上达到了 1550 TFLOPS 的峰值性能。让我们理解这个数字的含义：

H800 SXM5 的 FP8 Tensor Core 理论峰值约为 1979 TFLOPS
1550 TFLOPS 相当于约 78% 的硬件利用率
这是一个非常高的数字，考虑到实际 GEMM 运算中的内存带宽瓶颈和调度开销

6.3 不同矩阵形状的表现

GEMM 性能高度依赖于矩阵形状：

场景	M	N	K	典型效率
Dense 训练	8192	8192	8192	>90%
MoE expert (大)	1024	4096	4096	~85%
MoE expert (小)	128	4096	4096	~70%
Decoding	1	4096	4096	~50%

DeepGEMM 通过 grouped GEMM 和 masked GEMM 技术，在小 batch 场景下仍能保持较高效率。

第七章：DeepGEMM 的局限与未来

7.1 当前局限

硬件限制：仅支持 SM90 (Hopper) 和 SM100 (Blackwell)，不支持 Ampere 及更早架构
CUDA 版本要求：需要 CUDA 12.3+，推荐 12.9+
内存布局约束：对输入数据的内存布局有特定要求，需要预处理
FP8 精度限制：某些对精度敏感的场景可能需要保留 FP16/BF16 路径

7.2 未来发展方向

根据 DeepGEMM 的更新日志，我们可以预见以下发展方向：

更广泛的架构支持：可能扩展到更多 GPU 架构
NVRTC 完善：当前 NVRTC 支持是实验性的，未来可能成为默认选项
更多融合 kernel：类似 Mega MoE 的融合策略可能扩展到其他场景
FP4 普及：随着 Blackwell 的普及，FP4 支持将更加完善

7.3 对行业的影响

DeepGEMM 的开源标志着 AI 基础设施领域的几个重要趋势：

FP8 成为主流：2025-2026 年，FP8 将从"实验性特性"变为"生产默认"
领域专用库崛起：通用库（如 CUTLASS）面临领域专用库的挑战
开源生态竞争：DeepSeek 正在构建完整的开源 AI 基础设施栈，与 NVIDIA 的闭源方案竞争
JIT 编译范式：运行时编译可能成为高性能计算库的新标准

第八章：总结与思考

8.1 DeepGEMM 的核心价值

DeepGEMM 不仅仅是一个高性能 GEMM 库，它代表了 AI 计算优化的几个关键理念：

极简主义：用更少的代码实现更高的性能
领域专用：针对 LLM/MoE 场景深度优化，而非追求通用性
JIT 优先：运行时编译打破静态优化的限制
开源协作：通过开源社区持续迭代优化

8.2 对开发者的启示

对于从事 AI 基础设施开发的工程师，DeepGEMM 提供了以下启示：

深入理解硬件：只有充分理解 Tensor Core、内存层次结构，才能写出最优代码
敢于做减法：删除不必要的抽象和配置，专注于核心场景
拥抱 JIT：运行时编译是平衡灵活性和性能的有效手段
测量驱动优化：基于真实工作负载的性能数据进行优化

8.3 结语

在 AI 计算的黄金时代，DeepGEMM 展示了开源社区与硬件厂商共同推动技术进步的典范。当 1550 TFLOPS 成为新的起点，我们有理由期待，未来的 AI 基础设施将更加高效、开放和民主化。

DeepSeek 通过开源 DeepGEMM、DeepEP 等核心组件，正在重塑 AI 基础设施的竞争格局。这不仅是技术的胜利，更是开源精神的胜利——当知识被共享，创新的速度将超越任何单一组织的边界。

参考资源

本文基于 DeepGEMM 开源代码和公开技术资料撰写，旨在帮助开发者理解 FP8 GEMM 优化的核心技术与工程实践。