DeepSeek公司在近期举行的开源周活动中,持续推动其技术开放步伐,继MLA解码核FlashMLA和DeepEP代码库之后,再度迈出重要一步,正式对外公布了DeepGEMM代码库。
DeepGEMM是一款专为高效FP8通用矩阵乘法(GEMM)设计的工具库,其核心目标是强化V3/R1训练和推理任务的处理能力。该库不仅支持传统的矩阵运算,还特别针对专家混合(MoE)分组的GEMM运算进行了优化,全部采用CUDA语言编写。
据悉,DeepGEMM的一大亮点在于其安装流程的便捷性。用户无需进行繁琐的编译操作,而是依靠运行时的轻量级即时编译模块,动态生成所需的内核代码。这一设计极大地提升了使用的便利性。
目前,DeepGEMM主要适用于英伟达Hopper架构的硬件环境。为了解决FP8张量核心累加过程中可能出现的精度问题,该库创新性地采用了两级累加方法,确保能够充分利用CUDA核心的性能优势。DeepGEMM的代码设计极为精简,其核心功能高度集成在一个内核函数中,整个代码库的总代码量仅为约300行。
尽管体积小巧,但DeepGEMM在计算性能上的表现却不容小觑。经过DeepSeek团队的严格测试,该库在不同矩阵形状下的计算性能均表现出色,甚至能够与经过深度优化的专业库相媲美或超越。特别是在H800 GPU上使用NVCC 12.8编译器进行全面评估时,DeepGEMM的计算性能最高可达1358 TFLOPS,内存带宽峰值为2668 GB/s。与基于CUTLASS 3.6的优化实现相比,其性能提升幅度最高可达2.7倍。在分组GEMM(MoE模型)的连续性布局和掩码布局下,DeepGEMM的性能提升同样显著,可达1.2倍以上。
值得注意的是,使用DeepGEMM还需满足一定的环境要求。硬件方面,需要支持Hopper架构(sm_90a);操作系统则需安装Python 3.8或更高版本;CUDA版本需为12.3及以上(推荐使用12.8);同时,还需配备PyTorch 2.1及以上版本以及CUTLASS 3.6及以上版本。这些要求确保了DeepGEMM能够在最佳状态下运行,发挥出其卓越的计算性能。