专栏名称: 新智元
智能+中国主平台,致力于推动中国从互联网+迈向智能+新纪元。重点关注人工智能、机器人等前沿领域发展,关注人机融合、人工智能和机器人革命对人类社会与文明进化的影响,领航中国新智能时代。
目录
相关文章推荐
人工智能那点事  ·  DeepSeek:多地严禁其自动生成处方 ·  昨天  
爱可可-爱生活  ·  突破性的“一步扩散”生成模型 查看图片 ... ·  昨天  
心禅道  ·  deepseek推荐导致 AI搜索引擎优化 ... ·  2 天前  
心禅道  ·  deepseek推荐导致 AI搜索引擎优化 ... ·  2 天前  
爱可可-爱生活  ·  【[3.5k星]Bananas ... ·  2 天前  
51好读  ›  专栏  ›  新智元

DeepSeek-R2曝5月前上线!第三弹DeepGEMM 300行代码暴击专家优化内核

新智元  · 公众号  · AI  · 2025-02-26 10:16

正文



新智元报道

编辑:编辑部 HNYZ
【新智元导读】 DeepSeek开源第三弹,是支持稠密和MoE模型的FP8计算库——DeepGEMM,支持V3/R1训推。仅凭300行代码,就超过了专家优化的内核。开发者惊叹:DeepSeek有最好的GPU工程师,仿佛拥有某种编译器黑魔法!更令人兴奋的是,DeepSeek-R2有望在5月前提前发布。

第三天,DeepSeek发布了DeepGEMM。

这是一个支持稠密和MoE模型的FP8 GEMM(通用矩阵乘法)计算库,可为V3/R1的训练和推理提供强大支持。

仅用300行代码,DeepGEMM开源库就能超越专家精心调优的矩阵计算内核,为AI训练和推理带来史诗级的性能提升!

DeepGEMM库具有以下特征:

  • 在Hopper GPU上实现高达1350+ FP8 TFLOPS的算力

  • 极轻量级依赖,代码清晰易懂

  • 完全即时编译,即用即跑

  • 核心逻辑仅约300行代码,却在大多数矩阵规模下超越专家级优化内核

  • 同时支持密集布局和两种MoE布局

开发者惊叹道:才300行代码,就能打败专家优化的内核?!
要么是DeepSeek真的破解了GPU运算的天机,要么我们就是见证了有史以来最高级的编译器黑科技。

总之,这个DeepGEMM听起来简直是数学界的超级英雄,比飞快的计算器还要快。

它改变了我们使用FP8 GEMM库的方式,简单、快速、开源。这就是AI计算的未来!

同时,外媒还曝出了另一个重磅消息:原计划在5月初发布的DeepSeek-R2,现在发布时间将再次提前!

在DeepSeek-R2中,将实现更好的编码,还能用英语以外的语言进行推理。

业内人士预测,DeepSeek-R2的发布,将是AI行业的一个关键时刻。目前DeepSeek在创建高成本效益模型上的成功,已经打破了该领域少数主导玩家的垄断。

DeepSeek开源两天,前两个项目爆火程度难以想象。FlashMLA已在GitHub斩获近10k星标,DeepEP的星标已有5k。

DeepGEMM

DeepGEMM是一个专为清晰高效的FP8通用矩阵乘法(General Matrix Multiplications,GEMMs)设计的库,它采用了DeepSeek-V3中提出的细粒度缩放技术。

该库支持常规矩阵乘法和混合专家模型(Mix-of-Experts,MoE)分组矩阵乘法。DeepGEMM使用CUDA编写,无需在安装时进行编译,而是通过轻量级即时编译( Just-In-Time, JIT)模块在运行时编译所有内核。

目前,DeepGEMM仅支持NVIDIA Hopper张量核。为了解决FP8张量核在累加计算时的精度问题,该库采用了基于CUDA核心的两级累加(提升)技术。

虽然DeepGEMM借鉴了CUTLASS和CuTe的一些概念,但避免了过度依赖它们的模板或代数系统。

相反,该库追求设计简洁,仅包含一个核心内核函数,代码量仅约300行。这使其成为学习Hopper FP8矩阵乘法和优化技术的理想入门资源。

尽管采用轻量级设计,DeepGEMM在处理各种矩阵形状时的性能都能够达到甚至超越经专家调优的库。

性能

研究人员在配备NVCC 12.8的H800上测试了DeepSeek-V3/R1推理过程中,可能使用的所有矩阵形状(包括预填充和解码阶段,但不包括张量并行计算)。

所有性能提升指标均与基于CUTLASS 3.6内部精心优化的实现进行对比计算得出。

DeepGEMM在某些矩阵形状下的表现还不够理想,如果你对此感兴趣,可以提交优化相关的Pull Request(拉取请求)。

稠密模型的常规GEMM

下表展示了不同矩阵维度(M、N、K)下DeepGEMM库的性能数据,结果显示在某些配置(如 M=128, N=2112, K=7168)下实现了高达 2.4 倍的加速,反映了DeepGEMM在优化GPU矩阵计算方面的效率和灵活性。

MoE模型的分组GEMM(使用连续存储布局)

MoE模型的分组GEMM(使用掩码存储布局)

快速入门


要求

  • NVIDIA Hopper架构GPU(需支持sm_90a计算能力)

  • Python v3.8或更高版本

  • CUDA v12.3及以上版本(强烈建议使用v12.8或更新版本以获得最佳性能)

  • PyTorch v2.1及以上版本

  • CUTLASS v3.6或更高版本 (可通过Git子模块[submodule]方式克隆获取)

开发

下面代码是DeepGEMM项目的安装和测试指南。

首先,通过命令克隆仓库及其子模块。 然后,创建第三方库(CUTLASS和CuTe)的符号链接以便开发。 接着,测试JIT编译功能。 最后,测试所有GEMM实现。

# Submodule must be clonedgit clone --recursive [email protected]:deepseek-ai/DeepGEMM.git
# Make symbolic links for third-party (CUTLASS and CuTe) include directoriespython setup.py develop
# Test JIT compilationpython tests/test_jit.py
# Test all GEMM implements (normal, contiguous-grouped and masked-grouped)python tests/test_core.py

安装

下面代码使用脚本安装Python包,会将包及其依赖项安装到系统中以便在项目中使用。

python setup.py install

接下来,在你的Python项目中导入deep_gemm,就可以开始使用啦!

优化技术

注意:下面用🐳标记的是,CUTLASS中未包含的技术。

持久化线程束专用化

遵循CUTLASS的设计,DeepGEMM中的内核采用线程束(warp)专用化技术,实现了数据移动、张量核心MMA(矩阵乘累加)指令和CUDA核心提升操作的重叠执行。下图简要说明了这个过程:

TMA线程主要负责数据加载(Data load)和任务分发(TMA issue),用黄色和蓝色表示。数学线程则交替执行WGMA(Wavefront Matrix Multiply-Accumulate)计算(绿色)和数据提升(Promotion,黄色),展示了一种并行计算策略,其中数据加载与矩阵计算和优化操作协同工作,以提高效率和性能。

Hopper TMA特性

张量内存加速器(Tensor Memory Accelerator,TMA)是Hopper架构引入的新硬件特性,用于实现更快速的异步数据移动。具体来说,在以下方面使用TMA:

  • LHS(左矩阵)、LHS缩放因子和RHS(右矩阵)的TMA加载

  • 输出矩阵的TMA存储

  • LHS矩阵的TMA多播

  • TMA描述符预取






请到「今天看啥」查看全文