5分钟看懂Deepseek开源周之三：开源核弹！DeepGEMM凭FP8+MoE双杀：1350 TFLOPS碾压cuBLAS，大模型训练成本暴降50%
深度求索开源周第三天：DeepGEMM代表了大模型底层计算从“依赖硬件厂商优化”向“开源算法驱动”的范式转移，其技术方向将显著降低大模型训练门槛、加速稀疏架构创新，并推动行业形成更开放的软硬件协同生态。短期看，它可能成为MoE模型训练的事实标准；长期而言，其设计理念或影响下一代AI芯片的指令集架构设计。
前言
深度求索开源周第三天：DeepGEMM代表了大模型底层计算从“依赖硬件厂商优化”向“开源算法驱动”的范式转移，其技术方向将显著降低大模型训练门槛、加速稀疏架构创新，并推动行业形成更开放的软硬件协同生态。短期看，它可能成为MoE模型训练的事实标准；长期而言，其设计理念或影响下一代AI芯片的指令集架构设计。
https://github.com/deepseek-ai/open-infra-index/tree/main?tab=readme-ov-file开源第四天：DeepGEMM
在DeepSeek开源周的第三天介绍了DeepGEMM。这是一个FP8 GEMM库，支持密集和MoE（混合专家）GEMMs，为V3/R1训练和推理提供动力。
以下是DeepGEMM的主要特点：
- 在Hopper GPU上最高可达1350+ FP8 TFLOPS。
- 没有复杂的依赖关系，代码简洁如教程。
- 完全即时编译。
- 核心逻辑约300行代码，但在大多数矩阵尺寸上超越了专家调优的内核。
- 支持密集布局和两种MoE布局。
1. 计算效率革命与成本降低
FP8计算范式的突破：通过引入FP8精度，DeepGEMM在保持合理数值精度的同时，将计算吞吐量提升至1350+ TFLOPS（Hopper GPU），相比FP16/FP32训练，显存占用和通信带宽需求大幅下降。这将直接降低大模型训练与推理的硬件成本（如减少GPU使用量或缩短训练周期），尤其对千亿参数以上模型的训练经济性提升显著。
MoE架构的实用化加速：支持两种MoE布局的高效计算，解决了传统MoE模型中专家路由带来的计算碎片化问题。通过优化专家并行与数据并行的结合，MoE模型（如万亿参数级别）的训练和推理效率将逼近密集模型，进一步推动稀疏化专家模型成为行业主流选择。
2. 软硬件协同设计的新标杆
即时编译（JIT）与架构适配：完全依赖JIT而非预编译内核，使得DeepGEMM能动态适应不同GPU架构（如Hopper/Ampere）和多样化的矩阵尺寸。这种灵活性为未来新型硬件（如专有AI芯片）的快速适配提供了模板，可能加速行业从“通用硬件适配”向“算法-硬件协同设计”的转变。
代码极简主义的影响：核心逻辑仅300行代码却超越专家手工优化，表明算法创新（如分层分块策略、异步流水线设计）比传统硬件微调更具潜力。这可能促使更多研究转向计算图编译器优化和数学近似算法，而非依赖硬件厂商的封闭库。
3. 大模型架构创新的催化剂
动态稀疏计算的可行性提升：DeepGEMM对不规则MoE计算的支持，为动态专家分配（如训练时专家数量自适应变化）提供了底层支持。未来模型可能更激进地探索“动态稀疏化”架构（如实时专家淘汰/生成），突破现有静态MoE的设计限制。
混合精度训练的普及：FP8在训练中的稳定应用（需配合梯度量化策略），将推动行业重新评估训练精度需求。低精度训练可能从推理延伸至全流程，结合量化感知训练（QAT），催生“原生低精度大模型”新范式。
4. 开源生态与行业格局重塑
去中心化优化能力：依赖简洁代码和零复杂依赖的特性，使得中小团队甚至个人研究者可直接参与GEMM优化，打破传统大厂通过封闭内核库（如cuBLAS）构建的技术壁垒。开源社区可能出现“众包式性能优化”模式，加速计算库迭代。
框架耦合度降低：作为独立计算库，DeepGEMM若被主流框架（PyTorch/TensorFlow）集成，将推动行业标准化分离“计算后端”与“模型框架”，增强技术栈的模块化。开发者可更专注于模型结构设计，而非底层计算适配。
潜在挑战与应对
-
FP8数值稳定性管理：需配合缩放因子（Scale Factor）动态调整和溢出监控策略，行业可能衍生出自动化精度管理工具。
-
硬件生态依赖：当前性能数据基于Hopper架构，需验证其在存量Ampere显卡上的表现，可能倒逼云厂商加速硬件升级。
-
MoE算法-库协同优化：库的高效性要求模型设计者调整专家路由策略以匹配计算特性（如专家组粒度优化），可能催生新的模型设计方法论。
DeepGEMM
DeepGEMM 是一个专为简洁高效的FP8通用矩阵乘法（GEMM）设计的计算库，其核心实现了DeepSeek-V3提出的细粒度缩放技术。该库不仅支持常规矩阵运算，还可处理混合专家（MoE）分组GEMM场景。基于CUDA开发的DeepGEMM无需安装时编译，通过轻量级即时编译（JIT）模块在运行时动态生成所有计算内核。
目前该库专为NVIDIA Hopper张量核心优化。针对FP8张量核心累加精度不足的硬件限制，创新性地采用CUDA核心两级精度提升累加策略。虽然借鉴了CUTLASS和CuTe的部分设计思想，但摒弃了复杂的模板依赖和代数架构，转而以极简主义为核心哲学——单一核心计算函数仅约300行代码，堪称学习Hopper架构FP8矩阵乘法及优化技术的绝佳教材。
尽管架构极度轻量化，DeepGEMM在各类矩阵形态下的性能表现却达到甚至超越了专家手工调优的专业计算库水平。
性能表现
我们在配备 NVCC 12.8 的 H800 SXM5 GPU 上测试了 DeepSeek-V3/R1 推理中可能用到的所有计算形态（包括预填充和解码阶段，但不涉及张量并行）。所有加速指标均基于我们内部通过 CUTLASS 3.6 精心优化的实现进行对比计算。
DeepGEMM 在某些计算形态中表现未达预期，欢迎感兴趣的研究者提交优化 Pull Request（PR）以改进性能。
密集模型常规GEMM
| M | N | K | Computation | Memory bandwidth | Speedup |
|---|---|---|---|---|---|
| 64 | 2112 | 7168 | 206 TFLOPS | 1688 GB/s | 2.7x |
| 64 | 24576 | 1536 | 289 TFLOPS | 2455 GB/s | 1.7x |
| 64 | 32768 | 512 | 219 TFLOPS | 2143 GB/s | 1.8x |
| 64 | 7168 | 16384 | 336 TFLOPS | 2668 GB/s | 1.4x |
| 64 | 4096 | 7168 | 287 TFLOPS | 2320 GB/s | 1.4x |
| 64 | 7168 | 2048 | 295 TFLOPS | 2470 GB/s | 1.7x |
| 128 | 2112 | 7168 | 352 TFLOPS | 1509 GB/s | 2.4x |
| 128 | 24576 | 1536 | 535 TFLOPS | 2448 GB/s | 1.6x |
| 128 | 32768 | 512 | 358 TFLOPS | 2103 GB/s | 1.5x |
| 128 | 7168 | 16384 | 645 TFLOPS | 2604 GB/s | 1.4x |
| 128 | 4096 | 7168 | 533 TFLOPS | 2221 GB/s | 2.0x |
| 128 | 7168 | 2048 | 510 TFLOPS | 2277 GB/s | 1.7x |
| 4096 | 2112 | 7168 | 1058 TFLOPS | 527 GB/s | 1.1x |
| 4096 | 24576 | 1536 | 990 TFLOPS | 786 GB/s | 1.0x |
| 4096 | 32768 | 512 | 590 TFLOPS | 1232 GB/s | 1.0x |
| 4096 | 7168 | 16384 | 1358 TFLOPS | 343 GB/s | 1.2x |
| 4096 | 4096 | 7168 | 1304 TFLOPS | 500 GB/s | 1.1x |
| 4096 | 7168 | 2048 | 1025 TFLOPS | 697 GB/s | 1.1x |
小矩阵（M=64）：
2.7倍加速比 + 1688 GB/s高带宽 → 计算与访存双优化。
典型场景：Transformer解码阶段（单Token生成，M=序列长度），性能优势显著。
超大矩阵（M=4096）：
1358 TFLOPS逼近理论极限，但内存带宽骤降至343 GB/s → 显存墙效应显现。
加速比仅1.2x → 可能受限于Hopper架构的HBM3带宽瓶颈（如A100 HBM2e带宽1555 GB/s vs H800 HBM3 3.35 TB/s）。
MoE模型分组GEMM（连续布局）
| #Groups | M per group | N | K | Computation | Memory bandwidth | Speedup |
|---|---|---|---|---|---|---|
| 4 | 8192 | 4096 | 7168 | 1297 TFLOPS | 418 GB/s | 1.2x |
| 4 | 8192 | 7168 | 2048 | 1099 TFLOPS | 681 GB/s | 1.2x |
| 8 | 4096 | 4096 | 7168 | 1288 TFLOPS | 494 GB/s | 1.2x |
| 8 | 4096 | 7168 | 2048 | 1093 TFLOPS | 743 GB/s | 1.1x |
计算性能稳定在1200+ TFLOPS → 专家并行计算的高效融合。
低加速比（1.1-1.2x） → 现有MoE库（如Megablocks）已针对连续专家分配优化，DeepGEMM需进一步优化负载均衡策略。
MoE模型分组GEMM（掩码布局）
|
#Groups |
M per group |
N |
K |
Computation |
Memory bandwidth |
Speedup |
|---|---|---|---|---|---|---|
|
1 |
1024 |
4096 |
7168 |
1233 TFLOPS |
924 GB/s |
1.2x |
|
1 |
1024 |
7168 |
2048 |
925 TFLOPS |
968 GB/s |
1.2x |
|
2 |
512 |
4096 |
7168 |
1040 TFLOPS |
1288 GB/s |
1.2x |
|
2 |
512 |
7168 |
2048 |
916 TFLOPS |
1405 GB/s |
1.2x |
|
4 |
256 |
4096 |
7168 |
932 TFLOPS |
2064 GB/s |
1.1x |
|
4 |
256 |
7168 |
2048 |
815 TFLOPS |
2047 GB/s |
1.2x |
2047 GB/s超高带宽利用率 → 动态稀疏掩码（如Top-2专家选择）的内存访问模式优化成功。
加速比保持1.2x → 证明其对不规则计算模式的适应性，可能通过异步流水线隐藏路由延迟。
DeepGEMM在中小规模矩阵计算场景中展现出显著性能优势（最高加速比达2.7倍），但在超大矩阵（如M=4096）下加速比收窄至1.1-1.2倍，暗示其优化方向更聚焦于大模型推理典型负载（如KV Cache动态解码）。同时，其对MoE模型分组计算的稳定支持（全场景加速比≥1.1x），验证了其在稀疏专家架构下的通用性。
名词解释：
1. 基础矩阵参数
M：输入矩阵A的行数，通常对应 批量大小（Batch Size） 或 序列长度（Sequence Length）。
示例：在解码阶段，M=64可能表示同时处理64个Token的生成。N：权重矩阵B的列数，常代表 隐层维度（Hidden Dimension） 或 专家输出维度。
示例：N=7168可能对应MoE模型中单个专家的输出神经元数量。K：输入矩阵A的列数/权重矩阵B的行数，通常为 输入特征维度（Input Feature Dimension）。
示例：K=1536可视为Transformer中多头注意力机制的键/值向量总维度。2. 性能指标
Computation (TFLOPS)：每秒浮点运算次数（单位：万亿次），衡量GPU计算核心的利用率。
示例：1358 TFLOPS表示每秒执行1.358×10¹²次浮点运算，接近Hopper GPU的理论峰值（约1500 TFLOPS）。Memory Bandwidth (GB/s)：内存带宽（单位：千兆字节/秒），反映数据从显存到计算单元的传输速度。
示例：2143 GB/s表示每秒可搬运2143GB数据，若超过GPU显存带宽上限（如H800为3.35TB/s），则成为性能瓶颈。Speedup：相对于基准库（CUTLASS 3.6优化版本）的加速倍数。
示例：2.7x表示DeepGEMM在此场景下比CUTLASS快2.7倍。3. MoE模型专用参数
#Groups：专家分组的数量，对应 并行计算的专家子集数。
示例：#Groups=8表示将计算任务拆分为8个专家组并行处理。M per group：每个专家组处理的输入行数，影响 负载均衡与缓存利用率。
示例：M per group=256表示每个专家组处理256行输入数据。Contiguous Layout：连续内存布局，专家组的输入/输出在显存中连续存储，适合静态专家分配。
特点：高计算性能（如1288 TFLOPS），但带宽利用率较低（494 GB/s），因数据局部性较好。Masked Layout：掩码布局，通过稀疏掩码跳过无效计算，适合动态专家路由（如Top-K选择）。
特点：带宽利用率极高（如2047 GB/s），因需动态加载非连续数据，但计算性能略低（815 TFLOPS）。4. 典型场景映射
参数组合 对应模型阶段 技术挑战 M=64, N=7168, K=2048 Transformer解码（单Token生成） 小矩阵计算效率优化 M=4096, N=16384, K=7168 预填充（Prompt处理） 显存带宽瓶颈突破 #Groups=4, M per group=256 MoE动态路由（如Switch Transformer） 不规则计算并行化
安装条件
依赖项
系统要求
-
GPU架构：必须支持 Hopper架构的GPU（计算能力sm_90a及以上）
-
Python版本：3.8 或更高版本
-
CUDA版本：12.3 或更高版本（强烈推荐使用12.8及以上版本以获得最佳性能）
-
PyTorch版本：2.1 或更高版本
-
CUTLASS版本：3.6 或更高版本（可通过
Git submodule克隆）
Hopper架构 (sm_90a)
- NVIDIA H100/H800 GPU的硬件架构代号，需支持计算能力版本sm_90a。
CUDA 12.8
- NVIDIA并行计算平台版本，高版本（≥12.8）可充分发挥Hopper FP8张量核心性能。
CUTLASS 3.6
- NVIDIA开源的高性能矩阵计算库，需通过
git submodule add命令从仓库克隆依赖。
安装准备
# Submodule must be cloned
git clone --recursive git@github.com:deepseek-ai/DeepGEMM.git
# Make symbolic links for third-party (CUTLASS and CuTe) include directories
python setup.py develop
# Test JIT compilation
python tests/test_jit.py
# Test all GEMM implements (normal, contiguous-grouped and masked-grouped)
python tests/test_core.py安装启动
python setup.py install然后，在python脚本中：
import deep_gemm接口说明
核心特性
-
专用GEMM核
-
仅包含GEMM计算内核，要求输入矩阵的**左侧缩放因子（LHS Scaling Factor）**满足 TMA对齐（Tensor Memory Access对齐） 且 已转置，仅支持 NT格式（左侧矩阵非转置，右侧矩阵转置）。
-
不支持转置或FP8类型转换操作，需用户自行在前置计算核中实现或融合这些操作。
-
附带的PyTorch工具函数性能可能较低，核心优化聚焦于GEMM内核本身。
-
功能接口说明
1. 常规密集GEMM（非分组）
-
调用函数：
deep_gemm.gemm_fp8_fp8_bf16_nt-
功能：执行基础的非分组FP8 GEMM运算（左侧矩阵非转置，右侧矩阵转置）。
-
适用场景：标准全连接层或注意力机制中的矩阵乘法。
-
2. 分组GEMM（连续布局）
-
调用函数：
m_grouped_gemm_fp8_fp8_bf16_nt_contiguous-
设计特点：
-
仅对M轴分组，N和K维度固定，专为 MoE模型专家形状一致 的场景优化。
-
连续布局：在训练前向或推理预填充阶段，将不同专家处理的Token拼接为单一连续张量。
-
-
对齐要求：每个专家段的M轴长度需满足 GEMM M块对齐（通过
get_m_alignment_for_contiguous_layout()获取对齐值）。
-
3. 分组GEMM（掩码布局）
-
调用函数：
m_grouped_gemm_fp8_fp8_bf16_nt_masked-
适用场景：
-
推理解码阶段启用 CUDA图（CPU无法预知各专家处理的Token数量）。
-
通过 掩码张量（Mask Tensor） 动态跳过无效计算区域。
-
-
典型应用：以DeepEP低延迟核的输出作为输入，实现动态专家路由。
-
工具函数
-
硬件资源配置：
-
deep_gemm.set_num_sms：设置最大使用的 流式多处理器（SM）数量。 -
deep_gemm.get_num_sms：获取当前SM数量上限。
-
-
内存对齐工具：
-
deep_gemm.get_m_alignment_for_contiguous_layout：获取连续布局的 M轴对齐要求。 -
deep_gemm.get_tma_aligned_size：获取 TMA对齐所需尺寸。 -
deep_gemm.get_col_major_tma_aligned_tensor：生成 列主序TMA对齐张量。
-
环境变量
| 变量名 | 类型 | 功能说明 | 默认值 |
|---|---|---|---|
| DG_CACHE_DIR | 字符串 | 编译内核缓存目录 | $HOME/.deep_gemm |
| DG_NVCC_COMPILER | 字符串 | 指定NVCC编译器路径 | 从PyTorch CUDA_HOME自动探测 |
| DG_DISABLE_FFMA_INTERLEAVE | 0/1 | 禁用 FFMA指令交错优化（影响计算吞吐量） | 0（启用优化） |
| DG_PTXAS_VERBOSE | 0/1 | 显示PTX汇编编译器详细输出 | 0（关闭） |
| DG_PRINT_REG_REUSE | 0/1 | 打印FFMA指令的寄存器重用细节 | 0（关闭） |
| DG_JIT_PRINT_NVCC_COMMAND | 0/1 | 打印JIT编译时的NVCC命令 | 0（关闭） |
| DG_JIT_DEBUG | 0/1 | 输出JIT调试信息（如内核参数校验） | 0（关闭） |
优化技术详解
（标注🐳的为CUTLASS未采用的技术）
1. 持久化线程束特化（Persistent Warp-Specialization）
借鉴CUTLASS设计，DeepGEMM通过线程束（Warp）特化实现数据移动、张量核心MMA指令（矩阵乘加操作）与CUDA核心提升操作的重叠执行。这种设计类似“流水线并行”机制：
- 传统方式：数据加载完成后才进行计算，导致硬件资源闲置。
- DeepGEMM优化：数据加载与计算交替进行，GPU各单元始终处于忙碌状态，显著提升吞吐量。
设计示意图
Hopper TMA 特性
Tensor Memory Accelerator (TMA) 是Hopper架构引入的一种新硬件特性，旨在实现更快且异步的数据移动。具体来说，我们利用TMA进行以下操作：
- TMA加载：用于左侧矩阵（LHS）、左侧缩放因子和右侧矩阵（RHS）
- TMA存储：用于输出矩阵
- TMA多播：仅限于左侧矩阵
- TMA描述符预取
常见细节优化
- 使用stmatrix PTX指令
- 针对不同战线组定制寄存器数量控制
- 尽可能重叠操作，例如TMA存储与非TMA右侧矩阵缩放因子加载的重叠
- 统一且优化的块调度器
- 一个调度器适用于所有非分组和分组内核
- 光栅化以增强L2缓存重用
图片信息解释
- TMA warps：表示处理TMA相关的战线，包括TMA指令发出和数据加载。
- Math warps 0 和 Math warps 1：表示处理数学运算的战线，包括WGMMA（Warped General Matrix Multiplication Accumulate）和提升（Promotion）。
2. Hopper TMA特性深度利用
基于Hopper架构的张量内存加速器（TMA），实现异步高效数据传输：
- TMA加载：左操作数（LHS）、缩放因子及右操作数（RHS）矩阵的异步加载。
- TMA存储：输出矩阵的异步写回。
- TMA多播：LHS矩阵的单次加载多线程共享（减少重复传输）。
- TMA描述符预取：提前加载数据布局元信息，减少指令延迟。
3. 通用细节优化
- PTX指令级控制：采用
stmatrix指令实现线程束级矩阵存储优化。 - 寄存器分配策略：针对不同线程束组（Warpgroup）动态调整寄存器数量，避免资源竞争。
- 操作重叠🐳：TMA存储与非TMA的RHS缩放因子加载重叠执行，最大化硬件利用率。
4. 统一优化的块调度与光栅化
- 单调度器设计：适用于所有分组/非分组内核，简化资源管理。
- 光栅化（Rasterization）：优化线程块分布以增强L2缓存复用率，降低内存访问延迟。
5. 完全JIT设计🐳
- 动态编译优势：所有内核在运行时即时编译（JIT），将矩阵形状、块大小等参数作为编译常量，允许编译器深度优化。
- 自动参数选择：根据硬件特性动态选择最优块大小、线程束组数量及流水线阶段，无需手动调优。
- 小矩阵优化：完全展开MMA流水线指令，减少分支预测开销，对小规模计算效率提升显著。
6. 非对齐块大小支持🐳
- 问题背景：传统2的幂次方块大小（如128×128）可能导致流式多处理器（SM）利用率不足。
- DeepGEMM方案：支持112等非对齐块大小，例如处理M=256、N=7168矩阵时，SM利用率从112提升至128（总132个），计算效率提升14%。
7. FFMA SASS指令交错优化🐳
- 二进制指令修改：通过调整FFMA（融合乘加）指令的
yield位（控制线程束调度优先级）和reuse位（寄存器重用），增加MMA指令与提升操作的并行度。 - 效果：在细粒度缩放FP8 GEMM中，性能提升超过10%。
代码仓展示
仅仅5天斩获4.5K星，实力强劲！
文件目录：
deep_gemm：这是主要的代码目录，包含实现DeepGEMM库的核心功能。它支持密集和MoE（混合专家）GEMMs，并为V3/R1训练和推理提供动力。
figures：这个目录可能包含用于文档或演示的图表、图像或其他可视化内容。
tests：这个目录包含测试脚本和数据，用于验证DeepGEMM库的功能和性能。
third-party：这个目录包含第三方库或依赖项，这些库或依赖项对于DeepGEMM库的运行是必要的。
.gitignore：这是一个配置文件，用于指定在Git版本控制系统中忽略的文件或目录模式。
.gitmodules：这是一个配置文件，用于管理Git子模块，即作为独立项目嵌入到主项目中的其他Git仓库。
LICENSE：这是一个法律文件，定义了使用DeepGEMM库的许可条款和条件。
README.md：这是一个Markdown格式的文档文件，通常包含项目的介绍、安装指南、使用说明和其他重要信息。
setup.py：这是一个Python脚本，用于设置和打包项目，以便于安装和分发。
deep_gemm
fp8_gemm.cuh：
- 这个文件包含使用FP8（浮点数8位）格式的GEMM（General Matrix Multiplication）操作的CUDA实现。它支持密集和MoE（混合专家）GEMMs。
mma_utils.cuh：
- 这个文件包含矩阵乘法（Matrix Multiply Accumulate, MMA）相关的实用函数。这些函数可能用于优化矩阵乘法操作。
scheduler.cuh：
- 这个文件包含调度器相关的代码，用于管理任务的执行顺序和资源分配。调度器确保计算任务高效地运行在GPU上。
tma_utils.cuh：
- 这个文件包含与Tensor Memory Access（TMA）相关的实用函数。这些函数可能用于优化内存访问模式，以提高性能。
utils.cuh：
- 这个文件包含通用的辅助函数和工具，这些函数和工具在整个项目中可能会被多个模块使用。这些工具可能包括数据处理、错误处理等功能。
欢迎加入DeepSeek 技术社区。在这里，你可以找到志同道合的朋友，共同探索AI技术的奥秘。
更多推荐
7
0- 0
已为社区贡献10条内容
所有评论(0)