DeepGEMM：DeepSeek 推出的 FP8 矩阵乘法库

0 50

DeepGEMM简介

DeepGEMM 是由 DeepSeek 推出的一个高效 FP8 矩阵乘法（GEMM）库，专为 NVIDIA Hopper 架构 GPU 设计。它以简洁高效为理念，支持普通和混合专家（MoE）分组的 GEMM 运算，并通过细粒度缩放提升精度。该库采用 CUDA 编写，利用轻量级 JIT 编译技术，无需安装时编译，运行时动态生成内核。它充分利用 Hopper 架构的 Tensor Memory Accelerator（TMA）特性，结合非对齐块大小和 FFMA SASS 交错优化等技术，实现高性能矩阵运算。DeepGEMM 在多种矩阵形状下表现出色，性能与专家优化的库相当或更好，同时保持了极简的核心代码设计，易于学习和使用。

DeepGEMM主要功能

高效FP8矩阵乘法：DeepGEMM专注于8位浮点（FP8）矩阵乘法，通过细粒度缩放技术，能够在低精度计算中保持较高的精度，同时显著提升计算效率。
支持多种GEMM模式：
- 普通GEMM：适用于常规矩阵乘法，满足一般深度学习模型的需求。
- 分组GEMM：针对混合专家（MoE）模型设计，支持连续布局和掩码布局两种模式，优化了多专家共享矩阵形状的场景，提升计算效率。
运行时即时编译（JIT）：DeepGEMM采用即时编译技术，所有内核在运行时动态编译。这种方式可以根据具体的矩阵形状和块大小进行优化，节省寄存器资源，同时为编译器提供更多优化机会。
专为Hopper架构优化：该库专为NVIDIA Hopper架构设计，充分利用了Hopper的Tensor Memory Accelerator（TMA）特性，包括高效的数据加载、存储、多播和描述符预取，显著提升数据传输效率。
细粒度缩放与双级累加：通过细粒度缩放技术，结合CUDA核心的双级累加机制，将FP8计算结果提升到更高精度的格式（如BF16），有效解决了FP8精度不足的问题。
轻量级设计：DeepGEMM的核心代码简洁，仅约300行，易于理解和扩展。它避免了复杂模板或代数结构的依赖，降低了学习和优化的难度，同时保持了高性能。

DeepGEMM技术原理

即时编译（JIT）技术：DeepGEMM在运行时动态编译内核，而不是在安装时编译。这种方式可以根据具体的矩阵形状和块大小进行优化，节省寄存器资源，并为编译器提供更多优化机会，尤其在小矩阵形状下表现突出。
Hopper架构的TMA特性：DeepGEMM充分利用了Hopper架构的Tensor Memory Accelerator（TMA）特性，通过高效的TMA加载、存储、多播和描述符预取，显著提升了数据传输效率，减少了内存访问延迟。
细粒度缩放与双级累加：为了解决FP8精度不足的问题，DeepGEMM采用了细粒度缩放技术，并结合CUDA核心的双级累加机制，将FP8计算结果提升到更高精度的格式（如BF16），从而确保计算精度。
非标准块大小支持：DeepGEMM支持非标准块大小，能够更好地适应特定矩阵形状，提高硬件资源利用率，避免因块大小对齐导致的硬件资源浪费。
指令级优化：通过分析编译器生成的机器代码，DeepGEMM实现了特殊的指令排序优化，例如调整FFMA指令的执行顺序，提升了计算单元的并行效率。
统一调度系统：DeepGEMM采用特殊的排布策略，通过增强缓存重用效率，减少内存访问次数，从而提升整体性能。

DeepGEMM应用场景

大规模深度学习模型训练：DeepGEMM 通过优化的 FP8 矩阵乘法，能够显著加快大型深度学习模型的训练速度，尤其适合计算密集型任务，比如训练超大规模的语言模型或复杂的视觉模型。
实时推理加速：在对响应速度要求极高的场景中，例如自动驾驶、实时图像识别或语音处理，DeepGEMM 的高效内核可以快速完成矩阵运算，确保系统在极短时间内给出结果。
混合专家（MoE）模型优化：DeepGEMM 特别支持 MoE 架构中的分组矩阵乘法，通过优化连续布局和掩码布局，能够高效处理多专家模型中的复杂计算，提升整体性能。
分布式计算与大规模并行处理：在分布式训练环境中，DeepGEMM 可以高效处理多 GPU 或多节点之间的矩阵运算，帮助加速大规模系统的训练过程，例如在推荐系统或大规模机器学习任务中。
科学计算与高性能计算（HPC）：对于需要大量矩阵运算的科学计算任务，如物理模拟、气象预测或生物信息学分析，DeepGEMM 提供了强大的加速能力，帮助研究人员更快地完成复杂计算。
深度学习框架的底层优化：DeepGEMM 可以集成到主流的深度学习框架（如 PyTorch 或 TensorFlow）中，作为底层优化工具，提升框架的整体性能，让开发者在不改变现有代码的情况下享受加速效果。
低精度计算需求：DeepGEMM 通过细粒度缩放和双级累加技术，解决了低精度计算（如 FP8）的精度问题，同时保持高效率输出，适用于对计算资源有限制但对精度要求较高的场景。