DeepGEMM

DeepGEMM简介

DeepGEMM是由DeepSeek开源的高效FP8矩阵乘法（GEMM）库，目前仅支持NVIDIA Hopper架构的张量核心。该库支持普通和混合专家（MoE）分组的GEMM操作，采用即时编译（JIT）技术，实现运行时动态优化，无需安装时编译。通过细粒度缩放和CUDA核心双级累加技术，DeepGEMM有效解决了FP8精度不足的问题，并利用Hopper的Tensor Memory Accelerator（TMA）提升数据传输效率。其核心代码仅有约300行，具备高度可读性和可扩展性，在多种矩阵形状下表现出色。

DeepGEMM的核心功能

• 高效FP8矩阵乘法（GEMM）：专为FP8（8位浮点）矩阵乘法设计，支持细粒度缩放，提升计算性能与精度。
• 支持普通与分组GEMM：
  • 普通GEMM：适用于标准矩阵乘法。
  • 分组GEMM：支持MoE模型中的分组矩阵乘法，包括连续布局和掩码布局。
• JIT即时编译：内核在运行时动态编译，根据参数进行优化，提升性能。
• Hopper架构优化：充分利用TMA特性，提升数据传输效率。
• 细粒度缩放与双级累加：提高FP8计算精度，输出更高精度格式如BF16。
• 轻量级设计：核心代码简洁，便于学习、理解和扩展。

DeepGEMM项目信息

• GitHub仓库：https://github.com/deepseek-ai/DeepGEMM

DeepGEMM性能表现

• 普通GEMM性能
  • 最高加速比：某些矩阵形状下可达2.7倍，显著提升运算效率。
  • 计算性能：支持超过1000 TFLOPS的计算能力，接近Hopper架构GPU理论峰值。
• 分组GEMM性能
  • 加速比：在MoE模型中实现1.1-1.2倍加速。
  • 内存带宽优化：基于TMA特性，实现接近硬件极限的性能。

系统要求

• 硬件要求：
  • GPU架构：需支持NVIDIA Hopper架构（sm_90a），推荐使用H800或H100等型号。
  • CUDA兼容性：需CUDA 12.3或更高版本。
• 软件要求：
  • 操作系统：推荐Linux（如Ubuntu、CentOS）。
  • Python版本：Python 3.8及以上。
  • CUDA工具包：CUDA 12.3或更高版本。
  • PyTorch：PyTorch 2.1或更高。
  • CUTLASS库：CUTLASS 3.6或更高。
• 其他要求：
  • 常规编译工具（如gcc、make）。
  • torch.utils.cpp_extension模块。

应用场景

• 大规模AI模型推理：提升高维矩阵运算速度。
• 混合专家（MoE）模型：优化分组矩阵乘法，提高计算效率。
• 低精度计算：通过细粒度缩放提升FP8精度。
• 高性能计算：利用Hopper架构特性提升性能。
• 深度学习框架优化：作为底层库加速模型训练与推理。