DeepSeek代码开源第三弹:DeepGEMM代码库,V3/R1的训练推理动力

2月26日消息,在宣布开源MLA解码核FlashMLA以及DeepEP两款代码库后,DeepSeek在开源周的第三天宣布开放DeepGEMM代码库。

DeepSeek介绍,DeepGEMM是专为简洁高效的FP8通用矩阵乘法(GEMMs)而设计,它同时支持普通的和专家混合(MoE)分组的GEMM运算,为V3/R1训练和推理提供动力支持。该库使用CUDA编写,在安装过程中无需编译,通过在运行时使用轻量级即时编译模块来编译所有内核。

目前,DeepGEMM仅支持英伟达Hopper架构运算,为解决FP8张量核心累加不精确的问题,它采用了CUDA核心的两级累加(提升)方法。该代码库设计非常简洁,只有一个核心内核函数,代码量约为300行。

尽管其设计轻巧,DeepGEMM的性能在各种矩阵形状上与专家调优的库相匹配或超越。

DeepSeek团队在H800上使用NVCC 12.8测试了DeepSeek-V3/R1推理中可能使用的所有形状(包括预填充和解码,但没有张量并行)。
从测试结果来看,DeepGEMM计算性能最高可达1358 TFLOPS,内存宽带最高可达2668 GB/s。与基于CUTLASS 3.6的优化实现相比,可提速最高可达2.7倍。另外,分组GEMM(MoE模型)中连续性布局、掩码布局下可提速多达1.2倍。

另外,使用DeepGEMM需要的环境要求,包括:

* 必须支持Hopper架构的GPU,sm_90a

* Python 3.8及以上

* CUDA 12.3及以上(推荐12.8)

* PyTorch 2.1及以上

* CUTLASS 3.6及以上

 
举报 收藏 打赏 评论 0
24小时热闻
今日推荐
浙ICP备19001410号-1