DeepGEMM概述简介

DeepGEMM是DeepSeek开源的为高效简洁的FP8矩阵乘法（GEMM）设计的库，目前仅支持NVIDIA Hopper张量核心。DeepGEMM支持普通和混合专家（MoE）分组的GEMM操作，基于即时编译（JIT）技术，无需安装时编译，支持在运行时动态优化。DeepGEMM基于细粒度缩放和CUDA核心双级累加技术，解决FP8精度不足的问题，同时用Hopper的Tensor Memory Accelerator（TMA）特性提升数据传输效率。DeepGEMM核心代码仅约300行，易于学习和优化。DeepGEMM的轻量设计，在多种矩阵形状上均达到或超过专家级优化库。

DeepGEMM的功能特色

高效FP8矩阵乘法（GEMM）：DeepGEMM是专为FP8（8位浮点）矩阵乘法设计的高效库，支持细粒度缩放，显著提升矩阵运算的性能和精度。

支持普通和分组GEMM：

普通GEMM：适用于常规的矩阵乘法操作。

分组GEMM：支持混合专家（MoE）模型中的分组矩阵乘法，包括连续布局（contiguous layout）和掩码布局（masked layout），优化多专家共享形状的场景。

即时编译（JIT）设计：基于即时编译技术，所有内核在运行时动态编译，无需安装时编译。根据矩阵形状、块大小等参数进行优化，节省寄存器提升性能。

Hopper架构优化：专为NVIDIA Hopper架构设计，充分利用Tensor Memory Accelerator（TMA）特性，包括TMA加载、存储、多播和描述符预取，显著提升数据传输效率。

细粒度缩放和双级累加：为解决FP8精度不足的问题，DeepGEMM引入细粒度缩放技术，基于CUDA核心的双级累加机制，将FP8计算结果提升到更高精度的格式（如BF16），确保计算精度。

轻量级设计：核心代码简洁，仅约300行，易于理解和扩展。避免复杂模板或代数结构的依赖，降低学习和优化的难度。

DeepGEMM项目介绍

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

DeepGEMM的性能表现

普通GEMM（非分组）性能

最高加速比：在某些矩阵形状下，DeepGEMM能达到2.7倍的加速比，显著提升矩阵乘法的效率。

计算性能：在大规模矩阵运算中，DeepGEMM能够实现超过1000 TFLOPS的计算性能，接近Hopper架构GPU的理论峰值。

分组GEMM（MoE模型）性能

加速比：在分组GEMM中，DeepGEMM能实现1.1-1.2倍的加速比，显著提升MoE模型的训练和推理效率。

内存带宽优化：基于TMA特性，DeepGEMM在内存带宽利用上表现出色，达到接近硬件极限的性能。

连续布局（Contiguous Layout）

掩码布局（Masked Layout）

硬件要求：

GPU架构：必须支持NVIDIA Hopper架构，具体要求为支持sm_90a的GPU。推荐使用H800或H100等Hopper架构的GPU，专为FP8计算和Tensor Core优化设计。

CUDA兼容性：需要支持CUDA 12.3或更高版本，推荐使用CUDA 12.8或更高版本获得最佳性能。

软件要求：

操作系统推荐：使用Linux操作系统（如Ubuntu、CentOS等），因为CUDA和PyTorch在Linux上支持更好。

Python版本：Python 3.8或更高版本。

CUDA工具包：CUDA 12.3或更高版本。CUDA版本必须与GPU架构匹配，推荐使用12.8或更高版本以充分利用Hopper架构的特性。

PyTorch：PyTorch 2.1或更高版本。

CUTLASS库：CUTLASS 3.6或更高版本。

其他要求：

常规的编译工具（如gcc、make等）。

torch.utils.cpp_extension模块，用于CUDA扩展。

DeepGEMM能做什么？

大规模AI大模型推理：加速高维矩阵乘法，提升推理速度。

混合专家（MoE）模型：优化分组矩阵乘法，提高计算效率。

低精度计算：通过细粒度缩放解决FP8精度问题，保持高精度输出。

高性能计算：基于Hopper架构特性，提升矩阵运算效率。

深度学习框架优化：作为底层优化库，加速模型训练和推理。