北京,[日期] – DeepSeek 近日开源了其 FP8 通用矩阵乘法库 DeepGEMM,旨在为 NVIDIA Hopper 架构提供高效简洁的 FP8 矩阵乘法 (GEMM) 解决方案。这一举措有望显著加速大规模 AI 模型的推理和训练过程,尤其是在混合专家 (MoE) 模型等场景中。
DeepGEMM 是什么?
DeepGEMM 是一个专为高效 FP8 矩阵乘法设计的库,目前仅支持 NVIDIA Hopper 张量核心。它基于即时编译 (JIT) 技术,无需安装时编译,并支持运行时动态优化。DeepGEMM 的核心优势在于其细粒度缩放和 CUDA 核心双级累加技术,有效解决了 FP8 精度不足的问题,同时利用 Hopper 的 Tensor Memory Accelerator (TMA) 特性提升数据传输效率。更令人瞩目的是,DeepGEMM 的核心代码仅约 300 行,易于学习和优化。
主要功能亮点:
- 高效 FP8 矩阵乘法: 专为 FP8 (8 位浮点) 矩阵乘法设计,支持细粒度缩放,显著提升矩阵运算的性能和精度。
- 支持普通和分组 GEMM: 适用于常规矩阵乘法操作,并支持 MoE 模型中的分组矩阵乘法,优化多专家共享形状的场景。
- 即时编译 (JIT) 设计: 所有内核在运行时动态编译,无需安装时编译,并根据矩阵形状、块大小等参数进行优化。
- Hopper 架构优化: 充分利用 NVIDIA Hopper 架构的 TMA 特性,包括 TMA 加载、存储、多播和描述符预取,显著提升数据传输效率。
- 细粒度缩放和双级累加: 通过细粒度缩放技术和基于 CUDA 核心的双级累加机制,将 FP8 计算结果提升到更高精度的格式 (如 BF16),确保计算精度。
- 轻量级设计: 核心代码简洁,仅约 300 行,易于理解和扩展。
性能表现:
DeepGEMM 在多种矩阵形状上均达到或超过专家级优化库的性能。
- 普通 GEMM (非分组): 在某些矩阵形状下,DeepGEMM 能达到 2.7 倍的加速比,实现超过 1000 TFLOPS 的计算性能,接近 Hopper 架构 GPU 的理论峰值。
- 分组 GEMM (MoE 模型): DeepGEMM 能实现 1.1-1.2 倍的加速比,显著提升 MoE 模型的训练和推理效率。
系统要求:
- 硬件: NVIDIA Hopper 架构 GPU (支持 sm_90a),推荐 H800 或 H100。
- CUDA: CUDA 12.3 或更高版本 (推荐 12.8 或更高版本)。
- 软件: Linux 操作系统 (如 Ubuntu、CentOS),Python 3.8 或更高版本,PyTorch 2.1 或更高版本,CUTLASS 3.6 或更高版本。
应用场景:
DeepGEMM 适用于大规模 AI 模型推理、混合专家 (MoE) 模型、低精度计算、高性能计算和深度学习框架优化等场景。
项目地址:
GitHub 仓库:https://github.com/deepseek-ai/DeepGEMM
DeepSeek 的这一开源举措,无疑将推动 AI 领域的发展,为研究人员和开发者提供更高效、更便捷的 FP8 矩阵乘法解决方案。随着 DeepGEMM 的广泛应用,我们有望看到更多创新性的 AI 应用涌现。
参考文献:
- DeepGEMM GitHub 仓库:https://github.com/deepseek-ai/DeepGEMM
- AI工具集:https://www.qtool.ai/ai-project/deepgemm-deepseek-fp8-gemm/
Views: 0