Deepseek 开源周第三发：DeepGEMM - 优化 NVIDIA H800 性能的 FP8 GEMM 库 Deepseek 开源周第三发 DeepGEMM：DeepGEMM 是一个支持密集型和 MoE GEMM 的 FP8 GEMM 库，核心逻辑仅约300行代码，极限情况下可以将 NVIDIA H800 的计算性能提高 2.7 倍。本文将详细介绍 DeepGEMM 的设计原理、性能优势以及快速启动方法。介绍DeepGEMM 是一个专为高效 FP8 通用矩阵乘法（GEMM）设计的库。它支持普通和混合专家（MoE）分组 GEMM，并且核心逻辑仅约300行代码。该库专为 NVIDIA H800 GPU 优化，能够在某些情况下将计算性能提升至1350+ TFLOPS。DeepGEMM 的设计非常简洁，易于理解和优化，是学习 Hopper FP8 矩阵乘法和优化技术的极佳资源。核心功能1. 支持多种 GEMM 类型DeepGEMM 支持以下类型的 GEMM：密集型 GEMM：适用于普通模型。分组 GEMM（连续布局）：适用于 MoE 模型中的专家具有相同形状的场景。分组 GEMM（掩膜布局）：适用于推理解码阶段，支持掩码张量。2. 高性能DeepGEMM 在极限情况下可以将 NVIDIA H800 的计算性能提高 2.7 倍。以下是一些性能测试结果：MNK计算 (TFLOPS)内存带宽 (GB/秒)加速 (倍)642112716820616882.76424576153628924551.7643276851221921431.8..................快速启动系统要求Hopper 架构 GPU，支持 sm_90aPython 3.8 或更高版本CUDA 12.3 或更高版本（推荐 12.8 或更高版本）PyTorch 2.1 或更高版本CUTLASS 3.6 或更高版本（可以通过 Git 子模块克隆）安装步骤安装依赖：pip install torch cutlass克隆 DeepGEMM 仓库：git clone https://github.com/deepseek-ai/DeepGEMM.git cd DeepGEMM使用 GEMM 内核：正常密集 GEMM：from deep_gemm import gemm_fp8_fp8_bf16_nt result = gemm_fp8_fp8_bf16_nt(A, B, C)分组 GEMM（连续布局）：from deep_gemm import m_grouped_gemm_fp8_fp8_bf16_nt_contiguous result = m_grouped_gemm_fp8_fp8_bf16_nt_contiguous(A, B, C, group_m)分组 GEMM（掩膜布局）：from deep_gemm import m_grouped_gemm_fp8_fp8_bf16_nt_masked result = m_grouped_gemm_fp8_fp8_bf16_nt_masked(A, B, C, mask)优化细节1. 持久扭曲专业化DeepGEMM 中的内核经过了 warp 专门化，以实现重叠数据移动、张量核心 MMA 指令和 CUDA 核心提升。这种方法能够在不同 warpgroups 之间优化寄存器计数控制，并尽可能重叠操作。2. Hopper TMA 功能Hopper 引入了张量内存加速器（TMA），DeepGEMM 充分利用 TMA 来实现更快、更异步的数据移动。具体应用包括 LHS、LHS 缩放因子和 RHS 矩阵的 TMA 负载、TMA 存储输出矩阵等。3. 完全 JIT 设计DeepGEMM 采用完全即时（JIT）设计，安装时无需编译。所有内核均在运行时编译，具有更高的灵活性和性能。GEMM 形状、块大小和管道阶段数等参数被视为编译时常量，编译器可以进行更多优化。结论DeepGEMM 是一个简洁高效、易于理解和优化的 FP8 GEMM 库，专为 NVIDIA H800 GPU 优化。通过减少依赖、简化设计和充分利用 Hopper 架构的新功能，DeepGEMM 能够在多种矩阵形状下达到出色的性能。对于需要高性能 GEMM 计算的深度学习项目，DeepGEMM 是一个值得尝试的工具。{anote icon="fa-download" href="https://github.com/deepseek-ai/DeepGEMM" type="error" content="点此查看项目详情"/}